fix(sparse-mla-sm120): add topk=256 prefill-dsv4 dispatch (fixes #3828)

[OsakaTX]

📌 Description

Implements the prefill half of #3828, exactly as proposed there by @eous: the
SM120 sparse-MLA DSV4 prefill path (dispatch_dsv4_single in
csrc/sparse_mla_sm120_prefill.cu) has no topk == 256 branch — the dispatch
chain covers {128, 512, 1024, 2048} — so a prefill-shaped call
(num_tokens > 64) with topk=256 returns false and fails:

Check failed: (ok) is false: Unsupported sparse-MLA prefill configuration:
model=DSV4 num_heads=32 topk=256 page_block_size=64 topk_extra=0 extra_page_block_size=0

As #3828 notes, topk=256 is a valid shape (4 × BI(64), between the
instantiated 128 and 512); it's simply missing from the grid. The shape occurs
in practice via DeepSeek-V4-Flash-DSpark under tensor-parallel=2 — the decode
half of the same gap is addressed by #3817.

Change

• csrc/sparse_mla_sm120_prefill.cu: else if (topk == 256) DISPATCH_BY_NH_CM(FP8, 256);
  in dispatch_dsv4_single — FP8 compute mode per the existing comment (the
  BF16 QK fast path is reserved for the small K-loop at 128; FP8-512 already
  builds, so the smaller 256 tile is smem-safe).
• tests/attention/test_sparse_mla_sm120.py: add (16, 256) and (32, 256)
  to _DSV4_PREFILL_CONFIGS (16 exercises the MG_N_HG_T=1 routing, 32 is the
  motivating TP=2 shape).

🧪 Testing

Isolation A/B on GB10 (sm_121a), CUDA 13 — identical container and environment,
module JIT-built from source in both legs, only the two dispatch lines varied.
This complements the RTX PRO 6000 (SM120) serving validation reported in #3828 —
two architectures, same result:

• A (unmodified main + the new test configs): exactly the 8 topk=256
  prefill combos fail (num_tokens ∈ {128, 256} × with_sink ∈ {False, True} ×
  NH ∈ {16, 32}) with the error above; all 36 pre-existing prefill combos pass.
• B (this branch): all prefill combos pass, new shapes verified against the
  bf16 reference implementation.

🔗 Related

• Fixes [SM120] Sparse-MLA DSV4 decode+prefill kernels not instantiated for topk=256 #3828 (prefill half; analysis and fix design are @eous's — this PR
  implements it with test coverage and additional hardware validation)
• fix(sparse-mla-sm120): add TOPK=256 decode-dsv4 instantiation for DeepSeek-V4 DSpark #3817 — the decode half of the same missing-topk=256 gap (independently
  confirmed on my hardware in that thread)

Summary by CodeRabbit

• New Features
  • Expanded support for additional prefill configurations, including topk=256, for certain sparse attention workloads.
• Tests
  • Added test coverage for two more head-count and topk=256 combinations to improve validation of the new supported settings.

[gemini-code-assist]

Code Review

This pull request adds support for a topk value of 256 using FP8 precision in the dispatch_dsv4_single function within csrc/sparse_mla_sm120_prefill.cu. Additionally, it updates the test configurations in tests/attention/test_sparse_mla_sm120.py to include test cases for this new configuration. 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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 1a2cbe24-972f-476d-93dc-6f3f64512797

📥 Commits

Reviewing files that changed from the base of the PR and between 5eea6b4 and 1125246.

📒 Files selected for processing (2)

• csrc/sparse_mla_sm120_prefill.cu
• tests/attention/test_sparse_mla_sm120.py

---

📝 Walkthrough

Walkthrough

Added an else if (topk == 256) branch in dispatch_dsv4_single to dispatch the FP8 MG kernel variant for DSV4 single-cache prefill, and extended _DSV4_PREFILL_CONFIGS in the test suite with two new (num_heads, topk) configurations covering topk=256.

Changes

DSV4 Prefill topk=256 Support

Layer / File(s)	Summary
Add topk=256 dispatch branch and test coverage csrc/sparse_mla_sm120_prefill.cu, tests/attention/test_sparse_mla_sm120.py	Adds a new conditional branch dispatching the FP8 MG kernel path for topk == 256 in dispatch_dsv4_single, and adds (16, 256) and (32, 256) test configurations to validate the new path.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Suggested labels: run-ci

Suggested reviewers: saltyminty, yzh119, nv-yunzheq, bkryu, qsang-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Linked Issues check	⚠️ Warning	The PR only fixes the prefill branch, while #3828 also requires the decode topk=256 instantiations, so the full issue is not satisfied.	Add the decode-side topk=256 dispatches and coverage, or link this PR only to a prefill-only subtask.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly states the SM120 sparse-MLA prefill topk=256 dispatch fix and matches the main change.
Description check	✅ Passed	It covers the bug, the prefill fix, testing, and related issue, though the template's checklist section is not filled out.
Out of Scope Changes check	✅ Passed	The diff stays focused on the sparse-MLA SM120 prefill fix and matching tests, with no unrelated code changes.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[waynehacking8]

Heads up: #3834 (opened 07-04) implements this same prefill-half of #3828 -- same dispatch branch and the same (16,256)/(32,256) test configs -- but with BF16 CM instead of FP8, picked from measurement on an RTX PRO 6000 (BF16 beat FP8 by 2.5-9% at all six shapes I benched; the small-K-loop prologue argument in the PR body). ormandj's multi-day SM120 serving soak upthread on #3817 ran with the BF16 version of this branch. Not territorial about which PR lands, but the compute-mode choice is a real fork in the road and probably deserves a deliberate call by the maintainers.

[OsakaTX]

Heads up: #3834 (opened 07-04) implements this same prefill-half of #3828 -- same dispatch branch and the same (16,256)/(32,256) test configs -- but with BF16 CM instead of FP8, picked from measurement on an RTX PRO 6000 (BF16 beat FP8 by 2.5-9% at all six shapes I benched; the small-K-loop prologue argument in the PR body). ormandj's multi-day SM120 serving soak upthread on #3817 ran with the BF16 version of this branch. Not territorial about which PR lands, but the compute-mode choice is a real fork in the road and probably deserves a deliberate call by the maintainers.

Good catch, and I appreciate the friendliness. Haven't submitted many PRs before, mostly do private dev work. Agree the CM choice is the real decision and it belongs to the maintainers. Your RTX PRO 6000 data is compelling for sm_120. One datapoint is missing: sm_121a (GB10/DGX Spark), which is where the (32,256) shape actually bites in practice (DSv4 DSpark under vLLM TP=2). I'll try and bench both CM variants at the six shapes on GB10 hardware and post numbers here within a day or two. If BF16 holds on sm_121a as well, #3834 is the better patch and I'm happy to close this one in its favor. The fix landing matters more than whose branch carries it.

[waynehacking8]

sm_121a numbers would genuinely complete the picture -- it's a different chip and a different cubin variant (121a vs 120f), so I wouldn't assume my sm_120 ordering transfers. And the goodwill cuts both ways: if FP8 wins on GB10 I'm happy to flip #3834's CM or split it per-arch, whichever the maintainers prefer. Looking forward to your numbers.

[OsakaTX]

Closing this in favor of #3834. I benchmarked both compute modes for the topk == 256 DSv4 prefill dispatch on GB10 (sm_121a) and BF16 is the better default.

Method. #3834 and this PR differ by exactly one token in the dispatch (DISPATCH_BY_NH_CM(BF16, 256) vs FP8). To isolate that single variable with no base-commit confounds, I built both on the same checkout, forcing a clean JIT recompile from source each time (uninstalled flashinfer-jit-cache and flashinfer-cubin). KV input is identically packed via quantize_kv_dsv4 in both, so this measures the compute-mode path only, not KV memory. Timing is CUDA-event based, 60 warmup + 5x300 iters, median of reps (per-config std < 1 us). Both builds pass the (16, 256) and (32, 256) prefill correctness tests on sm_121a.

Result (BF16 time saved vs FP8, positive = BF16 faster):

heads	prefill tokens	full (256 valid idx)	half (128 valid, test setup)
16	128	+12.8%	+12.6%
16	256	+10.1%	+9.3%
16	512	+13.1%	+1.8%
16	1024	+5.3%	~0%
16	2048	+4.3%	+1.3%
32	128	+19.1%	+19.2%
32	256	+16.4%	+15.9%
32	512	+1.8%	+1.1%
32	1024	+0.2%	+1.3%
32	2048	+1.9%	-0.2%

BF16 is faster than or equal to FP8 in every configuration. The margin is largest at short prefills (about 10 to 19% at 128 to 256 tokens), where the kernel is latency-bound and the FP8 dequant path costs the most, and it narrows toward parity at long prefills (about 0 to 5% at 1024 to 2048 tokens) as the kernel becomes bandwidth-bound. FP8 shows no throughput advantage, so #3834 is the right choice for this dispatch on GB10. Thanks @waynehacking8.