| 1 | /* |
| 2 | * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. |
| 3 | */ |
| 4 | |
| 5 | #include <cstdint> |
| 6 | #include <cuda_bf16.h> |
| 7 | #include <cuda_runtime.h> |
| 8 | #include <array> |
| 9 | #include <optional> |
| 10 | |
| 11 | #include "../torch_utils.h" |
| 12 | #include "../../cuda_compat.h" |
| 13 | |
| 14 | namespace { |
| 15 | |
| 16 | constexpr int kDimK = 128; |
| 17 | constexpr int kDimV = 128; |
| 18 | constexpr int kKernelWidth = 4; |
| 19 | constexpr int kConvStateWidth = kKernelWidth - 1; |
| 20 | constexpr int kThreads = 256; |
| 21 | constexpr int kWarps = kThreads / 32; |
| 22 | constexpr int kChunkV = 32; |
| 23 | constexpr int kNumChunks = kDimV / kChunkV; |
| 24 | constexpr int kRowsPerWarp = kChunkV / kWarps; |
| 25 | |
| 26 | struct KdaDecodeStrides { |
| 27 | int64_t x_row; |
| 28 | int64_t beta_row; |
| 29 | int64_t onorm_row; |
| 30 | int64_t conv_slot; |
| 31 | int64_t state_slot; |
| 32 | }; |
| 33 | |
| 34 | __device__ __forceinline__ float bf16_load(const __nv_bfloat16* ptr, |
| 35 | int64_t idx) { |
| 36 | return __bfloat162float(ptr[idx]); |
| 37 | } |
| 38 | |
| 39 | __device__ __forceinline__ float bf16_load(const float* ptr, int64_t idx) { |
| 40 | return ptr[idx]; |
| 41 | } |
| 42 | |
| 43 | template <int kChannels> |
| 44 | __device__ __forceinline__ float conv_weight_load(const float* ptr, int channel, |
| 45 | int width) { |
| 46 | return ptr[width * kChannels + channel]; |
| 47 | } |
| 48 | |
| 49 | __device__ __forceinline__ __nv_bfloat16 bf16_store(float value) { |
| 50 | return __float2bfloat16(value); |
| 51 | } |
| 52 | |
| 53 | template <bool kUseCacheGlobalStore> |
| 54 | __device__ __forceinline__ void store_state_float4(float* ptr, float4 value) { |
| 55 | if constexpr (kUseCacheGlobalStore) { |
| 56 | __stcg(reinterpret_cast<float4*>(ptr), value); |
| 57 | } else { |
| 58 | *reinterpret_cast<float4*>(ptr) = value; |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | __device__ __forceinline__ float sigmoid_fast(float x) { |
| 63 | return 1.0f / (1.0f + __expf(-x)); |
| 64 | } |
| 65 | |
| 66 | __device__ __forceinline__ float silu_fast(float x) { |
| 67 | return x * sigmoid_fast(x); |
| 68 | } |
| 69 | |
| 70 | __device__ __forceinline__ float softplus_fast(float x) { |
| 71 | return x > 20.0f ? x : log1pf(__expf(x)); |
| 72 | } |
| 73 | |
| 74 | __device__ __forceinline__ float warp_reduce_sum(float value) { |
| 75 | #pragma unroll |
| 76 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 77 | value += __shfl_xor_sync(0xffffffffu, value, offset); |
| 78 | } |
| 79 | return value; |
| 80 | } |
| 81 | |
| 82 | __device__ __forceinline__ void cp_async_cg_16b(float* smem_ptr, |
| 83 | const float* gmem_ptr) { |
| 84 | uint32_t smem_addr = |
| 85 | static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr)); |
| 86 | asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" |
| 87 | : |
| 88 | : "r"(smem_addr), "l"(gmem_ptr)); |
| 89 | } |
| 90 | |
| 91 | __device__ __forceinline__ void cp_async_commit() { |
| 92 | asm volatile("cp.async.commit_group;\n" ::); |
| 93 | } |
| 94 | |
| 95 | __device__ __forceinline__ void cp_async_wait_all() { |
| 96 | asm volatile("cp.async.wait_all;\n" ::: "memory"); |
| 97 | } |
| 98 | |
| 99 | __device__ __forceinline__ void cp_async_wait_group_1() { |
| 100 | asm volatile("cp.async.wait_group 1;\n" ::: "memory"); |
| 101 | } |
| 102 | |
| 103 | template <int kCopyThreads> |
| 104 | __device__ __forceinline__ void cp_async_state_chunk_for(float* s_state, |
| 105 | const float* state, |
| 106 | int slot, int i_hv, |
| 107 | int HV, int chunk) { |
| 108 | constexpr int kFloat4PerChunk = kChunkV * kDimK / 4; |
| 109 | const int tid = threadIdx.x; |
| 110 | const int stage = chunk & 1; |
| 111 | const int v_base = chunk * kChunkV; |
| 112 | for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) { |
| 113 | const int elem = linear4 * 4; |
| 114 | const int row = elem / kDimK; |
| 115 | const int k = elem - row * kDimK; |
| 116 | float* dst = s_state + (stage * kChunkV + row) * kDimK + k; |
| 117 | const float* src = |
| 118 | state + ((slot * HV + i_hv) * kDimV + v_base + row) * kDimK + k; |
| 119 | cp_async_cg_16b(dst, src); |
| 120 | } |
| 121 | cp_async_commit(); |
| 122 | } |
| 123 | |
| 124 | __device__ __forceinline__ void cp_async_state_chunk(float* s_state, |
| 125 | const float* state, |
| 126 | int slot, int i_hv, int HV, |
| 127 | int chunk) { |
| 128 | cp_async_state_chunk_for<kThreads>(s_state, state, slot, i_hv, HV, chunk); |
| 129 | } |
| 130 | |
| 131 | __device__ __forceinline__ float block_reduce_sum(float value, float* scratch) { |
| 132 | const int tid = threadIdx.x; |
| 133 | const int lane = tid & 31; |
| 134 | const int warp = tid >> 5; |
| 135 | |
| 136 | float warp_total = warp_reduce_sum(value); |
| 137 | if (lane == 0) { |
| 138 | scratch[warp] = warp_total; |
| 139 | } |
| 140 | __syncthreads(); |
| 141 | |
| 142 | float block_total = 0.0f; |
| 143 | if (warp == 0) { |
| 144 | block_total = lane < kWarps ? scratch[lane] : 0.0f; |
| 145 | block_total = warp_reduce_sum(block_total); |
| 146 | if (lane == 0) { |
| 147 | scratch[0] = block_total; |
| 148 | } |
| 149 | } |
| 150 | __syncthreads(); |
| 151 | return scratch[0]; |
| 152 | } |
| 153 | |
| 154 | struct Sum2 { |
| 155 | float x; |
| 156 | float y; |
| 157 | }; |
| 158 | |
| 159 | __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { |
| 160 | #pragma unroll |
| 161 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 162 | x += __shfl_xor_sync(0xffffffffu, x, offset); |
| 163 | y += __shfl_xor_sync(0xffffffffu, y, offset); |
| 164 | } |
| 165 | return {x, y}; |
| 166 | } |
| 167 | |
| 168 | template <int kReduceWarps> |
| 169 | __device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, |
| 170 | float* scratch) { |
| 171 | const int lane = threadIdx.x & 31; |
| 172 | const int warp = threadIdx.x >> 5; |
| 173 | |
| 174 | const float warp_x = warp_reduce_sum(x); |
| 175 | const float warp_y = warp_reduce_sum(y); |
| 176 | if (lane == 0) { |
| 177 | scratch[warp] = warp_x; |
| 178 | scratch[kReduceWarps + warp] = warp_y; |
| 179 | } |
| 180 | __syncthreads(); |
| 181 | |
| 182 | float block_x = 0.0f; |
| 183 | float block_y = 0.0f; |
| 184 | if (warp == 0) { |
| 185 | block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; |
| 186 | block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; |
| 187 | block_x = warp_reduce_sum(block_x); |
| 188 | block_y = warp_reduce_sum(block_y); |
| 189 | if (lane == 0) { |
| 190 | scratch[0] = block_x; |
| 191 | scratch[1] = block_y; |
| 192 | } |
| 193 | } |
| 194 | __syncthreads(); |
| 195 | return {scratch[0], scratch[1]}; |
| 196 | } |
| 197 | |
| 198 | __device__ __forceinline__ Sum2 block_reduce_sum2(float x, float y, |
| 199 | float* scratch) { |
| 200 | return block_reduce_sum2_for<kWarps>(x, y, scratch); |
| 201 | } |
| 202 | |
| 203 | template <int kReduceWarps> |
| 204 | __device__ __forceinline__ float block_reduce_sum_active_for(float value, |
| 205 | float* scratch) { |
| 206 | const int lane = threadIdx.x & 31; |
| 207 | const int warp = threadIdx.x >> 5; |
| 208 | |
| 209 | float warp_total = 0.0f; |
| 210 | if (warp < kReduceWarps) { |
| 211 | warp_total = warp_reduce_sum(value); |
| 212 | } |
| 213 | if (lane == 0 && warp < kReduceWarps) { |
| 214 | scratch[warp] = warp_total; |
| 215 | } |
| 216 | __syncthreads(); |
| 217 | |
| 218 | float block_total = 0.0f; |
| 219 | if (warp == 0) { |
| 220 | block_total = lane < kReduceWarps ? scratch[lane] : 0.0f; |
| 221 | block_total = warp_reduce_sum(block_total); |
| 222 | if (lane == 0) { |
| 223 | scratch[0] = block_total; |
| 224 | } |
| 225 | } |
| 226 | __syncthreads(); |
| 227 | return scratch[0]; |
| 228 | } |
| 229 | |
| 230 | template <int kReduceWarps> |
| 231 | __device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, |
| 232 | float* scratch) { |
| 233 | const int lane = threadIdx.x & 31; |
| 234 | const int warp = threadIdx.x >> 5; |
| 235 | |
| 236 | float warp_x = 0.0f; |
| 237 | float warp_y = 0.0f; |
| 238 | if (warp < kReduceWarps) { |
| 239 | warp_x = warp_reduce_sum(x); |
| 240 | warp_y = warp_reduce_sum(y); |
| 241 | } |
| 242 | if (lane == 0 && warp < kReduceWarps) { |
| 243 | scratch[warp] = warp_x; |
| 244 | scratch[kReduceWarps + warp] = warp_y; |
| 245 | } |
| 246 | __syncthreads(); |
| 247 | |
| 248 | float block_x = 0.0f; |
| 249 | float block_y = 0.0f; |
| 250 | if (warp == 0) { |
| 251 | block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; |
| 252 | block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; |
| 253 | block_x = warp_reduce_sum(block_x); |
| 254 | block_y = warp_reduce_sum(block_y); |
| 255 | if (lane == 0) { |
| 256 | scratch[0] = block_x; |
| 257 | scratch[1] = block_y; |
| 258 | } |
| 259 | } |
| 260 | __syncthreads(); |
| 261 | return {scratch[0], scratch[1]}; |
| 262 | } |
| 263 | |
| 264 | template < |
| 265 | bool kApplyOnorm, bool kUseStaticDecodeLayout = false, int kFixedHeads = 0, |
| 266 | int kFixedValueHeads = 0, bool kUseHeadGrid = false, |
| 267 | bool kAccumulateOnormSumsq = false, bool kUseActiveQkReduction = false, |
| 268 | bool kUseCacheGlobalStore = false, bool kComputeOutputBeforeStore = false, |
| 269 | bool kSkipWarpSync = false, bool kPreloadOnormParams = false, |
| 270 | bool kPrefetchNextStateChunk = false, bool kUseActiveOnormReduction = false, |
| 271 | bool kUpdateConvState = false, bool kUseLowerBound = false, |
| 272 | bool kApplyBetaSigmoid = true, int kConvStateChannelStride = 1, |
| 273 | int kConvStateTapStride = 3 * kFixedHeads * kDimK> |
| 274 | __global__ |
| 275 | __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( |
| 276 | const __nv_bfloat16* __restrict__ x_q, |
| 277 | const __nv_bfloat16* __restrict__ x_k, |
| 278 | const __nv_bfloat16* __restrict__ x_v, const float* __restrict__ w_q_t, |
| 279 | const float* __restrict__ w_k_t, const float* __restrict__ w_v_t, |
| 280 | const float* __restrict__ bias_q, const float* __restrict__ bias_k, |
| 281 | const float* __restrict__ bias_v, __nv_bfloat16* __restrict__ cs_q, |
| 282 | __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v, |
| 283 | const float* __restrict__ a_log, const __nv_bfloat16* __restrict__ g, |
| 284 | const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ beta, |
| 285 | const __nv_bfloat16* __restrict__ onorm_g, |
| 286 | const float* __restrict__ onorm_weight, |
| 287 | const int* __restrict__ ssm_state_indices, |
| 288 | const int* __restrict__ cu_seqlens, float* __restrict__ state, |
| 289 | __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound, |
| 290 | float scale, float onorm_eps, KdaDecodeStrides strides) { |
| 291 | const int tid = threadIdx.x; |
| 292 | const int lane = tid & 31; |
| 293 | const int warp = tid >> 5; |
| 294 | int i_n; |
| 295 | int i_hv; |
| 296 | int i_h; |
| 297 | int bos; |
| 298 | int slot; |
| 299 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 |
| 300 | cudaGridDependencySynchronize(); |
| 301 | #endif |
| 302 | if constexpr (kUseStaticDecodeLayout) { |
| 303 | if constexpr (kUseHeadGrid) { |
| 304 | i_n = blockIdx.x; |
| 305 | i_hv = blockIdx.y; |
| 306 | } else { |
| 307 | const int nhv = blockIdx.x; |
| 308 | i_n = nhv / kFixedValueHeads; |
| 309 | i_hv = nhv - i_n * kFixedValueHeads; |
| 310 | } |
| 311 | i_h = i_hv; |
| 312 | bos = i_n; |
| 313 | slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; |
| 314 | } else { |
| 315 | const int nhv = blockIdx.x; |
| 316 | i_n = nhv / HV; |
| 317 | i_hv = nhv - i_n * HV; |
| 318 | const int hv_per_h = HV / H; |
| 319 | i_h = i_hv / hv_per_h; |
| 320 | |
| 321 | bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; |
| 322 | const int eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; |
| 323 | if (eos <= bos) { |
| 324 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 |
| 325 | cudaTriggerProgrammaticLaunchCompletion(); |
| 326 | #endif |
| 327 | return; |
| 328 | } |
| 329 | slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; |
| 330 | } |
| 331 | |
| 332 | constexpr int kLocalDim = kFixedHeads * kDimK; |
| 333 | const int hk_off = i_h * kDimK; |
| 334 | const int hv_off = i_hv * kDimV; |
| 335 | constexpr int hv_count = kFixedValueHeads; |
| 336 | float* const state_for_slot = state + slot * strides.state_slot; |
| 337 | const int64_t conv_slot_offset = slot * strides.conv_slot; |
| 338 | __nv_bfloat16* const cs_q_for_slot = cs_q + conv_slot_offset; |
| 339 | __nv_bfloat16* const cs_k_for_slot = cs_k + conv_slot_offset; |
| 340 | __nv_bfloat16* const cs_v_for_slot = cs_v + conv_slot_offset; |
| 341 | |
| 342 | __shared__ float s_state[2][kChunkV][kDimK]; |
| 343 | __shared__ float s_q[kDimK]; |
| 344 | __shared__ float s_k[kDimK]; |
| 345 | __shared__ float s_decay[kDimK]; |
| 346 | __shared__ float s_v[kDimV]; |
| 347 | __shared__ float s_o[kDimV]; |
| 348 | __shared__ float s_reduce[kThreads]; |
| 349 | __shared__ float s_beta; |
| 350 | float pre_onorm_gate = 0.0f; |
| 351 | float pre_onorm_weight = 0.0f; |
| 352 | |
| 353 | cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, 0); |
| 354 | |
| 355 | if constexpr (kUpdateConvState) { |
| 356 | if (tid < kDimK) { |
| 357 | const int k = tid; |
| 358 | const int hk = hk_off + k; |
| 359 | const int64_t xq_idx = bos * strides.x_row + i_h * kDimK + k; |
| 360 | const float exp_a = |
| 361 | __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); |
| 362 | |
| 363 | float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); |
| 364 | float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); |
| 365 | __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); |
| 366 | __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); |
| 367 | __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); |
| 368 | __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); |
| 369 | #pragma unroll |
| 370 | for (int w = 0; w < kConvStateWidth; ++w) { |
| 371 | const int cs_idx = |
| 372 | hk * kConvStateChannelStride + w * kConvStateTapStride; |
| 373 | const __nv_bfloat16 q_state = cs_q_for_slot[cs_idx]; |
| 374 | const __nv_bfloat16 k_state = cs_k_for_slot[cs_idx]; |
| 375 | q_acc += __bfloat162float(q_state) * |
| 376 | conv_weight_load<kLocalDim>(w_q_t, hk, w); |
| 377 | k_acc += __bfloat162float(k_state) * |
| 378 | conv_weight_load<kLocalDim>(w_k_t, hk, w); |
| 379 | if (w == 1) { |
| 380 | q_shift0 = q_state; |
| 381 | k_shift0 = k_state; |
| 382 | } else if (w == 2) { |
| 383 | q_shift1 = q_state; |
| 384 | k_shift1 = k_state; |
| 385 | } |
| 386 | } |
| 387 | const __nv_bfloat16 q_new = x_q[xq_idx]; |
| 388 | const __nv_bfloat16 k_new = x_k[xq_idx]; |
| 389 | q_acc += __bfloat162float(q_new) * |
| 390 | conv_weight_load<kLocalDim>(w_q_t, hk, kKernelWidth - 1); |
| 391 | k_acc += __bfloat162float(k_new) * |
| 392 | conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1); |
| 393 | |
| 394 | cs_q_for_slot[hk * kConvStateChannelStride] = q_shift0; |
| 395 | cs_q_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] = |
| 396 | q_shift1; |
| 397 | cs_q_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] = |
| 398 | q_new; |
| 399 | cs_k_for_slot[hk * kConvStateChannelStride] = k_shift0; |
| 400 | cs_k_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] = |
| 401 | k_shift1; |
| 402 | cs_k_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] = |
| 403 | k_new; |
| 404 | |
| 405 | s_q[k] = silu_fast(q_acc); |
| 406 | s_k[k] = silu_fast(k_acc); |
| 407 | |
| 408 | const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; |
| 409 | const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; |
| 410 | if constexpr (kUseLowerBound) { |
| 411 | s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); |
| 412 | } else { |
| 413 | s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); |
| 414 | } |
| 415 | } |
| 416 | } else { |
| 417 | if (tid < kDimK) { |
| 418 | const int k = tid; |
| 419 | const int hk = hk_off + k; |
| 420 | const float exp_a = |
| 421 | __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); |
| 422 | |
| 423 | float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); |
| 424 | float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); |
| 425 | #pragma unroll |
| 426 | for (int w = 0; w < kConvStateWidth; ++w) { |
| 427 | const int cs_idx = |
| 428 | hk * kConvStateChannelStride + w * kConvStateTapStride; |
| 429 | q_acc += bf16_load(cs_q_for_slot, cs_idx) * |
| 430 | conv_weight_load<kLocalDim>(w_q_t, hk, w); |
| 431 | k_acc += bf16_load(cs_k_for_slot, cs_idx) * |
| 432 | conv_weight_load<kLocalDim>(w_k_t, hk, w); |
| 433 | } |
| 434 | q_acc += bf16_load(x_q, bos * strides.x_row + i_h * kDimK + k) * |
| 435 | conv_weight_load<kLocalDim>(w_q_t, hk, kKernelWidth - 1); |
| 436 | k_acc += bf16_load(x_k, bos * strides.x_row + i_h * kDimK + k) * |
| 437 | conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1); |
| 438 | |
| 439 | s_q[k] = silu_fast(q_acc); |
| 440 | s_k[k] = silu_fast(k_acc); |
| 441 | |
| 442 | const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; |
| 443 | const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; |
| 444 | if constexpr (kUseLowerBound) { |
| 445 | s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); |
| 446 | } else { |
| 447 | s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); |
| 448 | } |
| 449 | } |
| 450 | } |
| 451 | |
| 452 | if constexpr (kUpdateConvState) { |
| 453 | if (tid < kDimV) { |
| 454 | const int v = tid; |
| 455 | const int hvv = hv_off + v; |
| 456 | const int64_t xv_idx = bos * strides.x_row + i_hv * kDimV + v; |
| 457 | |
| 458 | float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); |
| 459 | __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); |
| 460 | __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); |
| 461 | #pragma unroll |
| 462 | for (int w = 0; w < kConvStateWidth; ++w) { |
| 463 | const int cs_idx = |
| 464 | hvv * kConvStateChannelStride + w * kConvStateTapStride; |
| 465 | const __nv_bfloat16 v_state = cs_v_for_slot[cs_idx]; |
| 466 | v_acc += __bfloat162float(v_state) * |
| 467 | conv_weight_load<kLocalDim>(w_v_t, hvv, w); |
| 468 | if (w == 1) { |
| 469 | v_shift0 = v_state; |
| 470 | } else if (w == 2) { |
| 471 | v_shift1 = v_state; |
| 472 | } |
| 473 | } |
| 474 | const __nv_bfloat16 v_new = x_v[xv_idx]; |
| 475 | v_acc += __bfloat162float(v_new) * |
| 476 | conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1); |
| 477 | cs_v_for_slot[hvv * kConvStateChannelStride] = v_shift0; |
| 478 | cs_v_for_slot[hvv * kConvStateChannelStride + kConvStateTapStride] = |
| 479 | v_shift1; |
| 480 | cs_v_for_slot[hvv * kConvStateChannelStride + 2 * kConvStateTapStride] = |
| 481 | v_new; |
| 482 | s_v[v] = silu_fast(v_acc); |
| 483 | |
| 484 | if constexpr (kApplyOnorm && kPreloadOnormParams) { |
| 485 | const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; |
| 486 | pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); |
| 487 | pre_onorm_weight = onorm_weight[v]; |
| 488 | } |
| 489 | } |
| 490 | } else { |
| 491 | if (tid < kDimV) { |
| 492 | const int v = tid; |
| 493 | const int hvv = hv_off + v; |
| 494 | |
| 495 | float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); |
| 496 | #pragma unroll |
| 497 | for (int w = 0; w < kConvStateWidth; ++w) { |
| 498 | const int cs_idx = |
| 499 | hvv * kConvStateChannelStride + w * kConvStateTapStride; |
| 500 | v_acc += bf16_load(cs_v_for_slot, cs_idx) * |
| 501 | conv_weight_load<kLocalDim>(w_v_t, hvv, w); |
| 502 | } |
| 503 | v_acc += bf16_load(x_v, bos * strides.x_row + i_hv * kDimV + v) * |
| 504 | conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1); |
| 505 | s_v[v] = silu_fast(v_acc); |
| 506 | |
| 507 | if constexpr (kApplyOnorm && kPreloadOnormParams) { |
| 508 | const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; |
| 509 | pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); |
| 510 | pre_onorm_weight = onorm_weight[v]; |
| 511 | } |
| 512 | } |
| 513 | } |
| 514 | |
| 515 | if (tid == 0) { |
| 516 | const float beta_raw = bf16_load(beta, bos * strides.beta_row + i_hv); |
| 517 | if constexpr (kApplyBetaSigmoid) { |
| 518 | s_beta = sigmoid_fast(beta_raw); |
| 519 | } else { |
| 520 | s_beta = beta_raw; |
| 521 | } |
| 522 | } |
| 523 | __syncthreads(); |
| 524 | |
| 525 | if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { |
| 526 | cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, |
| 527 | 1); |
| 528 | } |
| 529 | |
| 530 | const float q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; |
| 531 | const float k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; |
| 532 | Sum2 qk_sum; |
| 533 | if constexpr (kUseActiveQkReduction) { |
| 534 | qk_sum = block_reduce_sum2_active_for<kDimK / 32>(q_sq, k_sq, s_reduce); |
| 535 | } else { |
| 536 | qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); |
| 537 | } |
| 538 | if (tid < kDimK) { |
| 539 | s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; |
| 540 | s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); |
| 541 | } |
| 542 | __syncthreads(); |
| 543 | |
| 544 | const int k_base = lane * 4; |
| 545 | const float4 q4 = *reinterpret_cast<const float4*>(s_q + k_base); |
| 546 | const float4 k4 = *reinterpret_cast<const float4*>(s_k + k_base); |
| 547 | const float4 decay4 = *reinterpret_cast<const float4*>(s_decay + k_base); |
| 548 | float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; |
| 549 | float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; |
| 550 | float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; |
| 551 | float o_sumsq = 0.0f; |
| 552 | |
| 553 | #pragma unroll |
| 554 | for (int chunk = 0; chunk < kNumChunks; ++chunk) { |
| 555 | if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { |
| 556 | if (chunk + 1 < kNumChunks) { |
| 557 | cp_async_wait_group_1(); |
| 558 | } else { |
| 559 | cp_async_wait_all(); |
| 560 | } |
| 561 | } else { |
| 562 | cp_async_wait_all(); |
| 563 | } |
| 564 | if constexpr (!kSkipWarpSync) { |
| 565 | __syncwarp(); |
| 566 | } |
| 567 | |
| 568 | if constexpr (!kPrefetchNextStateChunk) { |
| 569 | if (chunk + 1 < kNumChunks) { |
| 570 | cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, |
| 571 | hv_count, chunk + 1); |
| 572 | } |
| 573 | } |
| 574 | |
| 575 | #pragma unroll |
| 576 | for (int row = 0; row < kRowsPerWarp; row += 2) { |
| 577 | const int v_row_a = warp + row * kWarps; |
| 578 | const int v_row_b = warp + (row + 1) * kWarps; |
| 579 | const int v0 = chunk * kChunkV + v_row_a; |
| 580 | const int v1 = chunk * kChunkV + v_row_b; |
| 581 | float h_a_vals[4]; |
| 582 | float h_b_vals[4]; |
| 583 | float dot_hk_a = 0.0f; |
| 584 | float dot_hk_b = 0.0f; |
| 585 | |
| 586 | const float4 raw_h_a = *reinterpret_cast<const float4*>( |
| 587 | &s_state[chunk & 1][v_row_a][k_base]); |
| 588 | const float4 raw_h_b = *reinterpret_cast<const float4*>( |
| 589 | &s_state[chunk & 1][v_row_b][k_base]); |
| 590 | h_a_vals[0] = raw_h_a.x * r_decay[0]; |
| 591 | h_a_vals[1] = raw_h_a.y * r_decay[1]; |
| 592 | h_a_vals[2] = raw_h_a.z * r_decay[2]; |
| 593 | h_a_vals[3] = raw_h_a.w * r_decay[3]; |
| 594 | h_b_vals[0] = raw_h_b.x * r_decay[0]; |
| 595 | h_b_vals[1] = raw_h_b.y * r_decay[1]; |
| 596 | h_b_vals[2] = raw_h_b.z * r_decay[2]; |
| 597 | h_b_vals[3] = raw_h_b.w * r_decay[3]; |
| 598 | dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + |
| 599 | h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; |
| 600 | dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + |
| 601 | h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; |
| 602 | |
| 603 | const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b); |
| 604 | const float v_new0 = (s_v[v0] - dot_hk.x) * s_beta; |
| 605 | const float v_new1 = (s_v[v1] - dot_hk.y) * s_beta; |
| 606 | |
| 607 | float dot_hq_a = 0.0f; |
| 608 | float dot_hq_b = 0.0f; |
| 609 | const int state_idx_a = (i_hv * kDimV + v0) * kDimK + k_base; |
| 610 | const int state_idx_b = (i_hv * kDimV + v1) * kDimK + k_base; |
| 611 | const float h_a_0 = h_a_vals[0] + r_k[0] * v_new0; |
| 612 | const float h_a_1 = h_a_vals[1] + r_k[1] * v_new0; |
| 613 | const float h_a_2 = h_a_vals[2] + r_k[2] * v_new0; |
| 614 | const float h_a_3 = h_a_vals[3] + r_k[3] * v_new0; |
| 615 | const float h_b_0 = h_b_vals[0] + r_k[0] * v_new1; |
| 616 | const float h_b_1 = h_b_vals[1] + r_k[1] * v_new1; |
| 617 | const float h_b_2 = h_b_vals[2] + r_k[2] * v_new1; |
| 618 | const float h_b_3 = h_b_vals[3] + r_k[3] * v_new1; |
| 619 | if constexpr (kComputeOutputBeforeStore) { |
| 620 | dot_hq_a = |
| 621 | h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; |
| 622 | dot_hq_b = |
| 623 | h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; |
| 624 | store_state_float4<kUseCacheGlobalStore>( |
| 625 | state_for_slot + state_idx_a, |
| 626 | make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); |
| 627 | store_state_float4<kUseCacheGlobalStore>( |
| 628 | state_for_slot + state_idx_b, |
| 629 | make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); |
| 630 | } else { |
| 631 | store_state_float4<kUseCacheGlobalStore>( |
| 632 | state_for_slot + state_idx_a, |
| 633 | make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); |
| 634 | store_state_float4<kUseCacheGlobalStore>( |
| 635 | state_for_slot + state_idx_b, |
| 636 | make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); |
| 637 | dot_hq_a = |
| 638 | h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; |
| 639 | dot_hq_b = |
| 640 | h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; |
| 641 | } |
| 642 | |
| 643 | const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b); |
| 644 | if (lane == 0) { |
| 645 | s_o[v0] = dot_hq.x; |
| 646 | s_o[v1] = dot_hq.y; |
| 647 | if constexpr (kApplyOnorm && kAccumulateOnormSumsq) { |
| 648 | o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y; |
| 649 | } |
| 650 | } |
| 651 | } |
| 652 | |
| 653 | if constexpr (kPrefetchNextStateChunk) { |
| 654 | if (chunk + 2 < kNumChunks) { |
| 655 | cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, |
| 656 | hv_count, chunk + 2); |
| 657 | } |
| 658 | } |
| 659 | } |
| 660 | __syncthreads(); |
| 661 | |
| 662 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 |
| 663 | cudaTriggerProgrammaticLaunchCompletion(); |
| 664 | #endif |
| 665 | |
| 666 | if constexpr (kApplyOnorm) { |
| 667 | if constexpr (kAccumulateOnormSumsq) { |
| 668 | if (lane == 0) { |
| 669 | s_reduce[warp] = o_sumsq; |
| 670 | } |
| 671 | __syncthreads(); |
| 672 | |
| 673 | float total_sumsq = 0.0f; |
| 674 | if (warp == 0) { |
| 675 | total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; |
| 676 | total_sumsq = warp_reduce_sum(total_sumsq); |
| 677 | if (lane == 0) { |
| 678 | s_reduce[0] = total_sumsq; |
| 679 | } |
| 680 | } |
| 681 | __syncthreads(); |
| 682 | |
| 683 | if (tid < kDimV) { |
| 684 | const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; |
| 685 | const float raw_o = s_o[tid]; |
| 686 | const float rstd = |
| 687 | rsqrtf(s_reduce[0] / static_cast<float>(kDimV) + onorm_eps); |
| 688 | float gate; |
| 689 | float weight; |
| 690 | if constexpr (kPreloadOnormParams) { |
| 691 | gate = pre_onorm_gate; |
| 692 | weight = pre_onorm_weight; |
| 693 | } else { |
| 694 | const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; |
| 695 | gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); |
| 696 | weight = onorm_weight[tid]; |
| 697 | } |
| 698 | const float y = raw_o * rstd * weight * gate; |
| 699 | out[out_idx] = bf16_store(y); |
| 700 | } |
| 701 | } else { |
| 702 | const float raw_o = tid < kDimV ? s_o[tid] : 0.0f; |
| 703 | const float o_sq = raw_o * raw_o; |
| 704 | float sumsq; |
| 705 | if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) { |
| 706 | sumsq = block_reduce_sum_active_for<kDimV / 32>(o_sq, s_reduce); |
| 707 | } else { |
| 708 | sumsq = block_reduce_sum(o_sq, s_reduce); |
| 709 | } |
| 710 | |
| 711 | if (tid < kDimV) { |
| 712 | const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; |
| 713 | const float rstd = |
| 714 | rsqrtf(sumsq / static_cast<float>(kDimV) + onorm_eps); |
| 715 | float gate; |
| 716 | float weight; |
| 717 | if constexpr (kPreloadOnormParams) { |
| 718 | gate = pre_onorm_gate; |
| 719 | weight = pre_onorm_weight; |
| 720 | } else { |
| 721 | const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; |
| 722 | gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); |
| 723 | weight = onorm_weight[tid]; |
| 724 | } |
| 725 | const float y = raw_o * rstd * weight * gate; |
| 726 | out[out_idx] = bf16_store(y); |
| 727 | } |
| 728 | } |
| 729 | } else { |
| 730 | if (tid < kDimV) { |
| 731 | const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; |
| 732 | out[out_idx] = bf16_store(s_o[tid]); |
| 733 | } |
| 734 | } |
| 735 | } |
| 736 | |
| 737 | template <int kHeads, bool kApplyOnorm, bool kUpdateConvState, |
| 738 | bool kUseLowerBound, bool kApplyBetaSigmoid, |
| 739 | int kConvStateChannelStride, int kConvStateTapStride> |
| 740 | void launch_kda_decode_many_heads_raw( |
| 741 | const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, |
| 742 | const void* w_k_t, const void* w_v_t, const void* bias_q, |
| 743 | const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, |
| 744 | const float* a_log, const void* g, const float* dt_bias, const void* beta, |
| 745 | const void* onorm_g, const float* onorm_weight, |
| 746 | const int* ssm_state_indices, const int* cu_seqlens, float* state, |
| 747 | void* out, int B, int H, int HV, float lower_bound, float scale, |
| 748 | float onorm_eps, KdaDecodeStrides strides, cudaStream_t stream) { |
| 749 | auto kernel = &kda_decode_fusion_many_heads_kernel< |
| 750 | kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false, |
| 751 | false, true, true, true, kUpdateConvState, kUseLowerBound, |
| 752 | kApplyBetaSigmoid, kConvStateChannelStride, kConvStateTapStride>; |
| 753 | cudaLaunchConfig_t config{}; |
| 754 | config.gridDim = dim3(B, kHeads); |
| 755 | config.blockDim = dim3(kThreads); |
| 756 | config.dynamicSmemBytes = 0; |
| 757 | config.stream = stream; |
| 758 | cudaLaunchAttribute attrs[1]; |
| 759 | attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 760 | attrs[0].val.programmaticStreamSerializationAllowed = 1; |
| 761 | config.attrs = attrs; |
| 762 | config.numAttrs = 1; |
| 763 | cudaLaunchKernelEx(&config, kernel, |
| 764 | reinterpret_cast<const __nv_bfloat16*>(x_q), |
| 765 | reinterpret_cast<const __nv_bfloat16*>(x_k), |
| 766 | reinterpret_cast<const __nv_bfloat16*>(x_v), |
| 767 | reinterpret_cast<const float*>(w_q_t), |
| 768 | reinterpret_cast<const float*>(w_k_t), |
| 769 | reinterpret_cast<const float*>(w_v_t), |
| 770 | reinterpret_cast<const float*>(bias_q), |
| 771 | reinterpret_cast<const float*>(bias_k), |
| 772 | reinterpret_cast<const float*>(bias_v), |
| 773 | reinterpret_cast<__nv_bfloat16*>(cs_q), |
| 774 | reinterpret_cast<__nv_bfloat16*>(cs_k), |
| 775 | reinterpret_cast<__nv_bfloat16*>(cs_v), a_log, |
| 776 | reinterpret_cast<const __nv_bfloat16*>(g), dt_bias, |
| 777 | reinterpret_cast<const __nv_bfloat16*>(beta), |
| 778 | reinterpret_cast<const __nv_bfloat16*>(onorm_g), |
| 779 | onorm_weight, ssm_state_indices, cu_seqlens, state, |
| 780 | reinterpret_cast<__nv_bfloat16*>(out), B, H, HV, |
| 781 | lower_bound, scale, onorm_eps, strides); |
| 782 | } |
| 783 | |
| 784 | template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid> |
| 785 | void launch_kda_decode_many_heads_selected( |
| 786 | const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, |
| 787 | const void* w_k_t, const void* w_v_t, const void* bias_q, |
| 788 | const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, |
| 789 | const float* a_log, const void* g, const float* dt_bias, const void* beta, |
| 790 | const void* onorm_g, const float* onorm_weight, |
| 791 | const int* ssm_state_indices, const int* cu_seqlens, float* state, |
| 792 | void* out, int B, int H, int HV, bool update_conv_cache, |
| 793 | bool conv_state_is_ds, float lower_bound, float scale, float onorm_eps, |
| 794 | KdaDecodeStrides strides, cudaStream_t stream) { |
| 795 | #define LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, CS_CHANNEL_STRIDE, CS_TAP_STRIDE) \ |
| 796 | do { \ |
| 797 | if (update_conv_cache) { \ |
| 798 | launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, true, \ |
| 799 | kUseLowerBound, kApplyBetaSigmoid, \ |
| 800 | CS_CHANNEL_STRIDE, CS_TAP_STRIDE>( \ |
| 801 | x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ |
| 802 | cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ |
| 803 | ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ |
| 804 | scale, onorm_eps, strides, stream); \ |
| 805 | } else { \ |
| 806 | launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, false, \ |
| 807 | kUseLowerBound, kApplyBetaSigmoid, \ |
| 808 | CS_CHANNEL_STRIDE, CS_TAP_STRIDE>( \ |
| 809 | x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ |
| 810 | cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ |
| 811 | ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ |
| 812 | scale, onorm_eps, strides, stream); \ |
| 813 | } \ |
| 814 | } while (false) |
| 815 | |
| 816 | // SD packs (tap, channel) per slot while DS packs (channel, tap); both are |
| 817 | // baked into the kernel as compile-time strides. |
| 818 | #define LAUNCH_KDA_DECODE(NUM_HEADS) \ |
| 819 | do { \ |
| 820 | if (conv_state_is_ds) { \ |
| 821 | LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, kConvStateWidth, 1); \ |
| 822 | } else { \ |
| 823 | LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, 1, 3 * (NUM_HEADS) * kDimK); \ |
| 824 | } \ |
| 825 | } while (false) |
| 826 | |
| 827 | switch (H) { |
| 828 | case 12: |
| 829 | LAUNCH_KDA_DECODE(12); |
| 830 | break; |
| 831 | case 24: |
| 832 | LAUNCH_KDA_DECODE(24); |
| 833 | break; |
| 834 | case 48: |
| 835 | LAUNCH_KDA_DECODE(48); |
| 836 | break; |
| 837 | case 96: |
| 838 | LAUNCH_KDA_DECODE(96); |
| 839 | break; |
| 840 | default: |
| 841 | STD_TORCH_CHECK(false, "Unsupported number of heads: ", H); |
| 842 | } |
| 843 | |
| 844 | #undef LAUNCH_KDA_DECODE |
| 845 | #undef LAUNCH_KDA_DECODE_LAYOUT |
| 846 | } |
| 847 | |
| 848 | struct KdaDecodeLaunchParams { |
| 849 | const void* x_q; |
| 850 | const void* x_k; |
| 851 | const void* x_v; |
| 852 | const void* w_q_t; |
| 853 | const void* w_k_t; |
| 854 | const void* w_v_t; |
| 855 | const void* bias_q; |
| 856 | const void* bias_k; |
| 857 | const void* bias_v; |
| 858 | void* cs_q; |
| 859 | void* cs_k; |
| 860 | void* cs_v; |
| 861 | const float* a_log; |
| 862 | const void* g; |
| 863 | const float* dt_bias; |
| 864 | const void* beta; |
| 865 | const void* onorm_g; |
| 866 | const float* onorm_weight; |
| 867 | const int* ssm_state_indices; |
| 868 | const int* cu_seqlens; |
| 869 | float* state; |
| 870 | void* out; |
| 871 | int B; |
| 872 | int H; |
| 873 | int HV; |
| 874 | bool update_conv_cache; |
| 875 | bool conv_state_is_ds; |
| 876 | float lower_bound; |
| 877 | float scale; |
| 878 | float onorm_eps; |
| 879 | KdaDecodeStrides strides; |
| 880 | cudaStream_t stream; |
| 881 | }; |
| 882 | |
| 883 | template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid> |
| 884 | void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) { |
| 885 | launch_kda_decode_many_heads_selected<kApplyOnorm, kUseLowerBound, |
| 886 | kApplyBetaSigmoid>( |
| 887 | p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, |
| 888 | p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, |
| 889 | p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state, |
| 890 | p.out, p.B, p.H, p.HV, p.update_conv_cache, p.conv_state_is_ds, |
| 891 | p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream); |
| 892 | } |
| 893 | |
| 894 | template <bool kApplyOnorm, bool kUseLowerBound> |
| 895 | void dispatch_kda_decode_beta(const KdaDecodeLaunchParams& p, |
| 896 | bool apply_beta_sigmoid) { |
| 897 | if (apply_beta_sigmoid) { |
| 898 | launch_kda_decode_selected_backend<kApplyOnorm, kUseLowerBound, true>(p); |
| 899 | } else { |
| 900 | launch_kda_decode_selected_backend<kApplyOnorm, kUseLowerBound, false>(p); |
| 901 | } |
| 902 | } |
| 903 | |
| 904 | template <bool kApplyOnorm> |
| 905 | void dispatch_kda_decode_decay(const KdaDecodeLaunchParams& p, |
| 906 | bool use_lower_bound, bool apply_beta_sigmoid) { |
| 907 | if (use_lower_bound) { |
| 908 | dispatch_kda_decode_beta<kApplyOnorm, true>(p, apply_beta_sigmoid); |
| 909 | } else { |
| 910 | dispatch_kda_decode_beta<kApplyOnorm, false>(p, apply_beta_sigmoid); |
| 911 | } |
| 912 | } |
| 913 | |
| 914 | void dispatch_kda_decode_features(const KdaDecodeLaunchParams& p, |
| 915 | bool apply_onorm, bool use_lower_bound, |
| 916 | bool apply_beta_sigmoid) { |
| 917 | if (apply_onorm) { |
| 918 | dispatch_kda_decode_decay<true>(p, use_lower_bound, apply_beta_sigmoid); |
| 919 | } else { |
| 920 | dispatch_kda_decode_decay<false>(p, use_lower_bound, apply_beta_sigmoid); |
| 921 | } |
| 922 | } |
| 923 | |
| 924 | } // namespace |
| 925 | |
| 926 | extern "C" void launch_kda_decode_many_heads_cuda( |
| 927 | const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, |
| 928 | const void* w_k_t, const void* w_v_t, const void* bias_q, |
| 929 | const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, |
| 930 | const float* a_log, const void* g, const float* dt_bias, const void* beta, |
| 931 | const void* onorm_g, const float* onorm_weight, |
| 932 | const int* ssm_state_indices, const int* cu_seqlens, float* state, |
| 933 | void* out, int B, int H, int HV, bool apply_onorm, bool update_conv_cache, |
| 934 | bool use_lower_bound, bool apply_beta_sigmoid, float lower_bound, |
| 935 | float scale, float onorm_eps, const int64_t* raw_strides, |
| 936 | cudaStream_t stream) { |
| 937 | const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2], |
| 938 | raw_strides[3], raw_strides[4]}; |
| 939 | const bool conv_state_is_ds = |
| 940 | raw_strides[5] == kConvStateWidth && raw_strides[6] == 1; |
| 941 | const KdaDecodeLaunchParams params{x_q, |
| 942 | x_k, |
| 943 | x_v, |
| 944 | w_q_t, |
| 945 | w_k_t, |
| 946 | w_v_t, |
| 947 | bias_q, |
| 948 | bias_k, |
| 949 | bias_v, |
| 950 | cs_q, |
| 951 | cs_k, |
| 952 | cs_v, |
| 953 | a_log, |
| 954 | g, |
| 955 | dt_bias, |
| 956 | beta, |
| 957 | onorm_g, |
| 958 | onorm_weight, |
| 959 | ssm_state_indices, |
| 960 | cu_seqlens, |
| 961 | state, |
| 962 | out, |
| 963 | B, |
| 964 | H, |
| 965 | HV, |
| 966 | update_conv_cache, |
| 967 | conv_state_is_ds, |
| 968 | lower_bound, |
| 969 | scale, |
| 970 | onorm_eps, |
| 971 | strides, |
| 972 | stream}; |
| 973 | dispatch_kda_decode_features(params, apply_onorm, use_lower_bound, |
| 974 | apply_beta_sigmoid); |
| 975 | } |
| 976 | |
| 977 | void fused_kda_decode( |
| 978 | torch::stable::Tensor const& x, torch::stable::Tensor const& weight, |
| 979 | std::optional<torch::stable::Tensor> bias, |
| 980 | torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, |
| 981 | torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, |
| 982 | torch::stable::Tensor const& dt_bias, |
| 983 | torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, |
| 984 | torch::stable::Tensor& out, std::optional<double> lower_bound, |
| 985 | std::optional<torch::stable::Tensor> output_gate, |
| 986 | std::optional<torch::stable::Tensor> norm_weight, double norm_eps) { |
| 987 | using torch::headeronly::ScalarType; |
| 988 | constexpr int kHeadDim = 128; |
| 989 | constexpr int kConvWidth = 4; |
| 990 | |
| 991 | STD_TORCH_CHECK(x.is_cuda() && x.scalar_type() == ScalarType::BFloat16, |
| 992 | "x must be a CUDA bfloat16 tensor"); |
| 993 | STD_TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == ScalarType::Float, |
| 994 | "weight must be a CUDA float32 tensor"); |
| 995 | STD_TORCH_CHECK( |
| 996 | conv_state.is_cuda() && conv_state.scalar_type() == ScalarType::BFloat16, |
| 997 | "conv_state must be a CUDA bfloat16 tensor"); |
| 998 | STD_TORCH_CHECK( |
| 999 | raw_g.is_cuda() && raw_g.scalar_type() == ScalarType::BFloat16, |
| 1000 | "raw_g must be a CUDA bfloat16 tensor"); |
| 1001 | STD_TORCH_CHECK( |
| 1002 | raw_beta.is_cuda() && raw_beta.scalar_type() == ScalarType::BFloat16, |
| 1003 | "raw_beta must be a CUDA bfloat16 tensor"); |
| 1004 | STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float, |
| 1005 | "A_log must be a CUDA float32 tensor"); |
| 1006 | STD_TORCH_CHECK( |
| 1007 | dt_bias.is_cuda() && dt_bias.scalar_type() == ScalarType::Float, |
| 1008 | "dt_bias must be a CUDA float32 tensor"); |
| 1009 | STD_TORCH_CHECK(state.is_cuda() && state.scalar_type() == ScalarType::Float, |
| 1010 | "state must be a CUDA float32 tensor"); |
| 1011 | STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16, |
| 1012 | "out must be a CUDA bfloat16 tensor"); |
| 1013 | STD_TORCH_CHECK( |
| 1014 | state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int, |
| 1015 | "state_indices must be a CUDA int32 tensor"); |
| 1016 | |
| 1017 | STD_TORCH_CHECK(x.dim() == 2, "x must have shape [B, 3 * H * 128]"); |
| 1018 | int const batch_size = static_cast<int>(x.size(0)); |
| 1019 | int64_t const qkv_width = x.size(1); |
| 1020 | STD_TORCH_CHECK(qkv_width % (3 * kHeadDim) == 0, |
| 1021 | "x must have shape [B, 3 * H * 128]"); |
| 1022 | int64_t const num_heads = qkv_width / (3 * kHeadDim); |
| 1023 | STD_TORCH_CHECK( |
| 1024 | num_heads == 12 || num_heads == 24 || num_heads == 48 || num_heads == 96, |
| 1025 | "H must be 12, 24, 48, or 96, got ", num_heads); |
| 1026 | STD_TORCH_CHECK(batch_size > 0, |
| 1027 | "KDA decode fusion requires at least one row"); |
| 1028 | int const dim = num_heads * kHeadDim; |
| 1029 | |
| 1030 | STD_TORCH_CHECK(weight.dim() == 3 && weight.is_contiguous() && |
| 1031 | weight.size(0) == 3 && weight.size(1) == kConvWidth && |
| 1032 | weight.size(2) == dim, |
| 1033 | "weight must have shape [3, 4, H * 128]"); |
| 1034 | STD_TORCH_CHECK(conv_state.dim() == 3 && conv_state.size(1) == 3 * dim && |
| 1035 | conv_state.size(2) == kConvWidth - 1, |
| 1036 | "conv_state must have shape [slots, 3 * H * 128, 3]"); |
| 1037 | STD_TORCH_CHECK(raw_g.dim() == 4 && raw_g.size(0) == 1 && |
| 1038 | raw_g.size(1) == batch_size && |
| 1039 | raw_g.size(2) == num_heads && raw_g.size(3) == kHeadDim, |
| 1040 | "raw_g must have shape [1, B, H, 128]"); |
| 1041 | STD_TORCH_CHECK(raw_beta.dim() == 3 && raw_beta.size(0) == 1 && |
| 1042 | raw_beta.size(1) == batch_size && |
| 1043 | raw_beta.size(2) == num_heads, |
| 1044 | "raw_beta must have shape [1, B, H]"); |
| 1045 | STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_heads, |
| 1046 | "A_log must be contiguous with H elements"); |
| 1047 | STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == dim, |
| 1048 | "dt_bias must be contiguous with H * 128 elements"); |
| 1049 | STD_TORCH_CHECK( |
| 1050 | state_indices.is_contiguous() && state_indices.numel() == batch_size, |
| 1051 | "state_indices must be contiguous with B elements"); |
| 1052 | STD_TORCH_CHECK(state.dim() == 4 && state.size(1) == num_heads && |
| 1053 | state.size(2) == kHeadDim && state.size(3) == kHeadDim, |
| 1054 | "state must have shape [slots, H, 128, 128]"); |
| 1055 | STD_TORCH_CHECK(out.dim() == 4 && out.size(0) == 1 && |
| 1056 | out.size(1) == batch_size && out.size(2) == num_heads && |
| 1057 | out.size(3) == kHeadDim, |
| 1058 | "out must have shape [1, B, H, 128]"); |
| 1059 | STD_TORCH_CHECK(x.stride(1) == 1, |
| 1060 | "x must be contiguous in its channel dimension"); |
| 1061 | // Both conv-state cache layouts are supported: SD stores (state_len, dim) |
| 1062 | // per block while DS stores (dim, state_len). The kernel only needs the |
| 1063 | // per-plane strides. |
| 1064 | STD_TORCH_CHECK( |
| 1065 | conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && |
| 1066 | ((conv_state.stride(1) == 1 && conv_state.stride(2) == 3 * dim) || |
| 1067 | (conv_state.stride(1) == kConvWidth - 1 && |
| 1068 | conv_state.stride(2) == 1)), |
| 1069 | "conv_state must use the SD or DS cache layout"); |
| 1070 | STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim && |
| 1071 | state.stride(1) == kHeadDim * kHeadDim && |
| 1072 | state.stride(2) == kHeadDim && state.stride(3) == 1, |
| 1073 | "state must have contiguous [H, 128, 128] slot contents"); |
| 1074 | STD_TORCH_CHECK(raw_g.is_contiguous(), "raw_g must be contiguous"); |
| 1075 | STD_TORCH_CHECK(raw_beta.stride(2) == 1, |
| 1076 | "raw_beta must be contiguous in its head dimension"); |
| 1077 | STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); |
| 1078 | |
| 1079 | bool const apply_onorm = output_gate.has_value(); |
| 1080 | STD_TORCH_CHECK(apply_onorm == norm_weight.has_value(), |
| 1081 | "output_gate and norm_weight must be provided together"); |
| 1082 | void const* output_gate_ptr = nullptr; |
| 1083 | float const* norm_weight_ptr = nullptr; |
| 1084 | int64_t output_gate_row_stride = 0; |
| 1085 | if (apply_onorm) { |
| 1086 | STD_TORCH_CHECK(output_gate->is_cuda() && |
| 1087 | output_gate->scalar_type() == ScalarType::BFloat16, |
| 1088 | "output_gate must be a CUDA bfloat16 tensor"); |
| 1089 | bool const gate_is_3d = |
| 1090 | output_gate->dim() == 3 && output_gate->size(0) == batch_size && |
| 1091 | output_gate->size(1) == num_heads && output_gate->size(2) == kHeadDim; |
| 1092 | bool const gate_is_4d = |
| 1093 | output_gate->dim() == 4 && output_gate->size(0) == 1 && |
| 1094 | output_gate->size(1) == batch_size && |
| 1095 | output_gate->size(2) == num_heads && output_gate->size(3) == kHeadDim; |
| 1096 | STD_TORCH_CHECK(gate_is_3d || gate_is_4d, |
| 1097 | "output_gate must have shape [B, H, 128] or " |
| 1098 | "[1, B, H, 128]"); |
| 1099 | int const row_dim = gate_is_3d ? 0 : 1; |
| 1100 | STD_TORCH_CHECK(output_gate->stride(output_gate->dim() - 1) == 1, |
| 1101 | "output_gate must be contiguous in its last dimension"); |
| 1102 | STD_TORCH_CHECK(output_gate->stride(row_dim + 1) == kHeadDim, |
| 1103 | "output_gate must have contiguous head rows"); |
| 1104 | STD_TORCH_CHECK(norm_weight->is_cuda() && |
| 1105 | norm_weight->scalar_type() == ScalarType::Float, |
| 1106 | "norm_weight must be a CUDA float32 tensor"); |
| 1107 | STD_TORCH_CHECK( |
| 1108 | norm_weight->is_contiguous() && norm_weight->numel() == kHeadDim, |
| 1109 | "norm_weight must be contiguous with 128 elements"); |
| 1110 | STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative"); |
| 1111 | output_gate_ptr = output_gate->data_ptr(); |
| 1112 | norm_weight_ptr = static_cast<float const*>(norm_weight->data_ptr()); |
| 1113 | output_gate_row_stride = output_gate->stride(row_dim); |
| 1114 | } |
| 1115 | |
| 1116 | void const* bias_ptr = nullptr; |
| 1117 | if (bias.has_value()) { |
| 1118 | STD_TORCH_CHECK(bias->is_cuda() && bias->scalar_type() == ScalarType::Float, |
| 1119 | "bias must be a CUDA float32 tensor"); |
| 1120 | STD_TORCH_CHECK(bias->is_contiguous() && bias->numel() == 3 * dim, |
| 1121 | "bias must be contiguous with 3 * H * 128 elements"); |
| 1122 | bias_ptr = bias->data_ptr(); |
| 1123 | } |
| 1124 | |
| 1125 | auto const* x_ptr = static_cast<char const*>(x.data_ptr()); |
| 1126 | auto const* weight_ptr = static_cast<char const*>(weight.data_ptr()); |
| 1127 | auto* conv_ptr = static_cast<char*>(conv_state.data_ptr()); |
| 1128 | auto const* bias_bytes = static_cast<char const*>(bias_ptr); |
| 1129 | int64_t const segment_bytes = dim * sizeof(__nv_bfloat16); |
| 1130 | int64_t const weight_segment_bytes = |
| 1131 | dim * kConvWidth * static_cast<int64_t>(sizeof(float)); |
| 1132 | int64_t const conv_segment_bytes = |
| 1133 | dim * conv_state.stride(1) * sizeof(__nv_bfloat16); |
| 1134 | int64_t const bias_segment_bytes = dim * sizeof(float); |
| 1135 | std::array<int64_t, 7> const strides{ |
| 1136 | x.stride(0), raw_beta.stride(1), output_gate_row_stride, |
| 1137 | conv_state.stride(0), state.stride(0), conv_state.stride(1), |
| 1138 | conv_state.stride(2), |
| 1139 | }; |
| 1140 | bool const use_lower_bound = lower_bound.has_value(); |
| 1141 | float const lower_bound_value = |
| 1142 | use_lower_bound ? static_cast<float>(*lower_bound) : 0.0f; |
| 1143 | |
| 1144 | torch::stable::accelerator::DeviceGuard const device_guard( |
| 1145 | x.get_device_index()); |
| 1146 | cudaStream_t const stream = get_current_cuda_stream(x.get_device_index()); |
| 1147 | launch_kda_decode_many_heads_cuda( |
| 1148 | x_ptr, x_ptr + segment_bytes, x_ptr + 2 * segment_bytes, weight_ptr, |
| 1149 | weight_ptr + weight_segment_bytes, weight_ptr + 2 * weight_segment_bytes, |
| 1150 | bias_bytes, |
| 1151 | bias_bytes == nullptr ? nullptr : bias_bytes + bias_segment_bytes, |
| 1152 | bias_bytes == nullptr ? nullptr : bias_bytes + 2 * bias_segment_bytes, |
| 1153 | conv_ptr, conv_ptr + conv_segment_bytes, |
| 1154 | conv_ptr + 2 * conv_segment_bytes, |
| 1155 | static_cast<float const*>(a_log.data_ptr()), raw_g.data_ptr(), |
| 1156 | static_cast<float const*>(dt_bias.data_ptr()), raw_beta.data_ptr(), |
| 1157 | output_gate_ptr, norm_weight_ptr, |
| 1158 | static_cast<int const*>(state_indices.data_ptr()), nullptr, |
| 1159 | static_cast<float*>(state.data_ptr()), out.data_ptr(), batch_size, |
| 1160 | num_heads, num_heads, apply_onorm, true, use_lower_bound, true, |
| 1161 | lower_bound_value, 0.08838834764831845f, static_cast<float>(norm_eps), |
| 1162 | strides.data(), stream); |
| 1163 | cudaError_t const error = cudaGetLastError(); |
| 1164 | STD_TORCH_CHECK( |
| 1165 | error == cudaSuccess, |
| 1166 | "Kimi K3 KDA decode kernel launch failed: ", cudaGetErrorString(error)); |
| 1167 | } |