[Kernel][Perf] Support DSpark K=8 in fused GDN MTP decode - #53070
Draft
BabyDrangoner wants to merge 1 commit into
Draft
[Kernel][Perf] Support DSpark K=8 in fused GDN MTP decode#53070BabyDrangoner wants to merge 1 commit into
BabyDrangoner wants to merge 1 commit into
Conversation
Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
#52539 enabled the production Qwen GDN head ratios, but the fused MTP
post-convolution kernel still accepted at most eight verification tokens.
Qwen3.6 DSpark K=8 verifies
K + 1 = 9tokens, so all 30 GDN layers fellback to the Triton sequence.
This PR adds a width-9 specialization. Widths up to eight keep the existing
MaxTokens=8specialization; only the new path allocates the ninth token'sshared storage, and warp 0 handles that token in the two token-parallel
phases. The public operator ABI and build configuration are unchanged.
This is a narrow follow-up to #52539. A search of open vLLM PRs found no
other width-9 implementation.
Test Plan
python -m pytest -q tests/kernels/test_fused_gdn_post_conv.py python -m pytest -q tests/kernels/mamba/test_gdn_fused_mtp.py python -m pytest -q \ 'tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py::test_dspark_correctness_and_acceptance_rate[qwen3.6-speculators]' pre-commit run --files \ csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu \ vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py \ tests/kernels/test_fused_gdn_post_conv.py \ tests/kernels/mamba/test_gdn_fused_mtp.pyTest Result
accuracy
0.875, acceptance rate0.454, and acceptance length4.630.On one RTX PRO 6000 Blackwell (SM120), a fixed-input Qwen3.6 TP1 width-9
post-convolution benchmark used 60 paired ABBA samples per batch and mode:
Maximum output/full-state relative L2 was
4.174e-5 / 5.521e-8. Repeatedruns were deterministic, graph replay matched eager bitwise within each arm,
and the null state slot was unchanged. Width-2/8 outputs and full state were
bitwise equal to the base implementation; the extracted
MaxTokens=8device code was also byte-identical on the audited CUDA targets. These are
custom out-of-tree kernel-boundary measurements.
The official FP8 checkpoint was additionally tested with DSpark K=8 at TP1,
eager mode, greedy 64-token generation, and 15 clean single-load ABBA rounds:
All 300 measured requests had identical prompt IDs, output IDs, and
speculative work across arms, with zero cached tokens. At B2/B8/B16, final
output IDs also matched but speculative work differed, so those timings were
intentionally excluded.
AI assistance and human review
OpenAI Codex assisted with implementation, tests, experiments, and this
description.
changed line and confirm they can explain, test, and maintain it.
Essential Elements of an Effective PR Description Checklist