diff --git a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py index 6c0dfe1b8d..fcce640152 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py +++ b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py @@ -81,7 +81,7 @@ def _bench_graph(fn, repeats: int = NUM_KDA_LAYERS) -> float: class Inputs: - def __init__(self, num_tokens: int, num_heads: int) -> None: + def __init__(self, num_tokens: int, num_heads: int, conv_layout: str) -> None: torch.manual_seed(0) device = "cuda" dim = num_heads * HEAD_DIM @@ -98,9 +98,15 @@ class Inputs: for i in range(3) ] ) - self.conv_state = torch.randn( - num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE - ) + if conv_layout == "DS": + self.conv_state_t = torch.randn( + num_slots, 3 * dim, CONV_WIDTH - 1, device=device, dtype=DTYPE + ) + else: + conv_state = torch.randn( + num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE + ) + self.conv_state_t = conv_state.transpose(-1, -2) self.recurrent_state = torch.randn( num_slots, num_heads, HEAD_DIM, HEAD_DIM, device=device, dtype=torch.float32 ) @@ -120,7 +126,6 @@ class Inputs: self.state_indices = torch.arange( 1, num_tokens + 1, device=device, dtype=torch.int32 ) - self.conv_state_t = self.conv_state.transpose(-1, -2) self.out = torch.empty( 1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE ) @@ -237,6 +242,12 @@ def main() -> None: default=[12], help="KDA heads per rank (96 total / TP size)", ) + parser.add_argument( + "--layout", + choices=("SD", "DS"), + default="SD", + help="physical conv-state cache layout", + ) args = parser.parse_args() if not hasattr(torch.ops._C, "fused_kda_decode"): @@ -245,7 +256,10 @@ def main() -> None: props = torch.cuda.get_device_properties(0) bench = _bench_graph if args.graph else _bench mode = "cuda-graph replay" if args.graph else "eager dispatch" - print(f"device: {props.name} ({props.gcnArchName}) timing: {mode}") + print( + f"device: {props.name} ({props.gcnArchName}) timing: {mode} " + f"layout: {args.layout}" + ) print( f"{'heads':>6} {'tokens':>7} {'eager-norm':>11} {'triton-norm':>12} " f"{'fused us':>9} {'speedup':>8} {'state TB/s':>11} {'saved ms/step':>14}" @@ -253,7 +267,10 @@ def main() -> None: for num_heads in args.heads: for num_tokens in args.tokens: if args.graph and args.layers > 1: - layers = [Inputs(num_tokens, num_heads) for _ in range(args.layers)] + layers = [ + Inputs(num_tokens, num_heads, args.layout) + for _ in range(args.layers) + ] eager_ms = _bench_graph_layers( [functools.partial(triton_chain, i, False) for i in layers] ) @@ -267,7 +284,7 @@ def main() -> None: del layers torch.accelerator.empty_cache() else: - inp = Inputs(num_tokens, num_heads) + inp = Inputs(num_tokens, num_heads, args.layout) eager_ms = bench(functools.partial(triton_chain, inp, False)) triton_ms = bench(functools.partial(triton_chain, inp, True)) fused_ms = bench(functools.partial(fused, inp)) diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu index 0badeb9b77..00e5ce794d 100644 --- a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu +++ b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu @@ -261,15 +261,16 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, return {scratch[0], scratch[1]}; } -template +template < + bool kApplyOnorm, bool kUseStaticDecodeLayout = false, int kFixedHeads = 0, + int kFixedValueHeads = 0, bool kUseHeadGrid = false, + bool kAccumulateOnormSumsq = false, bool kUseActiveQkReduction = false, + bool kUseCacheGlobalStore = false, bool kComputeOutputBeforeStore = false, + bool kSkipWarpSync = false, bool kPreloadOnormParams = false, + bool kPrefetchNextStateChunk = false, bool kUseActiveOnormReduction = false, + bool kUpdateConvState = false, bool kUseLowerBound = false, + bool kApplyBetaSigmoid = true, int kConvStateChannelStride = 1, + int kConvStateTapStride = 3 * kFixedHeads * kDimK> __global__ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( const __nv_bfloat16* __restrict__ x_q, @@ -329,7 +330,6 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( } constexpr int kLocalDim = kFixedHeads * kDimK; - constexpr int kPackedDim = 3 * kLocalDim; const int hk_off = i_h * kDimK; const int hv_off = i_hv * kDimV; constexpr int hv_count = kFixedValueHeads; @@ -368,8 +368,10 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim]; - const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim]; + const int cs_idx = + hk * kConvStateChannelStride + w * kConvStateTapStride; + const __nv_bfloat16 q_state = cs_q_for_slot[cs_idx]; + const __nv_bfloat16 k_state = cs_k_for_slot[cs_idx]; q_acc += __bfloat162float(q_state) * conv_weight_load(w_q_t, hk, w); k_acc += __bfloat162float(k_state) * @@ -389,12 +391,16 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( k_acc += __bfloat162float(k_new) * conv_weight_load(w_k_t, hk, kKernelWidth - 1); - cs_q_for_slot[hk] = q_shift0; - cs_q_for_slot[hk + kPackedDim] = q_shift1; - cs_q_for_slot[hk + 2 * kPackedDim] = q_new; - cs_k_for_slot[hk] = k_shift0; - cs_k_for_slot[hk + kPackedDim] = k_shift1; - cs_k_for_slot[hk + 2 * kPackedDim] = k_new; + cs_q_for_slot[hk * kConvStateChannelStride] = q_shift0; + cs_q_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] = + q_shift1; + cs_q_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] = + q_new; + cs_k_for_slot[hk * kConvStateChannelStride] = k_shift0; + cs_k_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] = + k_shift1; + cs_k_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] = + k_new; s_q[k] = silu_fast(q_acc); s_k[k] = silu_fast(k_acc); @@ -418,7 +424,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const int cs_idx = hk + w * kPackedDim; + const int cs_idx = + hk * kConvStateChannelStride + w * kConvStateTapStride; q_acc += bf16_load(cs_q_for_slot, cs_idx) * conv_weight_load(w_q_t, hk, w); 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( __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim]; + const int cs_idx = + hvv * kConvStateChannelStride + w * kConvStateTapStride; + const __nv_bfloat16 v_state = cs_v_for_slot[cs_idx]; v_acc += __bfloat162float(v_state) * conv_weight_load(w_v_t, hvv, w); if (w == 1) { @@ -465,9 +474,11 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( const __nv_bfloat16 v_new = x_v[xv_idx]; v_acc += __bfloat162float(v_new) * conv_weight_load(w_v_t, hvv, kKernelWidth - 1); - cs_v_for_slot[hvv] = v_shift0; - cs_v_for_slot[hvv + kPackedDim] = v_shift1; - cs_v_for_slot[hvv + 2 * kPackedDim] = v_new; + cs_v_for_slot[hvv * kConvStateChannelStride] = v_shift0; + cs_v_for_slot[hvv * kConvStateChannelStride + kConvStateTapStride] = + v_shift1; + cs_v_for_slot[hvv * kConvStateChannelStride + 2 * kConvStateTapStride] = + v_new; s_v[v] = silu_fast(v_acc); if constexpr (kApplyOnorm && kPreloadOnormParams) { @@ -484,7 +495,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const int cs_idx = hvv + w * kPackedDim; + const int cs_idx = + hvv * kConvStateChannelStride + w * kConvStateTapStride; v_acc += bf16_load(cs_v_for_slot, cs_idx) * conv_weight_load(w_v_t, hvv, w); } @@ -723,7 +735,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( } template + bool kUseLowerBound, bool kApplyBetaSigmoid, + int kConvStateChannelStride, int kConvStateTapStride> void launch_kda_decode_many_heads_raw( const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, 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( auto kernel = &kda_decode_fusion_many_heads_kernel< kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false, false, true, true, true, kUpdateConvState, kUseLowerBound, - kApplyBetaSigmoid>; + kApplyBetaSigmoid, kConvStateChannelStride, kConvStateTapStride>; cudaLaunchConfig_t config{}; config.gridDim = dim3(B, kHeads); config.blockDim = dim3(kThreads); @@ -776,26 +789,39 @@ void launch_kda_decode_many_heads_selected( const float* a_log, const void* g, const float* dt_bias, const void* beta, const void* onorm_g, const float* onorm_weight, const int* ssm_state_indices, const int* cu_seqlens, float* state, - void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound, - float scale, float onorm_eps, KdaDecodeStrides strides, - cudaStream_t stream) { -#define LAUNCH_KDA_DECODE(NUM_HEADS) \ - do { \ - if (update_conv_cache) { \ - launch_kda_decode_many_heads_raw( \ - x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ - cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ - ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ - scale, onorm_eps, strides, stream); \ - } else { \ - launch_kda_decode_many_heads_raw( \ - x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ - cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ - ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ - scale, onorm_eps, strides, stream); \ - } \ + void* out, int B, int H, int HV, bool update_conv_cache, + bool conv_state_is_ds, float lower_bound, float scale, float onorm_eps, + KdaDecodeStrides strides, cudaStream_t stream) { +#define LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, CS_CHANNEL_STRIDE, CS_TAP_STRIDE) \ + do { \ + if (update_conv_cache) { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } else { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } \ + } while (false) + + // SD packs (tap, channel) per slot while DS packs (channel, tap); both are + // baked into the kernel as compile-time strides. +#define LAUNCH_KDA_DECODE(NUM_HEADS) \ + do { \ + if (conv_state_is_ds) { \ + LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, kConvStateWidth, 1); \ + } else { \ + LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, 1, 3 * (NUM_HEADS) * kDimK); \ + } \ } while (false) switch (H) { @@ -816,6 +842,7 @@ void launch_kda_decode_many_heads_selected( } #undef LAUNCH_KDA_DECODE +#undef LAUNCH_KDA_DECODE_LAYOUT } struct KdaDecodeLaunchParams { @@ -845,6 +872,7 @@ struct KdaDecodeLaunchParams { int H; int HV; bool update_conv_cache; + bool conv_state_is_ds; float lower_bound; float scale; float onorm_eps; @@ -859,8 +887,8 @@ void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) { 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, p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state, - p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale, - p.onorm_eps, p.strides, p.stream); + p.out, p.B, p.H, p.HV, p.update_conv_cache, p.conv_state_is_ds, + p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream); } template @@ -908,6 +936,8 @@ extern "C" void launch_kda_decode_many_heads_cuda( cudaStream_t stream) { const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2], raw_strides[3], raw_strides[4]}; + const bool conv_state_is_ds = + raw_strides[5] == kConvStateWidth && raw_strides[6] == 1; const KdaDecodeLaunchParams params{x_q, x_k, x_v, @@ -934,6 +964,7 @@ extern "C" void launch_kda_decode_many_heads_cuda( H, HV, update_conv_cache, + conv_state_is_ds, lower_bound, scale, onorm_eps, @@ -1027,10 +1058,15 @@ void fused_kda_decode( "out must have shape [1, B, H, 128]"); STD_TORCH_CHECK(x.stride(1) == 1, "x must be contiguous in its channel dimension"); - STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && - conv_state.stride(1) == 1 && - conv_state.stride(2) == 3 * dim, - "conv_state must use the SD cache layout"); + // Both conv-state cache layouts are supported: SD stores (state_len, dim) + // per block while DS stores (dim, state_len). The kernel only needs the + // per-plane strides. + STD_TORCH_CHECK( + conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && + ((conv_state.stride(1) == 1 && conv_state.stride(2) == 3 * dim) || + (conv_state.stride(1) == kConvWidth - 1 && + conv_state.stride(2) == 1)), + "conv_state must use the SD or DS cache layout"); STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim && state.stride(1) == kHeadDim * kHeadDim && state.stride(2) == kHeadDim && state.stride(3) == 1, @@ -1096,9 +1132,10 @@ void fused_kda_decode( int64_t const conv_segment_bytes = dim * conv_state.stride(1) * sizeof(__nv_bfloat16); int64_t const bias_segment_bytes = dim * sizeof(float); - std::array const strides{ + std::array const strides{ x.stride(0), raw_beta.stride(1), output_gate_row_stride, - conv_state.stride(0), state.stride(0), + conv_state.stride(0), state.stride(0), conv_state.stride(1), + conv_state.stride(2), }; bool const use_lower_bound = lower_bound.has_value(); float const lower_bound_value = diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu index 6711685809..800bb85c93 100644 --- a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu +++ b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu @@ -147,7 +147,8 @@ __device__ __forceinline__ Sum3 block_reduce_sum3(float x, float y, float z, scratch[3 * kGroups + 2]}; } -template __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel( 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( } return; } - constexpr int kPackedDim = 3 * kLocalDim; const int hk_off = i_h * kDimK; const int hv_off = i_hv * kDimV; float* const state_for_slot = state + slot * strides.state_slot; @@ -242,8 +242,10 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel( bf16_t k_shift1 = __float2bfloat16(0.0f); #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const bf16_t q_state = cs_q_for_slot[hk + w * kPackedDim]; - const bf16_t k_state = cs_k_for_slot[hk + w * kPackedDim]; + const int cs_idx = + hk * kConvStateChannelStride + w * kConvStateTapStride; + const bf16_t q_state = cs_q_for_slot[cs_idx]; + const bf16_t k_state = cs_k_for_slot[cs_idx]; q_acc += __bfloat162float(q_state) * conv_weight_load(w_q_t, hk, w); k_acc += __bfloat162float(k_state) * @@ -264,16 +266,21 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel( k_acc += __bfloat162float(k_new) * conv_weight_load(w_k_t, hk, kKernelWidth - 1); - cs_q_for_slot[hk] = q_shift0; - cs_q_for_slot[hk + kPackedDim] = q_shift1; - cs_q_for_slot[hk + 2 * kPackedDim] = q_new; - cs_k_for_slot[hk] = k_shift0; - cs_k_for_slot[hk + kPackedDim] = k_shift1; - cs_k_for_slot[hk + 2 * kPackedDim] = k_new; + cs_q_for_slot[hk * kConvStateChannelStride] = q_shift0; + cs_q_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] = + q_shift1; + cs_q_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] = + q_new; + cs_k_for_slot[hk * kConvStateChannelStride] = k_shift0; + cs_k_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] = + k_shift1; + cs_k_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] = + k_new; } else { #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const int cs_idx = hk + w * kPackedDim; + const int cs_idx = + hk * kConvStateChannelStride + w * kConvStateTapStride; q_acc += bf16_load(cs_q_for_slot, cs_idx) * conv_weight_load(w_q_t, hk, w); k_acc += bf16_load(cs_k_for_slot, cs_idx) * @@ -307,7 +314,9 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel( bf16_t v_shift1 = __float2bfloat16(0.0f); #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const bf16_t v_state = cs_v_for_slot[hvv + w * kPackedDim]; + const int cs_idx = + hvv * kConvStateChannelStride + w * kConvStateTapStride; + const bf16_t v_state = cs_v_for_slot[cs_idx]; v_acc += __bfloat162float(v_state) * conv_weight_load(w_v_t, hvv, w); if (w == 1) { @@ -320,13 +329,16 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel( const bf16_t v_new = x_v[xv_idx]; v_acc += __bfloat162float(v_new) * conv_weight_load(w_v_t, hvv, kKernelWidth - 1); - cs_v_for_slot[hvv] = v_shift0; - cs_v_for_slot[hvv + kPackedDim] = v_shift1; - cs_v_for_slot[hvv + 2 * kPackedDim] = v_new; + cs_v_for_slot[hvv * kConvStateChannelStride] = v_shift0; + cs_v_for_slot[hvv * kConvStateChannelStride + kConvStateTapStride] = + v_shift1; + cs_v_for_slot[hvv * kConvStateChannelStride + 2 * kConvStateTapStride] = + v_new; } else { #pragma unroll for (int w = 0; w < kConvStateWidth; ++w) { - const int cs_idx = hvv + w * kPackedDim; + const int cs_idx = + hvv * kConvStateChannelStride + w * kConvStateTapStride; v_acc += bf16_load(cs_v_for_slot, cs_idx) * conv_weight_load(w_v_t, hvv, w); } @@ -461,7 +473,8 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel( } template + bool kUseLowerBound, bool kApplyBetaSigmoid, + int kConvStateChannelStride, int kConvStateTapStride> void launch_kda_decode_raw( const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, const void* w_k_t, const void* w_v_t, const void* bias_q, @@ -471,7 +484,8 @@ void launch_kda_decode_raw( const int* ssm_state_indices, float* state, void* out, int B, float lower_bound, float scale, float onorm_eps, KdaDecodeStrides strides, hipStream_t stream) { - kda_decode_fusion_kernel <<>>( reinterpret_cast(x_q), @@ -517,6 +531,7 @@ struct KdaDecodeLaunchParams { int B; int H; bool update_conv_cache; + bool conv_state_is_ds; float lower_bound; float scale; float onorm_eps; @@ -527,19 +542,29 @@ struct KdaDecodeLaunchParams { template void dispatch_kda_decode_conv(const KdaDecodeLaunchParams& p) { -#define LAUNCH_KDA_DECODE(UPDATE_CONV) \ - launch_kda_decode_raw( \ - 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, \ - p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \ - p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \ +#define LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, CS_CHANNEL_STRIDE, \ + CS_TAP_STRIDE) \ + launch_kda_decode_raw( \ + 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, \ + p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \ + p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \ p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream) + // SD packs (tap, channel) per slot while DS packs (channel, tap); both are + // baked into the kernel as compile-time strides. +#define LAUNCH_KDA_DECODE(UPDATE_CONV) \ + if (p.conv_state_is_ds) { \ + LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, kConvStateWidth, 1); \ + } else { \ + LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, 1, 3 * kHeads * kDimK); \ + } if (p.update_conv_cache) { LAUNCH_KDA_DECODE(true); } else { LAUNCH_KDA_DECODE(false); } #undef LAUNCH_KDA_DECODE +#undef LAUNCH_KDA_DECODE_LAYOUT } template @@ -686,10 +711,12 @@ void fused_kda_decode( "out must have shape [1, B, H, 128]"); STD_TORCH_CHECK(x.stride(1) == 1, "x must be contiguous in its channel dimension"); - STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && - conv_state.stride(1) == 1 && - conv_state.stride(2) == 3 * dim, - "conv_state must use the SD cache layout"); + STD_TORCH_CHECK( + conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && + ((conv_state.stride(1) == 1 && conv_state.stride(2) == 3 * dim) || + (conv_state.stride(1) == kConvWidth - 1 && + conv_state.stride(2) == 1)), + "conv_state must use the SD or DS cache layout"); STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim && state.stride(1) == kHeadDim * kHeadDim && state.stride(2) == kHeadDim && state.stride(3) == 1, @@ -761,6 +788,8 @@ void fused_kda_decode( torch::stable::accelerator::DeviceGuard const device_guard( x.get_device_index()); + bool const conv_state_is_ds = + conv_state.stride(1) == kConvStateWidth && conv_state.stride(2) == 1; KdaDecodeLaunchParams const params{ x_ptr, x_ptr + segment_bytes, @@ -786,6 +815,7 @@ void fused_kda_decode( batch_size, static_cast(num_heads), true, + conv_state_is_ds, lower_bound_value, 0.08838834764831845f, static_cast(norm_eps), diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 3aba8aa3bc..e822ca919e 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -895,13 +895,18 @@ def test_kda_recoverssm_verify_and_group_commit( @pytest.mark.parametrize( - ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"), + ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm", "conv_layout"), [ - (12, 1, -5.0, True), - (12, 4, None, False), - (24, 4, None, False), - (48, 1, -5.0, True), - (96, 1, -5.0, True), + (12, 1, -5.0, True, "SD"), + (12, 4, None, False, "SD"), + (24, 4, None, False, "SD"), + (48, 1, -5.0, True, "SD"), + (96, 1, -5.0, True, "SD"), + (12, 1, -5.0, True, "DS"), + (12, 4, None, False, "DS"), + (24, 4, None, False, "DS"), + (48, 1, -5.0, True, "DS"), + (96, 1, -5.0, True, "DS"), ], ) @torch.inference_mode() @@ -910,6 +915,7 @@ def test_fused_kda_decode_correctness( num_seqs: int, lower_bound: float | None, fuse_output_norm: bool, + conv_layout: str, ): D, W = 128, 4 if not is_fused_kda_decode_supported( @@ -921,7 +927,7 @@ def test_fused_kda_decode_correctness( conv_state_dtype=torch.bfloat16, ): pytest.skip("Fused KDA decode is not supported on this platform") - torch.manual_seed(967 + num_heads + num_seqs) + torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS")) dim = num_heads * D slots = num_seqs + 2 packed_x_storage = torch.randn( @@ -929,13 +935,25 @@ def test_fused_kda_decode_correctness( ) packed_x = packed_x_storage[:, : 3 * dim] weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE) - conv_seed = 0.1 * torch.randn( - slots, - W - 1, - 3 * dim, - dtype=torch.bfloat16, - device=DEVICE, - ).transpose(1, 2) + if conv_layout == "DS": + # DS cache layout: per slot the taps are innermost + # (stride (W-1, 1)), matching VLLM_SSM_CONV_STATE_LAYOUT=DS. + conv_seed = 0.1 * torch.randn( + slots, + 3 * dim, + W - 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + else: + # SD cache layout: per slot the channels are innermost. + conv_seed = 0.1 * torch.randn( + slots, + W - 1, + 3 * dim, + dtype=torch.bfloat16, + device=DEVICE, + ).transpose(1, 2) raw_g = torch.randn( 1, num_seqs, @@ -1017,7 +1035,11 @@ def test_fused_kda_decode_correctness( conv_actual = torch.as_strided( cache_storage.view(torch.bfloat16), size=(slots, 3 * dim, W - 1), - stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim), + stride=( + page_bytes // torch.bfloat16.itemsize, + (W - 1) if conv_layout == "DS" else 1, + 1 if conv_layout == "DS" else 3 * dim, + ), ) state_actual = torch.as_strided( cache_storage.view(torch.float32), diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 9347751a25..83460967a6 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -153,6 +153,8 @@ def is_fused_kda_decode_supported( input_dtype: torch.dtype, conv_state_dtype: torch.dtype, ) -> bool: + # The fused kernel handles both conv-state cache layouts (SD and DS); the + # inner strides are selected from the tensor at launch time. if ( num_heads not in (12, 24, 48, 96) or head_dim != 128 @@ -160,7 +162,6 @@ def is_fused_kda_decode_supported( or num_spec != 0 or input_dtype != torch.bfloat16 or conv_state_dtype != torch.bfloat16 - or is_conv_state_dim_first() or not hasattr(torch.ops._C, "fused_kda_decode") ): return False