[Kimi K3][Kernel] Support DS conv-state layout in fused KDA decode kernel - #53396
Conversation
The fused KDA decode CUDA kernel previously required the SD conv-state cache layout, so deployments that pin VLLM_SSM_CONV_STATE_LAYOUT=DS (e.g. for NIXL P/D state transfer) silently fell back to the 4-op Triton decode chain for all 69 KDA layers. Make the conv-state inner strides launch parameters instead of compile time constants so both layouts run through the single fused kernel, and enable the fused path when the dim-first (DS) layout is active. Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
gcanlin
changed the title
|
Documentation preview: https://vllm--53396.org.readthedocs.build/en/53396/ |
4e3c60e to
6057745
Compare
6057745 to
bb09534
Compare
gcanlin
changed the title
|
Do you have microbenchmark results to see if this change causes any regression to the current SD codepath? |
Oh, it indeed causes regression in small batch, but not sure why have improvement for batch 256, 512. The root cause may be that this PR changes |
|
yea the strides can be constexpr, should be fine |
Dispatch SD and DS conv-state layouts to separate kernel instantiations so their address calculations remain compile-time constants. Extend the kernel benchmark to select either physical layout. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
|
@gau-nernst Thanks for suggestion! Now it doesn't cause any regression for SD layout and the new implementation is better than the old one for DS layout(I have updated the new data in PR description), so I run E2E benchmark again to see the benefits. @tjtanaa @AndreasKaratzas Could you please take a look for the ROCm part? If it's not suitable for ROCm, I can revert it. |
| bool kUseActiveOnormReduction = false, bool kUpdateConvState = false, | ||
| bool kUseLowerBound = false, bool kApplyBetaSigmoid = true> | ||
| bool kUseLowerBound = false, bool kApplyBetaSigmoid = true, | ||
| ConvStateLayout kConvStateLayout = ConvStateLayout::kSD> |
There was a problem hiding this comment.
personal style but why not just introduce 2 int strides as template parameters instead of ConvStateLayout enum?
There was a problem hiding this comment.
I asked codex to try replacing ConvStateLayout with two integer stride NTTPs. On B200, the stride-NTTP version regressed CUDA-graph replay latency by 18.3% for DS and 7.2% for SD at H_local=12 and B=128. 🧐 haven't found the root cause. Perhaps it changed the unrelated part. I get about 3-4% e2e improvement in the currently new version so prefer to keep it.
There was a problem hiding this comment.
That is very strange. Are you sure the two integer strides are template parameters? Look like it's even slower than dynamic strides? Can you double check? Thank you.
There was a problem hiding this comment.
K3 helped me push one new expected change. Now it's equivalent to the enum version.
There was a problem hiding this comment.
K3 is better than Codex? 🤯
There was a problem hiding this comment.
For this case I think it is. Or maybe I haven't given codex enough room to shine. 🤣
Replace the ConvStateLayout enum template parameter with two integer
stride NTTPs (channel/tap) per review feedback. The layout is still
detected at launch from conv_state strides and resolved inside the
head-count dispatch, so both strides stay compile-time constants:
SD = (1, 3 * heads * 128), DS = (kConvWidth - 1, 1).
Verified on B200: all 128 kda_decode instantiations emit SASS identical
to the enum version; CUDA-graph replay latency is unchanged (max |delta|
0.6% at H_local=12 over B in {1, 8, 32, 64, 128}, both layouts);
test_fused_kda_decode_correctness passes 10/10 for SD and DS.
Co-authored-by: Kimi Code <noreply@moonshot.cn>
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
| template <int kConvStateChannelStride, int kConvStateTapStride> | ||
| __device__ __forceinline__ int conv_state_offset(int channel, int tap) { | ||
| return channel * kConvStateChannelStride + tap * kConvStateTapStride; | ||
| } |
There was a problem hiding this comment.
again just style but can you inline this? thank you
Drop the conv_state_offset helper and write the stride arithmetic directly; generated SASS stays bit-identical to the enum version for all 128 instantiations, and test_fused_kda_decode_correctness passes 10/10 for SD and DS. Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #85738 for commit |
|
✅ @gcanlin, CI is now available for this PR.
|
|
@gcanlin Can you rebase onto main? The arm64 docker CI should be fixed by a recent commit iirc. |
|
/ci run |
|
✅ Triggered Buildkite CI #85744 for commit |
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> Signed-off-by: khushali9 <khushali.desai9@gmail.com>
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> Signed-off-by: mikeshawcode <michaelwshaw2@gmail.com>
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> Signed-off-by: mikeshawcode <michaelwshaw2@gmail.com>
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
…rnel (vllm-project#53396) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
4 participants
Summary
The fused KDA decode kernel (
ops.fused_kda_decode) previously required the SD conv-state cache layout, so any deployment that pinsVLLM_SSM_CONV_STATE_LAYOUT=DS— required for NIXL point-to-point state transfer in P/D-disaggregated serving — silently bypassed the fused path and fell back to the multi-op Triton decode chain for all 69 KDA layers of the model.This PR supports both physical cache layouts by dispatching to separate SD- and DS-specialized fused-kernel instantiations. Keeping each layout's address calculation compile-time constant avoids the extra integer instructions and register pressure caused by runtime inner strides. Both layouts now use the fused path (in-proj GEMM → fused KDA decode → out-proj GEMM per layer and step), and numerics are bit-identical for identical inputs.
Update (head
0e67480f): per review feedback, the layout specialization is now expressed as two integer conv-state inner strides (channel, tap) as non-type template parameters instead of a layout enum. The generated device code is unchanged — see the parity data below.Motivation
VLLM_SSM_CONV_STATE_LAYOUT=DS(NIXL PD transfer), and several recent fixes treat DS as the supported production layout. Under DS, the fused kernel was never used.Performance
DS specialized fused path vs Triton fallback
B200, CUDA graph replay,
H_local=12,K=V=128, 50 warmup iterations, and 300 measured iterations. Values are the average of two median runs. The fallback uses the Triton gated RMSNorm kernel rather than its eager PyTorch decomposition.SD regression and DS specialization check
B200, CUDA graph replay,
H_local=12,K=V=128, 50 warmup iterations, and 300 measured iterations. Values are the average of two median runs with reversed execution order, using binaries produced by the same build tree.Compile-time specialization fully recovers the original SD performance and improves DS performance by up to 10.95%. For the representative H_local=12 kernel, runtime strides increase register use from 64 to 72 registers per thread, while both specialized variants use 64. The additional kernel instances increase the extension binary by 392 KB (0.95%).
Stride-NTTP parity check (review feedback, head
0e67480f)The layout enum was replaced by two integer strides as template parameters, per review feedback. The layout is still detected at launch from the conv_state strides and resolved inside the head-count dispatch (SD =
<1, 3 * H * 128>, DS =<kConvWidth - 1, 1>), so both strides remain compile-time constants.Codegen parity proof: a normalized
cuobjdumpcomparison of all 128kda_decode_fusion_many_heads_kernelinstantiations in_C_stable_libtorchshows bit-identical SASS between the enum version and this rewrite (identical REG/STACK/SHARED, zero spills), so measured performance is identical by construction.B200, CUDA graph replay,
H_local=12,K=V=128, 50 warmup / 300 measured iterations, fused-path per-layer latency:Production-shaped check (69 per-layer state buffer sets, B=128, DS): 37.04 us/layer fused vs 49.33 us/layer Triton chain (1.33x).
An earlier stride-NTTP attempt was measured to regress 7-18%, but that attempt bundled unrelated changes; this clean rewrite matches the enum version exactly.
E2E: 2 nodes x 8xB200, TP8+PP2, FP8 KV cache, DS layout,
FULL_DECODE_ONLYcudagraphsvllm bench serve, random dataset, 128 input / 512 output tokens,128 requests, max-concurrency 64, temperature 0, and no profiler. Each side used
one full-load pilot followed by two measured runs. Both sides used the same vLLM
nightly image (
sha256:95bed119f39e...); the candidate overlaid this PR atcommit
40459de9d3. The current head0e67480femits bit-identical kernel SASS (verified above), so these numbers carry over; a fresh E2E run on it will follow.Every measured run completed all 128 requests with zero failures. The
output-throughput spread was 0.46% for the baseline and 0.53% for this PR, so
the 3.84% improvement is well above the observed run-to-run variation. Both
nodes logged
Fused KDA decode kernel (conv+KDA+norm) is enabled.for thecandidate, with no errors or runtime JIT-compilation warnings.
Acc:
Kernel trace verification
With the DS layout pinned, the patched serving stack runs
kda_decode_fusion_many_heads_kernelfor every KDA decode layer (5076 calls in the profiled window, 7.2us/call average, 1.26% of window kernel time); the decode-side_causal_conv1d_update_kernel,fused_recurrent_kda_packed_decodetriton kernels and the separate gated-norm kernel disappear from the decode hot path (prefill-only_causal_conv1d_fwd_kerneland FlashKDA prefill remain, as intended). In the baseline trace the fallback trio costs ~3.5% of kernel time.Test plan
10 passed, 51 deselected; the targeted cases cover both SD and DS layouts, H=12/24/48/96, output normalization, and lower-bound decay. Re-run at head0e67480f(stride-NTTP rewrite):10 passed, 51 deselected.Server (as used for the E2E numbers above):