@tduncan / Nvidia-Optimization-Stack-VLLM / commits / aa640684cc

[Kimi K3][Kernel] Support DS conv-state layout in fused KDA decode kernel (#53396)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>

Canlin Guo committed Aug 27, 2026 at 14:54 UTC aa640684cc95f1e4b14fdf38298ea8577a70de56
5 files changed +213 -106
benchmarks/kernels/benchmark_kimi_k3_kda_decode.py
+25 -8
@@ -81,7 +81,7 @@ def _bench_graph(fn, repeats: int = NUM_KDA_LAYERS) -> float:
81
82
83 class Inputs:
84 - def __init__(self, num_tokens: int, num_heads: int) -> None:
84 + def __init__(self, num_tokens: int, num_heads: int, conv_layout: str) -> None:
85 torch.manual_seed(0)
86 device = "cuda"
87 dim = num_heads * HEAD_DIM
@@ -98,9 +98,15 @@ class Inputs:
98 for i in range(3)
99 ]
100 )
101 - self.conv_state = torch.randn(
102 - num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE
103 - )
101 + if conv_layout == "DS":
102 + self.conv_state_t = torch.randn(
103 + num_slots, 3 * dim, CONV_WIDTH - 1, device=device, dtype=DTYPE
104 + )
105 + else:
106 + conv_state = torch.randn(
107 + num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE
108 + )
109 + self.conv_state_t = conv_state.transpose(-1, -2)
110 self.recurrent_state = torch.randn(
111 num_slots, num_heads, HEAD_DIM, HEAD_DIM, device=device, dtype=torch.float32
112 )
@@ -120,7 +126,6 @@ class Inputs:
126 self.state_indices = torch.arange(
127 1, num_tokens + 1, device=device, dtype=torch.int32
128 )
123 - self.conv_state_t = self.conv_state.transpose(-1, -2)
129 self.out = torch.empty(
130 1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE
131 )
@@ -237,6 +242,12 @@ def main() -> None:
242 default=[12],
243 help="KDA heads per rank (96 total / TP size)",
244 )
245 + parser.add_argument(
246 + "--layout",
247 + choices=("SD", "DS"),
248 + default="SD",
249 + help="physical conv-state cache layout",
250 + )
251 args = parser.parse_args()
252
253 if not hasattr(torch.ops._C, "fused_kda_decode"):
@@ -245,7 +256,10 @@ def main() -> None:
256 props = torch.cuda.get_device_properties(0)
257 bench = _bench_graph if args.graph else _bench
258 mode = "cuda-graph replay" if args.graph else "eager dispatch"
248 - print(f"device: {props.name} ({props.gcnArchName}) timing: {mode}")
259 + print(
260 + f"device: {props.name} ({props.gcnArchName}) timing: {mode} "
261 + f"layout: {args.layout}"
262 + )
263 print(
264 f"{'heads':>6} {'tokens':>7} {'eager-norm':>11} {'triton-norm':>12} "
265 f"{'fused us':>9} {'speedup':>8} {'state TB/s':>11} {'saved ms/step':>14}"
@@ -253,7 +267,10 @@ def main() -> None:
267 for num_heads in args.heads:
268 for num_tokens in args.tokens:
269 if args.graph and args.layers > 1:
256 - layers = [Inputs(num_tokens, num_heads) for _ in range(args.layers)]
270 + layers = [
271 + Inputs(num_tokens, num_heads, args.layout)
272 + for _ in range(args.layers)
273 + ]
274 eager_ms = _bench_graph_layers(
275 [functools.partial(triton_chain, i, False) for i in layers]
276 )
@@ -267,7 +284,7 @@ def main() -> None:
284 del layers
285 torch.accelerator.empty_cache()
286 else:
270 - inp = Inputs(num_tokens, num_heads)
287 + inp = Inputs(num_tokens, num_heads, args.layout)
288 eager_ms = bench(functools.partial(triton_chain, inp, False))
289 triton_ms = bench(functools.partial(triton_chain, inp, True))
290 fused_ms = bench(functools.partial(fused, inp))
csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu
+91 -54
@@ -261,15 +261,16 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y,
261 return {scratch[0], scratch[1]};
262 }
263
264 -template <bool kApplyOnorm, bool kUseStaticDecodeLayout = false,
265 - int kFixedHeads = 0, int kFixedValueHeads = 0,
266 - bool kUseHeadGrid = false, bool kAccumulateOnormSumsq = false,
267 - bool kUseActiveQkReduction = false, bool kUseCacheGlobalStore = false,
268 - bool kComputeOutputBeforeStore = false, bool kSkipWarpSync = false,
269 - bool kPreloadOnormParams = false,
270 - bool kPrefetchNextStateChunk = false,
271 - bool kUseActiveOnormReduction = false, bool kUpdateConvState = false,
272 - bool kUseLowerBound = false, bool kApplyBetaSigmoid = true>
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,
@@ -329,7 +330,6 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
330 }
331
332 constexpr int kLocalDim = kFixedHeads * kDimK;
332 - constexpr int kPackedDim = 3 * kLocalDim;
333 const int hk_off = i_h * kDimK;
334 const int hv_off = i_hv * kDimV;
335 constexpr int hv_count = kFixedValueHeads;
@@ -368,8 +368,10 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
368 __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f);
369 #pragma unroll
370 for (int w = 0; w < kConvStateWidth; ++w) {
371 - const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim];
372 - const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim];
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) *
@@ -389,12 +391,16 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
391 k_acc += __bfloat162float(k_new) *
392 conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
393
392 - cs_q_for_slot[hk] = q_shift0;
393 - cs_q_for_slot[hk + kPackedDim] = q_shift1;
394 - cs_q_for_slot[hk + 2 * kPackedDim] = q_new;
395 - cs_k_for_slot[hk] = k_shift0;
396 - cs_k_for_slot[hk + kPackedDim] = k_shift1;
397 - cs_k_for_slot[hk + 2 * kPackedDim] = k_new;
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);
@@ -418,7 +424,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
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) {
421 - const int cs_idx = hk + w * kPackedDim;
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) *
@@ -453,7 +460,9 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
460 __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f);
461 #pragma unroll
462 for (int w = 0; w < kConvStateWidth; ++w) {
456 - const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim];
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) {
@@ -465,9 +474,11 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
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);
468 - cs_v_for_slot[hvv] = v_shift0;
469 - cs_v_for_slot[hvv + kPackedDim] = v_shift1;
470 - cs_v_for_slot[hvv + 2 * kPackedDim] = v_new;
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) {
@@ -484,7 +495,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
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) {
487 - const int cs_idx = hvv + w * kPackedDim;
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 }
@@ -723,7 +735,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
735 }
736
737 template <int kHeads, bool kApplyOnorm, bool kUpdateConvState,
726 - bool kUseLowerBound, bool kApplyBetaSigmoid>
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,
@@ -736,7 +749,7 @@ void launch_kda_decode_many_heads_raw(
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,
739 - kApplyBetaSigmoid>;
752 + kApplyBetaSigmoid, kConvStateChannelStride, kConvStateTapStride>;
753 cudaLaunchConfig_t config{};
754 config.gridDim = dim3(B, kHeads);
755 config.blockDim = dim3(kThreads);
@@ -776,26 +789,39 @@ void launch_kda_decode_many_heads_selected(
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,
779 - void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound,
780 - float scale, float onorm_eps, KdaDecodeStrides strides,
781 - cudaStream_t stream) {
782 -#define LAUNCH_KDA_DECODE(NUM_HEADS) \
783 - do { \
784 - if (update_conv_cache) { \
785 - launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, true, \
786 - kUseLowerBound, kApplyBetaSigmoid>( \
787 - x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
788 - cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
789 - ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
790 - scale, onorm_eps, strides, stream); \
791 - } else { \
792 - launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, false, \
793 - kUseLowerBound, kApplyBetaSigmoid>( \
794 - x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
795 - cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
796 - ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
797 - scale, onorm_eps, strides, stream); \
798 - } \
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) {
@@ -816,6 +842,7 @@ void launch_kda_decode_many_heads_selected(
842 }
843
844 #undef LAUNCH_KDA_DECODE
845 +#undef LAUNCH_KDA_DECODE_LAYOUT
846 }
847
848 struct KdaDecodeLaunchParams {
@@ -845,6 +872,7 @@ struct KdaDecodeLaunchParams {
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;
@@ -859,8 +887,8 @@ void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) {
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,
862 - p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale,
863 - p.onorm_eps, p.strides, p.stream);
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>
@@ -908,6 +936,8 @@ extern "C" void launch_kda_decode_many_heads_cuda(
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,
@@ -934,6 +964,7 @@ extern "C" void launch_kda_decode_many_heads_cuda(
964 H,
965 HV,
966 update_conv_cache,
967 + conv_state_is_ds,
968 lower_bound,
969 scale,
970 onorm_eps,
@@ -1027,10 +1058,15 @@ void fused_kda_decode(
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");
1030 - STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
1031 - conv_state.stride(1) == 1 &&
1032 - conv_state.stride(2) == 3 * dim,
1033 - "conv_state must use the SD cache layout");
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,
@@ -1096,9 +1132,10 @@ void fused_kda_decode(
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);
1099 - std::array<int64_t, 5> const strides{
1135 + std::array<int64_t, 7> const strides{
1136 x.stride(0), raw_beta.stride(1), output_gate_row_stride,
1101 - conv_state.stride(0), state.stride(0),
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 =
csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu
+58 -28
@@ -147,7 +147,8 @@ __device__ __forceinline__ Sum3 block_reduce_sum3(float x, float y, float z,
147 scratch[3 * kGroups + 2]};
148 }
149
150 -template <bool kApplyOnorm, int kFixedHeads, bool kUpdateConvState,
150 +template <bool kApplyOnorm, int kConvStateChannelStride,
151 + int kConvStateTapStride, int kFixedHeads, bool kUpdateConvState,
152 bool kUseLowerBound, bool kApplyBetaSigmoid>
153 __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
154 const bf16_t* __restrict__ x_q, const bf16_t* __restrict__ x_k,
@@ -186,7 +187,6 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
187 }
188 return;
189 }
189 - constexpr int kPackedDim = 3 * kLocalDim;
190 const int hk_off = i_h * kDimK;
191 const int hv_off = i_hv * kDimV;
192 float* const state_for_slot = state + slot * strides.state_slot;
@@ -242,8 +242,10 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
242 bf16_t k_shift1 = __float2bfloat16(0.0f);
243 #pragma unroll
244 for (int w = 0; w < kConvStateWidth; ++w) {
245 - const bf16_t q_state = cs_q_for_slot[hk + w * kPackedDim];
246 - const bf16_t k_state = cs_k_for_slot[hk + w * kPackedDim];
245 + const int cs_idx =
246 + hk * kConvStateChannelStride + w * kConvStateTapStride;
247 + const bf16_t q_state = cs_q_for_slot[cs_idx];
248 + const bf16_t k_state = cs_k_for_slot[cs_idx];
249 q_acc += __bfloat162float(q_state) *
250 conv_weight_load<kLocalDim>(w_q_t, hk, w);
251 k_acc += __bfloat162float(k_state) *
@@ -264,16 +266,21 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
266 k_acc += __bfloat162float(k_new) *
267 conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
268
267 - cs_q_for_slot[hk] = q_shift0;
268 - cs_q_for_slot[hk + kPackedDim] = q_shift1;
269 - cs_q_for_slot[hk + 2 * kPackedDim] = q_new;
270 - cs_k_for_slot[hk] = k_shift0;
271 - cs_k_for_slot[hk + kPackedDim] = k_shift1;
272 - cs_k_for_slot[hk + 2 * kPackedDim] = k_new;
269 + cs_q_for_slot[hk * kConvStateChannelStride] = q_shift0;
270 + cs_q_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] =
271 + q_shift1;
272 + cs_q_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] =
273 + q_new;
274 + cs_k_for_slot[hk * kConvStateChannelStride] = k_shift0;
275 + cs_k_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] =
276 + k_shift1;
277 + cs_k_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] =
278 + k_new;
279 } else {
280 #pragma unroll
281 for (int w = 0; w < kConvStateWidth; ++w) {
276 - const int cs_idx = hk + w * kPackedDim;
282 + const int cs_idx =
283 + hk * kConvStateChannelStride + w * kConvStateTapStride;
284 q_acc += bf16_load(cs_q_for_slot, cs_idx) *
285 conv_weight_load<kLocalDim>(w_q_t, hk, w);
286 k_acc += bf16_load(cs_k_for_slot, cs_idx) *
@@ -307,7 +314,9 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
314 bf16_t v_shift1 = __float2bfloat16(0.0f);
315 #pragma unroll
316 for (int w = 0; w < kConvStateWidth; ++w) {
310 - const bf16_t v_state = cs_v_for_slot[hvv + w * kPackedDim];
317 + const int cs_idx =
318 + hvv * kConvStateChannelStride + w * kConvStateTapStride;
319 + const bf16_t v_state = cs_v_for_slot[cs_idx];
320 v_acc += __bfloat162float(v_state) *
321 conv_weight_load<kLocalDim>(w_v_t, hvv, w);
322 if (w == 1) {
@@ -320,13 +329,16 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
329 const bf16_t v_new = x_v[xv_idx];
330 v_acc += __bfloat162float(v_new) *
331 conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
323 - cs_v_for_slot[hvv] = v_shift0;
324 - cs_v_for_slot[hvv + kPackedDim] = v_shift1;
325 - cs_v_for_slot[hvv + 2 * kPackedDim] = v_new;
332 + cs_v_for_slot[hvv * kConvStateChannelStride] = v_shift0;
333 + cs_v_for_slot[hvv * kConvStateChannelStride + kConvStateTapStride] =
334 + v_shift1;
335 + cs_v_for_slot[hvv * kConvStateChannelStride + 2 * kConvStateTapStride] =
336 + v_new;
337 } else {
338 #pragma unroll
339 for (int w = 0; w < kConvStateWidth; ++w) {
329 - const int cs_idx = hvv + w * kPackedDim;
340 + const int cs_idx =
341 + hvv * kConvStateChannelStride + w * kConvStateTapStride;
342 v_acc += bf16_load(cs_v_for_slot, cs_idx) *
343 conv_weight_load<kLocalDim>(w_v_t, hvv, w);
344 }
@@ -461,7 +473,8 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
473 }
474
475 template <int kHeads, bool kApplyOnorm, bool kUpdateConvState,
464 - bool kUseLowerBound, bool kApplyBetaSigmoid>
476 + bool kUseLowerBound, bool kApplyBetaSigmoid,
477 + int kConvStateChannelStride, int kConvStateTapStride>
478 void launch_kda_decode_raw(
479 const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
480 const void* w_k_t, const void* w_v_t, const void* bias_q,
@@ -471,7 +484,8 @@ void launch_kda_decode_raw(
484 const int* ssm_state_indices, float* state, void* out, int B,
485 float lower_bound, float scale, float onorm_eps, KdaDecodeStrides strides,
486 hipStream_t stream) {
474 - kda_decode_fusion_kernel<kApplyOnorm, kHeads, kUpdateConvState,
487 + kda_decode_fusion_kernel<kApplyOnorm, kConvStateChannelStride,
488 + kConvStateTapStride, kHeads, kUpdateConvState,
489 kUseLowerBound, kApplyBetaSigmoid>
490 <<<dim3(kHeads, B), dim3(kThreads), 0, stream>>>(
491 reinterpret_cast<const bf16_t*>(x_q),
@@ -517,6 +531,7 @@ struct KdaDecodeLaunchParams {
531 int B;
532 int H;
533 bool update_conv_cache;
534 + bool conv_state_is_ds;
535 float lower_bound;
536 float scale;
537 float onorm_eps;
@@ -527,19 +542,29 @@ struct KdaDecodeLaunchParams {
542 template <int kHeads, bool kApplyOnorm, bool kUseLowerBound,
543 bool kApplyBetaSigmoid>
544 void dispatch_kda_decode_conv(const KdaDecodeLaunchParams& p) {
530 -#define LAUNCH_KDA_DECODE(UPDATE_CONV) \
531 - launch_kda_decode_raw<kHeads, kApplyOnorm, UPDATE_CONV, kUseLowerBound, \
532 - kApplyBetaSigmoid>( \
533 - 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, \
534 - p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \
535 - p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \
545 +#define LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, CS_CHANNEL_STRIDE, \
546 + CS_TAP_STRIDE) \
547 + launch_kda_decode_raw<kHeads, kApplyOnorm, UPDATE_CONV, kUseLowerBound, \
548 + kApplyBetaSigmoid, CS_CHANNEL_STRIDE, CS_TAP_STRIDE>( \
549 + 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, \
550 + p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \
551 + p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \
552 p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream)
553 + // SD packs (tap, channel) per slot while DS packs (channel, tap); both are
554 + // baked into the kernel as compile-time strides.
555 +#define LAUNCH_KDA_DECODE(UPDATE_CONV) \
556 + if (p.conv_state_is_ds) { \
557 + LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, kConvStateWidth, 1); \
558 + } else { \
559 + LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, 1, 3 * kHeads * kDimK); \
560 + }
561 if (p.update_conv_cache) {
562 LAUNCH_KDA_DECODE(true);
563 } else {
564 LAUNCH_KDA_DECODE(false);
565 }
566 #undef LAUNCH_KDA_DECODE
567 +#undef LAUNCH_KDA_DECODE_LAYOUT
568 }
569
570 template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid>
@@ -686,10 +711,12 @@ void fused_kda_decode(
711 "out must have shape [1, B, H, 128]");
712 STD_TORCH_CHECK(x.stride(1) == 1,
713 "x must be contiguous in its channel dimension");
689 - STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
690 - conv_state.stride(1) == 1 &&
691 - conv_state.stride(2) == 3 * dim,
692 - "conv_state must use the SD cache layout");
714 + STD_TORCH_CHECK(
715 + conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
716 + ((conv_state.stride(1) == 1 && conv_state.stride(2) == 3 * dim) ||
717 + (conv_state.stride(1) == kConvWidth - 1 &&
718 + conv_state.stride(2) == 1)),
719 + "conv_state must use the SD or DS cache layout");
720 STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim &&
721 state.stride(1) == kHeadDim * kHeadDim &&
722 state.stride(2) == kHeadDim && state.stride(3) == 1,
@@ -761,6 +788,8 @@ void fused_kda_decode(
788
789 torch::stable::accelerator::DeviceGuard const device_guard(
790 x.get_device_index());
791 + bool const conv_state_is_ds =
792 + conv_state.stride(1) == kConvStateWidth && conv_state.stride(2) == 1;
793 KdaDecodeLaunchParams const params{
794 x_ptr,
795 x_ptr + segment_bytes,
@@ -786,6 +815,7 @@ void fused_kda_decode(
815 batch_size,
816 static_cast<int>(num_heads),
817 true,
818 + conv_state_is_ds,
819 lower_bound_value,
820 0.08838834764831845f,
821 static_cast<float>(norm_eps),
tests/models/kimi_k3/test_kda.py
+37 -15
@@ -895,13 +895,18 @@ def test_kda_recoverssm_verify_and_group_commit(
895
896
897 @pytest.mark.parametrize(
898 - ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"),
898 + ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm", "conv_layout"),
899 [
900 - (12, 1, -5.0, True),
901 - (12, 4, None, False),
902 - (24, 4, None, False),
903 - (48, 1, -5.0, True),
904 - (96, 1, -5.0, True),
900 + (12, 1, -5.0, True, "SD"),
901 + (12, 4, None, False, "SD"),
902 + (24, 4, None, False, "SD"),
903 + (48, 1, -5.0, True, "SD"),
904 + (96, 1, -5.0, True, "SD"),
905 + (12, 1, -5.0, True, "DS"),
906 + (12, 4, None, False, "DS"),
907 + (24, 4, None, False, "DS"),
908 + (48, 1, -5.0, True, "DS"),
909 + (96, 1, -5.0, True, "DS"),
910 ],
911 )
912 @torch.inference_mode()
@@ -910,6 +915,7 @@ def test_fused_kda_decode_correctness(
915 num_seqs: int,
916 lower_bound: float | None,
917 fuse_output_norm: bool,
918 + conv_layout: str,
919 ):
920 D, W = 128, 4
921 if not is_fused_kda_decode_supported(
@@ -921,7 +927,7 @@ def test_fused_kda_decode_correctness(
927 conv_state_dtype=torch.bfloat16,
928 ):
929 pytest.skip("Fused KDA decode is not supported on this platform")
924 - torch.manual_seed(967 + num_heads + num_seqs)
930 + torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS"))
931 dim = num_heads * D
932 slots = num_seqs + 2
933 packed_x_storage = torch.randn(
@@ -929,13 +935,25 @@ def test_fused_kda_decode_correctness(
935 )
936 packed_x = packed_x_storage[:, : 3 * dim]
937 weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE)
932 - conv_seed = 0.1 * torch.randn(
933 - slots,
934 - W - 1,
935 - 3 * dim,
936 - dtype=torch.bfloat16,
937 - device=DEVICE,
938 - ).transpose(1, 2)
938 + if conv_layout == "DS":
939 + # DS cache layout: per slot the taps are innermost
940 + # (stride (W-1, 1)), matching VLLM_SSM_CONV_STATE_LAYOUT=DS.
941 + conv_seed = 0.1 * torch.randn(
942 + slots,
943 + 3 * dim,
944 + W - 1,
945 + dtype=torch.bfloat16,
946 + device=DEVICE,
947 + )
948 + else:
949 + # SD cache layout: per slot the channels are innermost.
950 + conv_seed = 0.1 * torch.randn(
951 + slots,
952 + W - 1,
953 + 3 * dim,
954 + dtype=torch.bfloat16,
955 + device=DEVICE,
956 + ).transpose(1, 2)
957 raw_g = torch.randn(
958 1,
959 num_seqs,
@@ -1017,7 +1035,11 @@ def test_fused_kda_decode_correctness(
1035 conv_actual = torch.as_strided(
1036 cache_storage.view(torch.bfloat16),
1037 size=(slots, 3 * dim, W - 1),
1020 - stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim),
1038 + stride=(
1039 + page_bytes // torch.bfloat16.itemsize,
1040 + (W - 1) if conv_layout == "DS" else 1,
1041 + 1 if conv_layout == "DS" else 3 * dim,
1042 + ),
1043 )
1044 state_actual = torch.as_strided(
1045 cache_storage.view(torch.float32),
vllm/models/kimi_k3/nvidia/kda.py
+2 -1
@@ -153,6 +153,8 @@ def is_fused_kda_decode_supported(
153 input_dtype: torch.dtype,
154 conv_state_dtype: torch.dtype,
155 ) -> bool:
156 + # The fused kernel handles both conv-state cache layouts (SD and DS); the
157 + # inner strides are selected from the tensor at launch time.
158 if (
159 num_heads not in (12, 24, 48, 96)
160 or head_dim != 128
@@ -160,7 +162,6 @@ def is_fused_kda_decode_supported(
162 or num_spec != 0
163 or input_dtype != torch.bfloat16
164 or conv_state_dtype != torch.bfloat16
163 - or is_conv_state_dim_first()
165 or not hasattr(torch.ops._C, "fused_kda_decode")
166 ):
167 return False