fix(sparse-mla-sm120): add TOPK=256 DSV4 prefill instantiation - #3834
waynehacking8 wants to merge 1 commit into
Conversation
dispatch_dsv4_single only instantiated topk in {128, 512, 1024, 2048};
topk=256 (a valid 4x BI(64) shape, used by DeepSeek-V4 DSpark configs)
fell through to the unsupported-configuration error in the prefill
orchestrator. Add the 256 branch with BF16 compute mode: measured on an
RTX PRO 6000 (SM120), BF16 beats FP8 at topk=256 for every tested shape
(nh 32/128 x nt 128/1024/4096, 2.5-9% faster), matching the existing
small-K-loop comment - the FP8 Q-quantize prologue doesn't amortise at
K=256. Complements PR flashinfer-ai#3817, which adds the decode half of flashinfer-ai#3828.
Test configs (16,256) and (32,256) cover the NHG=1 and NHG=2 routes;
full tests/attention/test_sparse_mla_sm120.py passes (145 tests).
AI-assisted.
Signed-off-by: waynehacking8 <waynehacking8@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for a topk value of 256 using the BF16 data type in the sparse MLA prefill CUDA kernel dispatch logic. It also updates the test configurations in test_sparse_mla_sm120.py to include cases with topk=256 (specifically, (16, 256) and (32, 256)). There are no review comments, and I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughAdds a new dispatch branch for topk == 256 in the DSV4 single-cache prefill dispatcher, routing it through the existing BF16 DISPATCH_BY_NH_CM macro. Corresponding test parametrization is extended with (16, 256) and (32, 256) configuration tuples. Changestopk=256 Dispatch Support
Estimated code review effort: 1 (Trivial) | ~3 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Tested together with #3817 on SM120 (details in my comment there: 4 quality runs of a 35-conversation agent workload with per-turn degeneration scanning, 0 degenerate turns / 0 parser failures, decode benched cc1-64, several days of serving). Both needed as a pair for dspark on these cards. |
|
Thanks for running this at that scale -- a multi-day serving soak with per-turn degeneration scanning is exactly the evidence this pair needed on real SM120 hardware. Heads up that OsakaTX's #3896 implements the same dispatch branch with FP8 compute mode (per the in-source comment convention); this PR picked BF16 CM from measurement on RTX PRO 6000 (2.5-9% faster than FP8 at all six shapes I benched, numbers in the PR body), so whichever the maintainers take, the CM choice is worth a deliberate look. |
|
Reverified the carried TOPK=256 prefill instantiation against current vLLM main DeepSeek-V4-Flash-DSpark, 2x RTX PRO 6000 Max-Q (SM120), TP2. Exact server-tokenizer targets: 8,194 tokens in 0.990 s (8,274 tok/s, n=10); 65,538 in 7.956 s (8,238 tok/s, n=2); 131,074 in 17.461 s (7,507 tok/s, n=1). One 960,164-token prompt also completed without error. Benchmark SHA256: Current |
|
Confirming this on GB10 (sm_121a) as well. I benchmarked the BF16 dispatch here against an FP8 variant of the same |
|
Thanks for closing the loop with the FP8 comparison on GB10 - good to have the CM choice backed by measurements on both sm_120 and sm_121 now. |
|
Head |
|
This PR is superseded by and will be closed after #4380 |
## 📌 Description This PR consolidates the community work needed to run DeepSeek-V4-Flash-0731 with DSpark through the SM120/SM121 sparse MLA backend. A downstream DSpark configuration has 133 active sparse entries per query: 128 sliding-window entries plus 5 draft entries. The index tensor is padded to 192 because the kernel consumes 64-entry tiles. Other integrations use the 256 bucket. Neither shape was fully represented in the existing SM120 decode and prefill dispatch matrices, so decode-sized requests could miss the standalone kernel and fall through to the prefill-only orchestrator, which aborts on `num_tokens <= 64`. This change: - adds DSV4 decode instantiations for `topk=192` and `topk=256` with 8, 16, 32, 64, and 128 query heads; - adds DSV4 single-cache prefill dispatch for `topk=192` and `topk=256` using BF16 QK; - supports TP8 prefill by padding H8 into the existing 16-head MG tile while gating all global Q, sink, output, and LSE accesses; - enables the padded H8 MG path for both single-cache and dual-cache prefill, including both supported secondary-cache page sizes and runtime top-k lengths; - keeps the Python and CUDA dispatch matrices and user-facing documentation synchronized; - raises an actionable Python `ValueError` when an unsupported decode shape would otherwise enter the prefill-only kernel and abort the process; - fixes the sparse MLA benchmark KV cache view to match its default HND layout; and - expands correctness and benchmark coverage for the new dispatch shapes, active-length truncation, HND/NHD layouts, caller-owned workspaces, and output-buffer reuse. ## 🔍 Related Issues and Community PRs After coordination with the maintainers, this PR is intended as the single consolidated replacement for the overlapping open implementations: - #4309 — DSV4 `topk=192` decode/prefill and active-length coverage. - #4372 — broad DSV4 `topk=256` decode/prefill dispatch and public API coverage. - #3989 — combined DSpark `topk=256` fix and benchmark HND layout correction. - #3817 — DSV4 `topk=256` decode instantiations and decode-to-prefill fail-fast behavior. - #3834 — DSV4 `topk=256` BF16 prefill dispatch. It also covers the implementation scope of the earlier closed alternatives #3896 and #3923. Related reports include #3828, #3988, and #4336. The NVFP4 KV-cache extension in #3937 is a separate feature and is not superseded by this PR. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I used the repository-compatible `prek` hook runner. - [x] All hooks relevant to the changed files pass, including clang-format, mypy, ruff check, and ruff format. - [x] `git diff --check` passes. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] `pytest -q tests/attention/test_sparse_mla_sm120.py`: 309 passed on SM120. - [x] `compute-sanitizer --tool memcheck` on representative H8 single-cache, dual-cache PBS64, dual-cache PBS2, and runtime-length cases: 4 passed, `ERROR SUMMARY: 0 errors`. - [x] The complete `benchmarks/bench_sparse_mla_sm120.py` sweep runs successfully, including the corrected HND path and H8/H16 single- and dual-cache cases. ## Reviewer Notes The main review focus is the padded H8 MG specialization. The kernel still computes a 16-head tile, but only the first 8 rows may touch caller-owned global tensors; the remaining rows are zero-filled in shared memory and excluded from sink, output, and LSE accesses. Tests cover sinks, active-length truncation, the public HND API, dual-cache full-tile and runtime-length paths, and both secondary-cache page sizes. AI tools were used to assist with implementation, testing, and PR preparation. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Expanded sparse MLA support for additional head counts and top-k values, including 192, 256, and 2048. * Added support for eight-head configurations and more prefill, token-count, page-size, and dual-cache scenarios. * Improved compatibility with HND and NHD layouts. * **Bug Fixes** * Unsupported decode shapes are now rejected before processing. * Improved truncation handling for top-k and extra top-k lengths. * **Tests** * Expanded coverage for layouts, buffer reuse, workspace, sequence lengths, decode, and prefill scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: David Orman <ormandj@corenode.com> Co-authored-by: JimZhang-lab <60953754+JimZhang-lab@users.noreply.github.com> Co-authored-by: Albert Catalan Tatjer <ac.tatjer@gmail.com> Co-authored-by: carlosmolina0615 <carlosmolina0615@gmail.com> Co-authored-by: Wei-Cheng (Wayne) Chiu <waynehacking8@gmail.com> Co-authored-by: OsakaTX <OsakaTX@users.noreply.github.com> Co-authored-by: Martin Vit <martin@voipmonitor.org>
5 participants
Description
The SM120 sparse-MLA DSV4 prefill dispatch (
dispatch_dsv4_singleincsrc/sparse_mla_sm120_prefill.cu) only instantiatestopk in {128, 512, 1024, 2048}, sotopk=256(a valid4 x BI(64)shape, hit by DeepSeek-V4 + DSpark configs) fails withUnsupported sparse-MLA prefill configuration: model=DSV4 num_heads=32 topk=256 .... This adds the missingtopk == 256branch plus test coverage.This is the prefill half of #3828; PR #3817 covers the decode half. No overlap: this PR touches only the prefill
.cu, and the test configs are distinct from #3817's decode entries in the same test file (trivial rebase for whichever lands second).Compute-mode choice: BF16, not FP8. Measured on an RTX PRO 6000 (SM120, CUDA 13), BF16 beats FP8 at
topk=256for every shape tested, consistent with the existing small-K-loop comment (the FP8 Q-quantize prologue doesn't amortise at K=256):Verified on the same card: the issue's prefill repro (num_heads=32, topk=256, num_tokens=128) fails on current main and runs + matches the reference with this change.
Related Issues
Prefill portion of #3828 (decode portion: #3817).
Pull Request Checklist
Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.Tests
unittest, etc.): new configs(16, 256)and(32, 256)exercise the NHG=1 and NHG=2 routes; fulltests/attention/test_sparse_mla_sm120.py= 145 passed on SM120.Reviewer Notes
dispatch_dsv4_dualis intentionally untouched: the reporter's trace hastopk_extra=0, which routes throughdispatch_dsv4_single.Summary by CodeRabbit