[SM120] Use exact query-head widths for DeepSeek-V4 sparse MLA decode - #36655
Conversation
|
/rerun-failed-ci |
Additional accuracy qualification: AIME25 repeat-8I ran a matched baseline/candidate AIME25 repeat-8 evaluation with the official Setup:
Command: sgl-eval run aime25 \
--base-url http://127.0.0.1:PORT/v1 \
--load-preset-from-model-id deepseek-ai/DeepSeek-V4-Flash-0731 \
--n-repeats 8 --num-threads 64 \
--out-dir OUTPUT_DIR
All five main misses were 200K-token no-answer truncations. The patch had four such truncations plus one normal-stop stochastic miss (AIME25 problem 14, repeat 3: As a diagnostic only, the patch generated 0.42% more completion tokens while finishing 8.14% sooner, for 1,110.69 versus 1,015.97 output tok/s (+9.32%). The dedicated frozen decode-only A/B in the PR body remains the performance authority. Both accepted runs completed all eight 30-record shards with zero HTTP/evaluator errors. Archived server logs contain no traceback, OOM, runtime/assertion failure, unsupported sparse-MLA shape, or request retraction. |
Hi @Fridge003 Baizhou, I discussed with the PR author @AliceChenyy before and the conclusion is the two PRs should be complementary, #29927 is about removing padding for the prefill kernel while this PR is for the decode kernel, Alice can comment on this too @AliceChenyy . |
|
Please resolve conflicts |
…ecode Port of upstream PR sgl-project#36655 to the ormandj-v0.8.1-rc.10 stack. On SM120 with the FlashInfer backend, TP>1 decode currently pads each rank's query heads to 64 before dispatching the sparse MLA kernel, making the kx kernel process four times as many heads (e.g. 16 real heads become 64 on TP4). Recent FlashInfer DSV4 decode dispatchers support the native per-rank widths. Select the native width only when the installed FlashInfer dispatcher explicitly supports the active decode shape (capability discovery fails closed), and keep 64-head padding for prefill-sized batches, unsupported head counts, and other backends. Retain the branch's existing SM120 prefill behavior (no padding above the 64-token decode cutoff). Keep one fallback-width attention-sink allocation and return exact-width views from it so prefill/decode alternation and CUDA graph capture do not reallocate or invalidate a captured pointer. Preserve the no-argument _local_attn_sink() contract used by DSpark. Adds registered unit coverage for head selection and sink-view stability, plus an SM120 exact-vs-padded kernel equivalence test.
|
Please solve conflicts |
4f5a119 to
8729eba
Compare
Skip unsupported 8-head and 16-head subtests independently so a supported width retains its numerical coverage. Keep the numerical operations and tolerances unchanged.
|
@Fridge003 conflict resolved and the unittest removed, extended the numerical test a little bit, let me know if you have more comments, THX a lot! |
|
/rerun-test test/registered/kernels/ops/attention/test_flash_mla_backends.py test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200.py test/registered/e2e/models/test_deepseek_v4_flash_fp4_megamoe_b200.py test/registered/e2e/models/test_deepseek_v4_flash_fp4_h200.py |
|
Results for 🚀 🚀 🚀 |
3 participants
Motivation
DeepSeek-V4 on SM120 pads TP4 decode queries from 16 real heads to 64 even
though FlashInfer supports the native width. This makes sparse attention
process four times as many query heads. This change selects the native width
for supported FlashInfer decode shapes and supplies a matching attention sink.
The native-width prefill path introduced by #29927 is preserved. This PR adds
the corresponding optimization for decode-sized batches.
Modifications
get_platform().is_sm120for the platform check.heads for supported batches up to the decode token limit.
padding for unsupported decode shapes/backends.
stable pointers across shape transitions and DSpark's no-argument contract.
for H8/H16, T=1/64/65, and single/dual cache inputs.
Validation of the current revision
Revision
8729ebafc1778481e64507bc1f057cba4c300715is rebased ontoc0b790cf7fe6dc1516c3ec1f23de9b37baca654a.On an exclusively held SM120 GPU, using the revised SGLang source with
FlashInfer 0.6.18 and PyTorch 2.13.0+cu130:
All 12 native-versus-padded comparisons passed at
atol=rtol=5e-2, includingthe decode/prefill boundary and attention sinks. The existing skipped test
probes a legacy FlashInfer module name; the new test exercised the actual
FlashInfer implementation. GPU ownership remained exclusive throughout.
All applicable pre-commit hooks passed on the three changed files, including
formatting, lint, AST parsing, and CI registration checks.
Historical serving and performance validation
The following results were measured on the earlier PR revision
4f5a119dc1against base
8005df61d3. They have not been rerun for the current rebase.Official
sgl-project/sgl-evalAIME25 repeat-8, using the same four RTX PRO6000 GPUs and model/server settings, produced 235/240 (97.92%) pass@1,
100% pass@8, and 100% majority@8 on both variants, with zero request errors.
See the full accuracy qualification.
Unprofiled decode-only A/B used DeepSeek-V4-Flash, pure TP4, FP8 KV, page size
256, FlashInfer sparse MLA and
flashinfer_mxfp4MoE, speculative decodingdisabled, and CUDA graphs/autotuning enabled. Each case had one full-output
warmup wave, then five timed requests per concurrency slot with 1,024 output
tokens and byte-identical prompts across variants.
Every point completed its exact request/token totals, reached its declared
concurrency, and had zero request errors. These are one-repeat synthetic-KV
decode diagnostics, not PD or full-capacity measurements.
Checklist
CI States
Latest PR Test (Base): 🚫 Run #34521154609
Latest PR Test (Extra): 🚫 Run #34521154456
Latest PR Test (AMD ROCm 10): 🚫 Run #34521154825