Skip to content

[Kernel][Perf] Support DSpark K=8 in fused GDN MTP decode - #53070

Draft
BabyDrangoner wants to merge 1 commit into
vllm-project:mainfrom
BabyDrangoner:perf/qwen36-gdn-dspark-width9
Draft

[Kernel][Perf] Support DSpark K=8 in fused GDN MTP decode#53070
BabyDrangoner wants to merge 1 commit into
vllm-project:mainfrom
BabyDrangoner:perf/qwen36-gdn-dspark-width9

Conversation

@BabyDrangoner

@BabyDrangoner BabyDrangoner commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

Draft: final human line-by-line review is pending.

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 = 9 tokens, so all 30 GDN layers fell
back to the Triton sequence.

This PR adds a width-9 specialization. Widths up to eight keep the existing
MaxTokens=8 specialization; only the new path allocates the ninth token's
shared 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.py

Test Result

  • Low-level post-convolution tests: 77 passed.
  • GDN MTP model-path tests: 16 passed.
  • Full pre-commit on all four changed files: passed.
  • The official 200-question NVFP4 DSpark test: passed, with GSM8K
    accuracy 0.875, acceptance rate 0.454, and acceptance length 4.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:

Mode Paired-median speedup Paired outcomes
Eager 1.017x-1.644x 240 wins, 0 ties, 0 losses
CUDA Graph 1.009x-1.359x 239 wins, 1 tie, 0 losses

Maximum output/full-state relative L2 was 4.174e-5 / 5.521e-8. Repeated
runs 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=8
device 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:

Batch Wall speedup TPOT speedup Wall / TPOT wins
1 1.121x 1.135x 15/15 / 15/15
4 1.119x 1.130x 15/15 / 15/15

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.

  • Before marking this PR ready, the human submitter will review every
    changed line and confirm they can explain, test, and maintain it.

Essential Elements of an Effective PR Description Checklist
  • The purpose and scope are described.
  • Exact test commands and results are included.
  • No documentation update is required for this internal kernel extension.

Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

1 participant