From 584f64ac8dac5c0346ae244ed4edd1ac545c0b1a Mon Sep 17 00:00:00 2001 From: PanZezhong Date: Fri, 7 Aug 2026 06:26:20 +0000 Subject: [PATCH] optimize: KDA resusing gated delta rule kernel --- .../ops/kimi_delta_attention/cuda/kernel.cuh | 171 +++++++++++ .../nvidia/kimi_delta_attention_nvidia.cu | 101 ++++++ .../cuda/kernel.cuh | 289 +++++++----------- .../cuda/recurrent_delta_rule_common.cuh | 283 +++++++++++++++++ .../recurrent_gated_delta_rule_nvidia.cu | 11 +- 5 files changed, 669 insertions(+), 186 deletions(-) create mode 100644 src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh diff --git a/src/infiniop/ops/kimi_delta_attention/cuda/kernel.cuh b/src/infiniop/ops/kimi_delta_attention/cuda/kernel.cuh index 2e68cd5ca..5abf02e2a 100644 --- a/src/infiniop/ops/kimi_delta_attention/cuda/kernel.cuh +++ b/src/infiniop/ops/kimi_delta_attention/cuda/kernel.cuh @@ -4,6 +4,8 @@ #include #include +#include "../../recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh" + template __device__ inline float kdaLoadAsFloat(const T *ptr, ptrdiff_t offset) { return static_cast(ptr[offset]); @@ -54,6 +56,175 @@ __device__ inline float kdaBlockReduceSum(float value, float *scratch) { return result; } +template +struct KimiDeltaRuleGatePolicy { + const Tgate *g; + const Tgate *beta; + const float *A_log; + const float *dt_bias; + float lower_bound; + ptrdiff_t g_s0; + ptrdiff_t g_s1; + ptrdiff_t g_s2; + ptrdiff_t beta_s0; + ptrdiff_t beta_s1; + ptrdiff_t beta_s2; + ptrdiff_t A_log_s0; + ptrdiff_t dt_bias_s0; + + __device__ void prepare(int token_batch, + int64_t token_idx, + int key_head_idx, + int, + Tcompute *decay, + Tcompute *beta_out) const { + if (threadIdx.x == 0) { + decay[0] = expf(A_log[static_cast(key_head_idx) * A_log_s0]); + const ptrdiff_t beta_offset = static_cast(token_batch) * beta_s0 + static_cast(token_idx) * beta_s1 + static_cast(key_head_idx) * beta_s2; + beta_out[0] = static_cast( + op::recurrent_gated_delta_rule::cuda::sigmoid( + op::recurrent_gated_delta_rule::cuda::loadAsFloat( + beta, beta_offset))); + } + __syncthreads(); + const Tcompute a_log_exp = decay[0]; + const ptrdiff_t gate_base = static_cast(token_batch) * g_s0 + static_cast(token_idx) * g_s1 + static_cast(key_head_idx) * g_s2; + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(D); + key_dim_idx += blockDim.x) { + const Tcompute raw_gate = static_cast( + op::recurrent_gated_delta_rule::cuda::loadAsFloat( + g, gate_base + key_dim_idx)) + + static_cast( + dt_bias[static_cast(key_head_idx) * dt_bias_s0 + key_dim_idx]); + decay[key_dim_idx] = expf( + static_cast(lower_bound) * op::recurrent_gated_delta_rule::cuda::sigmoid(a_log_exp * raw_gate)); + } + } +}; + +template +__global__ void kimiDeltaAttentionWarpCudaKernel( + Tdata *out, + Tdata *initial_state, + Tdata *final_state, + const Tdata *q, + const Tdata *k, + const Tdata *v, + const Tgate *g, + const Tgate *beta, + const float *A_log, + const float *dt_bias, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + bool cu_seqlens_i64, + bool initial_state_indices_i64, + bool final_state_indices_i64, + bool use_qk_l2norm, + bool has_cu_seqlens, + bool indexed_state_pool, + size_t total_tokens, + size_t pool_size, + float scale, + float lower_bound, + ptrdiff_t out_s0, + ptrdiff_t out_s1, + ptrdiff_t out_s2, + ptrdiff_t initial_s0, + ptrdiff_t initial_s1, + ptrdiff_t initial_s2, + ptrdiff_t initial_s3, + ptrdiff_t final_s0, + ptrdiff_t final_s1, + ptrdiff_t final_s2, + ptrdiff_t final_s3, + ptrdiff_t q_s0, + ptrdiff_t q_s1, + ptrdiff_t q_s2, + ptrdiff_t k_s0, + ptrdiff_t k_s1, + ptrdiff_t k_s2, + ptrdiff_t v_s0, + ptrdiff_t v_s1, + ptrdiff_t v_s2, + ptrdiff_t g_s0, + ptrdiff_t g_s1, + ptrdiff_t g_s2, + ptrdiff_t beta_s0, + ptrdiff_t beta_s1, + ptrdiff_t beta_s2, + ptrdiff_t A_log_s0, + ptrdiff_t dt_bias_s0) { + extern __shared__ char shared_memory[]; + const KimiDeltaRuleGatePolicy gate_policy{ + g, + beta, + A_log, + dt_bias, + lower_bound, + g_s0, + g_s1, + g_s2, + beta_s0, + beta_s1, + beta_s2, + A_log_s0, + dt_bias_s0, + }; + op::recurrent_gated_delta_rule::cuda::recurrentDeltaRuleWarpSequence< + Tdata, + Tcompute, + D, + D, + WARPS_PER_BLOCK>( + out, + initial_state, + final_state, + q, + k, + v, + cu_seqlens, + initial_state_indices, + final_state_indices, + cu_seqlens_i64, + initial_state_indices_i64, + final_state_indices_i64, + use_qk_l2norm, + has_cu_seqlens, + indexed_state_pool, + total_tokens, + pool_size, + gridDim.y, + 1, + static_cast(scale), + out_s0, + out_s1, + out_s2, + initial_s0, + initial_s1, + initial_s2, + initial_s3, + final_s0, + final_s1, + final_s2, + final_s3, + q_s0, + q_s1, + q_s2, + k_s0, + k_s1, + k_s2, + v_s0, + v_s1, + v_s2, + gate_policy, + reinterpret_cast(shared_memory)); +} + template __global__ void kimiDeltaAttentionDecodeCudaKernel( Tdata *out, diff --git a/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu b/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu index 4f1a8705e..8b7ba6bff 100644 --- a/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu +++ b/src/infiniop/ops/kimi_delta_attention/nvidia/kimi_delta_attention_nvidia.cu @@ -66,6 +66,88 @@ infiniStatus_t Descriptor::create( return INFINI_STATUS_SUCCESS; } +template +static infiniStatus_t launch_warp_sequence(const KimiDeltaAttentionInfo &info, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *A_log, + const void *dt_bias, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + cudaStream_t stream) { + constexpr size_t D = 128; + constexpr size_t WARPS_PER_BLOCK = 8; + constexpr size_t NUM_THREADS = WARPS_PER_BLOCK * 32; + const dim3 grid( + static_cast(info.B), + static_cast(info.H), + static_cast((D + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK)); + const dim3 block(NUM_THREADS); + const size_t shared = (D * 3 + NUM_THREADS + 1) * sizeof(float); + + kimiDeltaAttentionWarpCudaKernel + <<>>( + static_cast(out), + static_cast(initial_state), + static_cast(final_state), + static_cast(q), + static_cast(k), + static_cast(v), + static_cast(g), + static_cast(beta), + static_cast(A_log), + static_cast(dt_bias), + cu_seqlens, + initial_state_indices, + final_state_indices, + info.cu_seqlens_dtype == INFINI_DTYPE_I64, + info.initial_state_indices_dtype == INFINI_DTYPE_I64, + info.final_state_indices_dtype == INFINI_DTYPE_I64, + info.use_qk_l2norm, + info.has_cu_seqlens, + info.indexed_state_pool, + info.total_tokens, + info.pool_size, + info.scale, + info.lower_bound, + info.out_strides[0], + info.out_strides[1], + info.out_strides[2], + info.initial_state_strides[0], + info.initial_state_strides[1], + info.initial_state_strides[2], + info.initial_state_strides[3], + info.final_state_strides.empty() ? 0 : info.final_state_strides[0], + info.final_state_strides.empty() ? 0 : info.final_state_strides[1], + info.final_state_strides.empty() ? 0 : info.final_state_strides[2], + info.final_state_strides.empty() ? 0 : info.final_state_strides[3], + info.q_strides[0], + info.q_strides[1], + info.q_strides[2], + info.k_strides[0], + info.k_strides[1], + info.k_strides[2], + info.v_strides[0], + info.v_strides[1], + info.v_strides[2], + info.g_strides[0], + info.g_strides[1], + info.g_strides[2], + info.beta_strides[0], + info.beta_strides[1], + info.beta_strides[2], + info.A_log_strides[0], + info.dt_bias_strides[0]); + return INFINI_STATUS_SUCCESS; +} + template static infiniStatus_t launch_fallback(const KimiDeltaAttentionInfo &info, void *out, @@ -82,6 +164,25 @@ static infiniStatus_t launch_fallback(const KimiDeltaAttentionInfo &info, const void *initial_state_indices, const void *final_state_indices, cudaStream_t stream) { + if (info.D == 128) { + return launch_warp_sequence( + info, + out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + A_log, + dt_bias, + cu_seqlens, + initial_state_indices, + final_state_indices, + stream); + } + constexpr int threads = 256; dim3 grid(static_cast(info.B), static_cast(info.H), static_cast(info.D)); size_t shared = info.is_decode ? threads * sizeof(float) : (info.D * 3 + threads) * sizeof(float); diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh b/src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh index 9e952eb5a..df5529b8f 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh +++ b/src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh @@ -1,45 +1,48 @@ #ifndef __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ #define __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ -#include -#include - -__device__ inline int64_t loadStateIndex( - const void *indices, - bool is_i64, - int batch_idx, - int fallback) { - if (indices == nullptr) { - return static_cast(fallback); - } - return is_i64 - ? static_cast(indices)[batch_idx] - : static_cast(static_cast(indices)[batch_idx]); -} - -template -__device__ inline float loadAsFloat(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ inline float loadAsFloat(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ inline float loadAsFloat(const cuda_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -__device__ inline float warpReduceSum(float value) { -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - value += __shfl_down_sync(0xffffffff, value, offset, 32); +#include "recurrent_delta_rule_common.cuh" + +template +struct ScalarGatedDeltaRulePolicy { + const Tgate *g; + const Tgate *beta; + ptrdiff_t g_s0; + ptrdiff_t g_s1; + ptrdiff_t g_s2; + ptrdiff_t beta_s0; + ptrdiff_t beta_s1; + ptrdiff_t beta_s2; + + __device__ void prepare(int token_batch, + int64_t token_idx, + int, + int value_head_idx, + Tcompute *decay, + Tcompute *beta_out) const { + if (threadIdx.x == 0) { + const ptrdiff_t gate_offset = static_cast(token_batch) * g_s0 + static_cast(token_idx) * g_s1 + static_cast(value_head_idx) * g_s2; + const ptrdiff_t beta_offset = static_cast(token_batch) * beta_s0 + static_cast(token_idx) * beta_s1 + static_cast(value_head_idx) * beta_s2; + decay[0] = expf(static_cast( + op::recurrent_gated_delta_rule::cuda::loadAsFloat(g, gate_offset))); + beta_out[0] = static_cast( + op::recurrent_gated_delta_rule::cuda::loadAsFloat(beta, beta_offset)); + } + __syncthreads(); + const Tcompute scalar_decay = decay[0]; + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(Dk); + key_dim_idx += blockDim.x) { + decay[key_dim_idx] = scalar_decay; + } } - return __shfl_sync(0xffffffff, value, 0, 32); -} -template +}; + +template __device__ void recurrentGatedDeltaRuleIndexedPoolWarpKernel( Tdata *out, Tdata *initial_state, @@ -54,7 +57,9 @@ __device__ void recurrentGatedDeltaRuleIndexedPoolWarpKernel( bool initial_state_indices_i64, bool final_state_indices_i64, bool use_qk_l2norm, - size_t Hk, + bool indexed_state_pool, + size_t pool_size, + size_t num_key_heads, size_t value_heads_per_key_head, ptrdiff_t out_s0, ptrdiff_t out_s1, @@ -81,150 +86,66 @@ __device__ void recurrentGatedDeltaRuleIndexedPoolWarpKernel( ptrdiff_t g_s2, ptrdiff_t beta_s0, ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - constexpr int WARP_SIZE = 32; - constexpr int NUM_THREADS = WARPS_PER_BLOCK * WARP_SIZE; - - const int batch_idx = blockIdx.x; - const int value_head_idx = blockIdx.y; - const int warp_idx = threadIdx.x / WARP_SIZE; - const int lane_idx = threadIdx.x & (WARP_SIZE - 1); - const int value_dim_idx = blockIdx.z * WARPS_PER_BLOCK + warp_idx; - const int key_head_idx = value_head_idx / static_cast(value_heads_per_key_head); - - if (key_head_idx >= static_cast(Hk)) { - return; - } - - constexpr int seq_idx = 0; - const ptrdiff_t q_base = static_cast(batch_idx) * q_s0 + seq_idx * q_s1 + static_cast(key_head_idx) * q_s2; - const ptrdiff_t k_base = static_cast(batch_idx) * k_s0 + seq_idx * k_s1 + static_cast(key_head_idx) * k_s2; - - extern __shared__ char shared_mem_char[]; - Tcompute *shared_mem = reinterpret_cast(shared_mem_char); - Tcompute *q_local = shared_mem; - Tcompute *k_local = q_local + Dk; - Tcompute *norm_val = k_local + Dk; - - for (int i = threadIdx.x; i < static_cast(Dk); i += NUM_THREADS) { - q_local[i] = static_cast(loadAsFloat(q, q_base + i)); - k_local[i] = static_cast(loadAsFloat(k, k_base + i)); - } - - if (use_qk_l2norm) { - __syncthreads(); - Tcompute sum_sq = 0.0f; - for (int i = threadIdx.x; i < static_cast(Dk); i += NUM_THREADS) { - sum_sq += q_local[i] * q_local[i]; - } - norm_val[threadIdx.x] = sum_sq; - __syncthreads(); - if (threadIdx.x == 0) { - Tcompute total_sum_sq = 0.0f; - for (int i = 0; i < NUM_THREADS; ++i) { - total_sum_sq += norm_val[i]; - } - norm_val[0] = rsqrtf(total_sum_sq + 1e-6f); - } - __syncthreads(); - const Tcompute r_norm_q = norm_val[0]; - - for (int i = threadIdx.x; i < static_cast(Dk); i += NUM_THREADS) { - q_local[i] *= r_norm_q; - } - - sum_sq = 0.0f; - for (int i = threadIdx.x; i < static_cast(Dk); i += NUM_THREADS) { - sum_sq += k_local[i] * k_local[i]; - } - norm_val[threadIdx.x] = sum_sq; - __syncthreads(); - if (threadIdx.x == 0) { - Tcompute total_sum_sq = 0.0f; - for (int i = 0; i < NUM_THREADS; ++i) { - total_sum_sq += norm_val[i]; - } - norm_val[0] = rsqrtf(total_sum_sq + 1e-6f); - } - __syncthreads(); - const Tcompute r_norm_k = norm_val[0]; - - for (int i = threadIdx.x; i < static_cast(Dk); i += NUM_THREADS) { - k_local[i] *= r_norm_k; - } - } - - const Tcompute scale = rsqrtf(static_cast(Dk)); - for (int i = threadIdx.x; i < static_cast(Dk); i += NUM_THREADS) { - q_local[i] *= scale; - } - __syncthreads(); - - if (value_dim_idx >= static_cast(Dv)) { - return; - } - - int64_t read_slot = loadStateIndex(initial_state_indices, initial_state_indices_i64, batch_idx, batch_idx); - int64_t write_slot = final_state_indices == nullptr - ? static_cast(batch_idx) - : loadStateIndex(final_state_indices, final_state_indices_i64, batch_idx, batch_idx); - - const ptrdiff_t out_base = static_cast(batch_idx) * out_s0 + seq_idx * out_s1 + static_cast(value_head_idx) * out_s2; - if (read_slot < 0 || write_slot < 0) { - if (lane_idx == 0) { - out[out_base + value_dim_idx] = static_cast(0.0f); - } - return; - } - - const ptrdiff_t v_base = static_cast(batch_idx) * v_s0 + seq_idx * v_s1 + static_cast(value_head_idx) * v_s2; - const ptrdiff_t gate_offset = static_cast(batch_idx) * g_s0 + seq_idx * g_s1 + static_cast(value_head_idx) * g_s2; - const ptrdiff_t beta_offset = static_cast(batch_idx) * beta_s0 + seq_idx * beta_s1 + static_cast(value_head_idx) * beta_s2; - - const ptrdiff_t initial_base = static_cast(read_slot) * initial_s0 - + static_cast(value_head_idx) * initial_s1 - + static_cast(value_dim_idx) * initial_s2; - - Tdata *final_state_target = final_state_indices == nullptr ? final_state : initial_state; - const ptrdiff_t final_base = final_state_indices == nullptr - ? static_cast(batch_idx) * final_s0 - + static_cast(value_head_idx) * final_s1 - + static_cast(value_dim_idx) * final_s2 - : static_cast(write_slot) * initial_s0 - + static_cast(value_head_idx) * initial_s1 - + static_cast(value_dim_idx) * initial_s2; - const ptrdiff_t final_k_stride = final_state_indices == nullptr ? final_s3 : initial_s3; - - const Tcompute g_t = expf(static_cast(loadAsFloat(g, gate_offset))); - const Tcompute beta_t = static_cast(loadAsFloat(beta, beta_offset)); - - Tcompute kv_mem = 0.0f; - Tcompute hq_mem = 0.0f; - Tcompute kq_mem = 0.0f; - for (int dk_idx = lane_idx; dk_idx < static_cast(Dk); dk_idx += WARP_SIZE) { - const Tcompute h_prev = static_cast(loadAsFloat(initial_state, initial_base + static_cast(dk_idx) * initial_s3)); - const Tcompute k_t = k_local[dk_idx]; - const Tcompute q_t = q_local[dk_idx]; - kv_mem += (h_prev * g_t) * k_t; - hq_mem += h_prev * q_t; - kq_mem += k_t * q_t; - } - kv_mem = warpReduceSum(kv_mem); - hq_mem = warpReduceSum(hq_mem); - kq_mem = warpReduceSum(kq_mem); - - const Tcompute v_t = static_cast(loadAsFloat(v, v_base + value_dim_idx)); - const Tcompute delta = (v_t - kv_mem) * beta_t; - - if (lane_idx == 0) { - const Tcompute out_val = g_t * hq_mem + delta * kq_mem; - out[out_base + value_dim_idx] = static_cast(out_val); - } - - for (int dk_idx = lane_idx; dk_idx < static_cast(Dk); dk_idx += WARP_SIZE) { - const Tcompute h_prev = static_cast(loadAsFloat(initial_state, initial_base + static_cast(dk_idx) * initial_s3)); - const Tcompute h_final = (h_prev * g_t) + (k_local[dk_idx] * delta); - final_state_target[final_base + static_cast(dk_idx) * final_k_stride] = static_cast(h_final); - } + ptrdiff_t beta_s2, + Tcompute *shared) { + const ScalarGatedDeltaRulePolicy gate_policy{ + g, + beta, + g_s0, + g_s1, + g_s2, + beta_s0, + beta_s1, + beta_s2, + }; + op::recurrent_gated_delta_rule::cuda::recurrentDeltaRuleWarpSequence< + Tdata, + Tcompute, + Dk, + Dv, + WARPS_PER_BLOCK>( + out, + initial_state, + final_state, + q, + k, + v, + nullptr, + initial_state_indices, + final_state_indices, + false, + initial_state_indices_i64, + final_state_indices_i64, + use_qk_l2norm, + false, + indexed_state_pool, + 1, + pool_size, + num_key_heads, + value_heads_per_key_head, + rsqrtf(static_cast(Dk)), + out_s0, + out_s1, + out_s2, + initial_s0, + initial_s1, + initial_s2, + initial_s3, + final_s0, + final_s1, + final_s2, + final_s3, + q_s0, + q_s1, + q_s2, + k_s0, + k_s1, + k_s2, + v_s0, + v_s1, + v_s2, + gate_policy, + shared); } + #endif // __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh b/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh new file mode 100644 index 000000000..27931bb17 --- /dev/null +++ b/src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh @@ -0,0 +1,283 @@ +#ifndef __RECURRENT_DELTA_RULE_COMMON_CUH__ +#define __RECURRENT_DELTA_RULE_COMMON_CUH__ + +#include +#include + +namespace op::recurrent_gated_delta_rule::cuda { + +template +__device__ inline float loadAsFloat(const T *ptr, ptrdiff_t offset) { + return static_cast(ptr[offset]); +} + +template <> +__device__ inline float loadAsFloat(const half *ptr, ptrdiff_t offset) { + return __half2float(ptr[offset]); +} + +template <> +__device__ inline float loadAsFloat<__nv_bfloat16>(const __nv_bfloat16 *ptr, ptrdiff_t offset) { + return __bfloat162float(ptr[offset]); +} + +__device__ inline int64_t loadOptionalIndex(const void *indices, + bool is_i64, + int index, + int fallback) { + if (indices == nullptr) { + return static_cast(fallback); + } + return is_i64 + ? static_cast(indices)[index] + : static_cast(static_cast(indices)[index]); +} + +template +__device__ inline Tcompute warpReduceSum(Tcompute value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + return __shfl_sync(0xffffffff, value, 0); +} + +template +__device__ inline Tcompute blockReduceSum(Tcompute value, Tcompute *scratch) { + scratch[threadIdx.x] = value; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + scratch[threadIdx.x] += scratch[threadIdx.x + stride]; + } + __syncthreads(); + } + const Tcompute result = scratch[0]; + __syncthreads(); + return result; +} + +__device__ inline float sigmoid(float value) { + if (value >= 0.0f) { + const float exp_neg = expf(-value); + return 1.0f / (1.0f + exp_neg); + } + const float exp_pos = expf(value); + return exp_pos / (1.0f + exp_pos); +} + +// GatePolicy prepares one decay value per key dimension and a scalar beta for +// the current token. This keeps the state-update implementation shared while +// allowing scalar GDR and vector-decay KDA to retain their own gate semantics. +template +__device__ void recurrentDeltaRuleWarpSequence( + Tdata *out, + Tdata *initial_state, + Tdata *final_state, + const Tdata *q, + const Tdata *k, + const Tdata *v, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + bool cu_seqlens_i64, + bool initial_state_indices_i64, + bool final_state_indices_i64, + bool use_qk_l2norm, + bool has_cu_seqlens, + bool indexed_state_pool, + size_t total_tokens, + size_t pool_size, + size_t num_key_heads, + size_t value_heads_per_key_head, + Tcompute query_scale, + ptrdiff_t out_s0, + ptrdiff_t out_s1, + ptrdiff_t out_s2, + ptrdiff_t initial_s0, + ptrdiff_t initial_s1, + ptrdiff_t initial_s2, + ptrdiff_t initial_s3, + ptrdiff_t final_s0, + ptrdiff_t final_s1, + ptrdiff_t final_s2, + ptrdiff_t final_s3, + ptrdiff_t q_s0, + ptrdiff_t q_s1, + ptrdiff_t q_s2, + ptrdiff_t k_s0, + ptrdiff_t k_s1, + ptrdiff_t k_s2, + ptrdiff_t v_s0, + ptrdiff_t v_s1, + ptrdiff_t v_s2, + GatePolicy gate_policy, + Tcompute *shared) { + constexpr int WARP_SIZE = 32; + constexpr int NUM_THREADS = WARPS_PER_BLOCK * WARP_SIZE; + constexpr int STATE_VALUES_PER_LANE = (Dk + WARP_SIZE - 1) / WARP_SIZE; + + const int batch_idx = blockIdx.x; + const int value_head_idx = blockIdx.y; + const int key_head_idx = value_head_idx / static_cast(value_heads_per_key_head); + const int warp_idx = threadIdx.x / WARP_SIZE; + const int lane_idx = threadIdx.x & (WARP_SIZE - 1); + const int value_dim_idx = blockIdx.z * WARPS_PER_BLOCK + warp_idx; + const bool valid_value_dim = value_dim_idx < static_cast(Dv); + + if (key_head_idx >= static_cast(num_key_heads)) { + return; + } + + int64_t token_begin = 0; + int64_t token_end = static_cast(total_tokens); + if (has_cu_seqlens) { + token_begin = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx, 0); + token_end = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx + 1, 0); + if (token_begin < 0 || token_end < token_begin || token_end > static_cast(total_tokens)) { + return; + } + } + + int64_t read_slot = batch_idx; + int64_t write_slot = batch_idx; + if (indexed_state_pool) { + read_slot = loadOptionalIndex( + initial_state_indices, initial_state_indices_i64, batch_idx, batch_idx); + write_slot = final_state_indices == nullptr + ? static_cast(batch_idx) + : loadOptionalIndex( + final_state_indices, final_state_indices_i64, batch_idx, batch_idx); + if (read_slot < 0 || write_slot < 0 || read_slot >= static_cast(pool_size) || write_slot >= static_cast(pool_size)) { + if (valid_value_dim && lane_idx == 0) { + const int token_batch = has_cu_seqlens ? 0 : batch_idx; + for (int64_t token_idx = token_begin; token_idx < token_end; ++token_idx) { + const ptrdiff_t out_base = static_cast(token_batch) * out_s0 + static_cast(token_idx) * out_s1 + static_cast(value_head_idx) * out_s2; + out[out_base + value_dim_idx] = static_cast(0.0f); + } + } + return; + } + } + + const ptrdiff_t initial_base = static_cast(read_slot) * initial_s0 + static_cast(value_head_idx) * initial_s1 + static_cast(value_dim_idx) * initial_s2; + + Tdata *final_state_target = final_state_indices == nullptr ? final_state : initial_state; + const ptrdiff_t final_base = final_state_indices == nullptr + ? static_cast(batch_idx) * final_s0 + static_cast(value_head_idx) * final_s1 + static_cast(value_dim_idx) * final_s2 + : static_cast(write_slot) * initial_s0 + static_cast(value_head_idx) * initial_s1 + static_cast(value_dim_idx) * initial_s2; + const ptrdiff_t final_k_stride = final_state_indices == nullptr ? final_s3 : initial_s3; + + Tcompute state[STATE_VALUES_PER_LANE]; +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + state[i] = valid_value_dim && key_dim_idx < static_cast(Dk) + ? static_cast(loadAsFloat( + initial_state, + initial_base + static_cast(key_dim_idx) * initial_s3)) + : static_cast(0); + } + + Tcompute *q_local = shared; + Tcompute *k_local = q_local + Dk; + Tcompute *decay_local = k_local + Dk; + Tcompute *reduction_scratch = decay_local + Dk; + Tcompute *beta_shared = reduction_scratch + NUM_THREADS; + + const int token_batch = has_cu_seqlens ? 0 : batch_idx; + for (int64_t token_idx = token_begin; token_idx < token_end; ++token_idx) { + const ptrdiff_t q_base = static_cast(token_batch) * q_s0 + static_cast(token_idx) * q_s1 + static_cast(key_head_idx) * q_s2; + const ptrdiff_t k_base = static_cast(token_batch) * k_s0 + static_cast(token_idx) * k_s1 + static_cast(key_head_idx) * k_s2; + + Tcompute q_sum = 0; + Tcompute k_sum = 0; + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(Dk); + key_dim_idx += NUM_THREADS) { + const Tcompute q_value = static_cast(loadAsFloat(q, q_base + key_dim_idx)); + const Tcompute k_value = static_cast(loadAsFloat(k, k_base + key_dim_idx)); + q_local[key_dim_idx] = q_value; + k_local[key_dim_idx] = k_value; + q_sum += q_value * q_value; + k_sum += k_value * k_value; + } + q_sum = blockReduceSum(q_sum, reduction_scratch); + k_sum = blockReduceSum(k_sum, reduction_scratch); + + const Tcompute q_norm = use_qk_l2norm + ? rsqrtf(q_sum + static_cast(1e-6)) + : static_cast(1); + const Tcompute k_norm = use_qk_l2norm + ? rsqrtf(k_sum + static_cast(1e-6)) + : static_cast(1); + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(Dk); + key_dim_idx += NUM_THREADS) { + q_local[key_dim_idx] *= q_norm * query_scale; + k_local[key_dim_idx] *= k_norm; + } + + gate_policy.prepare( + token_batch, + token_idx, + key_head_idx, + value_head_idx, + decay_local, + beta_shared); + __syncthreads(); + + Tcompute kv_memory = 0; + Tcompute hq_memory = 0; + Tcompute kq_memory = 0; +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + if (valid_value_dim && key_dim_idx < static_cast(Dk)) { + const Tcompute decayed_state = state[i] * decay_local[key_dim_idx]; + const Tcompute k_value = k_local[key_dim_idx]; + const Tcompute q_value = q_local[key_dim_idx]; + kv_memory += decayed_state * k_value; + hq_memory += decayed_state * q_value; + kq_memory += k_value * q_value; + } + } + kv_memory = warpReduceSum(kv_memory); + hq_memory = warpReduceSum(hq_memory); + kq_memory = warpReduceSum(kq_memory); + + Tcompute delta = 0; + if (valid_value_dim && lane_idx == 0) { + const ptrdiff_t v_base = static_cast(token_batch) * v_s0 + static_cast(token_idx) * v_s1 + static_cast(value_head_idx) * v_s2; + const Tcompute v_value = static_cast(loadAsFloat(v, v_base + value_dim_idx)); + delta = (v_value - kv_memory) * beta_shared[0]; + const ptrdiff_t out_base = static_cast(token_batch) * out_s0 + static_cast(token_idx) * out_s1 + static_cast(value_head_idx) * out_s2; + out[out_base + value_dim_idx] = static_cast(hq_memory + delta * kq_memory); + } + delta = __shfl_sync(0xffffffff, delta, 0); + +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + if (valid_value_dim && key_dim_idx < static_cast(Dk)) { + state[i] = state[i] * decay_local[key_dim_idx] + k_local[key_dim_idx] * delta; + } + } + __syncthreads(); + } + +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + if (valid_value_dim && key_dim_idx < static_cast(Dk)) { + final_state_target[final_base + static_cast(key_dim_idx) * final_k_stride] = static_cast(state[i]); + } + } +} + +} // namespace op::recurrent_gated_delta_rule::cuda + +#endif // __RECURRENT_DELTA_RULE_COMMON_CUH__ diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu b/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu index 7697ad07d..367ab66f6 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu +++ b/src/infiniop/ops/recurrent_gated_delta_rule/nvidia/recurrent_gated_delta_rule_nvidia.cu @@ -16,6 +16,8 @@ INFINIOP_CUDA_KERNEL recurrentGatedDeltaRuleIndexedPoolWarp( bool initial_state_indices_i64, bool final_state_indices_i64, bool use_qk_l2norm, + bool indexed_state_pool, + size_t pool_size, size_t Hk, size_t value_heads_per_key_head, ptrdiff_t out_s0, @@ -44,11 +46,13 @@ INFINIOP_CUDA_KERNEL recurrentGatedDeltaRuleIndexedPoolWarp( ptrdiff_t beta_s0, ptrdiff_t beta_s1, ptrdiff_t beta_s2) { + extern __shared__ char shared_memory[]; recurrentGatedDeltaRuleIndexedPoolWarpKernel( out, initial_state, final_state, q, k, v, g, beta, initial_state_indices, final_state_indices, initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, + indexed_state_pool, pool_size, Hk, value_heads_per_key_head, out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, @@ -57,7 +61,8 @@ INFINIOP_CUDA_KERNEL recurrentGatedDeltaRuleIndexedPoolWarp( k_s0, k_s1, k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, - beta_s0, beta_s1, beta_s2); + beta_s0, beta_s1, beta_s2, + reinterpret_cast(shared_memory)); } namespace op { namespace recurrent_gated_delta_rule { @@ -114,7 +119,7 @@ infiniStatus_t launchIndexedPoolWarpKernelTyped( constexpr size_t NUM_THREADS = WARPS_PER_BLOCK * 32; dim3 grid(uint32_t(_info.B), uint32_t(_info.Hv), uint32_t((_info.Dv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK)); dim3 block(NUM_THREADS); - size_t shared_mem_size = (Dk + Dk + NUM_THREADS) * sizeof(float); + size_t shared_mem_size = (Dk * 3 + NUM_THREADS + 1) * sizeof(float); auto final_s0 = _info.final_state_strides.empty() ? 0 : _info.final_state_strides[0]; auto final_s1 = _info.final_state_strides.empty() ? 0 : _info.final_state_strides[1]; @@ -136,6 +141,8 @@ infiniStatus_t launchIndexedPoolWarpKernelTyped( initial_state_indices_i64, final_state_indices_i64, _info.use_qk_l2norm, + _info.indexed_state_pool, + _info.pool_size, _info.Hk, _info.value_heads_per_key_head, _info.out_strides[0],