Fix SM120 DSV4 topk-256 decode dispatch

[voipmonitor]

Summary

Add the missing SM120 DSV4 sparse-MLA decode instantiations for topk=256 and fail explicitly when a decode-sized DSV4 shape is unsupported instead of silently falling through to the generic paged-attention path.

Why

DeepSeek-V4-Flash-DSpark can build non-causal sparse-attention metadata with a 128-token context plus a padded draft block. That produces a topk=256 decode bucket. The existing DSV4 dispatch table supports 128, 512, and 1024, so valid DSpark decode requests miss the specialized SM120 kernel.

The added instantiations cover the DSV4 head counts used by TP1/TP2/TP4 (8, 16, 32, and 64). The Python dispatch table is kept in sync with the C++ switch.

For decode-sized inputs, an unsupported shape now raises a diagnostic containing model type, token count, head count, top-k, QK width, page-block size, and extra top-k. Prefill-sized inputs retain the existing generic path.

Validation

• git diff --check
• python3 -m py_compile flashinfer/mla/_sparse_mla_sm120.py tests/attention/test_sparse_mla_sm120.py
• Added (num_heads=32, topk=256) to the existing SM120 DSV4 decode matrix.
• Confirmed the dispatch is present in the image used for TP2 DSpark end-to-end validation.

Summary by CodeRabbit

• New Features

  • Added SM120 sparse MLA decode support for TOPK=256 across 8, 16, 32, and 64 attention heads.

• Bug Fixes

  • Unsupported decode configurations now return a clear error instead of silently using a fallback path.

• Tests

  • Added coverage for the 32-head, TOPK=256 configuration.

[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: b0f89b29-2906-4bad-a2a7-a0070e1cc988

📥 Commits

Reviewing files that changed from the base of the PR and between 422f6c9 and e939be5.

📒 Files selected for processing (3)

• csrc/sparse_mla_sm120_decode_dsv4.cu
• flashinfer/mla/_sparse_mla_sm120.py
• tests/attention/test_sparse_mla_sm120.py

---

📝 Walkthrough

Walkthrough

SM120 DSV4 decode dispatch now supports TOPK=256 for 8, 16, 32, and 64 heads in CUDA and Python routing. Unsupported decode shapes raise an explicit error, and tests add coverage for the (32, 256) configuration.

Changes

SM120 DSV4 TOPK 256 support

Layer / File(s)	Summary
CUDA dispatch expansion csrc/sparse_mla_sm120_decode_dsv4.cu	Adds DSV4 dispatch branches for (8,256), (16,256), (32,256), and (64,256), and updates the adjacent TOPK documentation.
Python routing and validation flashinfer/mla/_sparse_mla_sm120.py, tests/attention/test_sparse_mla_sm120.py	Adds matching Python dispatch entries, raises RuntimeError for unsupported decode shapes, and tests (32,256).

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related issues

• flashinfer-ai/flashinfer issue 3828 — Covers the missing DSV4 TOPK=256 CUDA dispatch and Python routing addressed here.

Possibly related PRs

• flashinfer-ai/flashinfer#3545 — Expands related DSv4 shape validation, connecting to the newly supported decode configurations.

Suggested reviewers: saltyminty, sricketts, aleozlx, yzh119, cyx-6, kahyunnam, nv-yunzheq

Sequence Diagram(s)

sequenceDiagram
  participant _paged_attention
  participant _DECODE_DSV4_DISPATCH
  participant launch_sparse_mla_decode_dsv4
  participant launch_decode_dsv4_impl
  _paged_attention->>_DECODE_DSV4_DISPATCH: match num_heads and topk
  _DECODE_DSV4_DISPATCH->>launch_sparse_mla_decode_dsv4: route supported TOPK=256 shape
  launch_sparse_mla_decode_dsv4->>launch_decode_dsv4_impl: dispatch selected head count and TOPK=256

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding SM120 DSV4 topk-256 decode dispatch support.
Description check	✅ Passed	The description covers what changed, why, and how it was validated, though it omits the template's related issues and checklist sections.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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.}

[gemini-code-assist]

Code Review

This pull request introduces support for a TOPK value of 256 in the SM120 sparse-MLA decode kernel for DeepSeek-V4 (DSV4), which accommodates DSpark's 128-token context and padded draft block. The changes update the CUDA dispatch macros, the Python dispatch set, and add a corresponding test case. Additionally, a runtime check is introduced to raise an error when num_tokens is less than or equal to _DECODE_MAX_TOKENS. There are no review comments, so no feedback is provided.

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.

[voipmonitor]

Closing in favor of the earlier and broader #3817, which covers the same DSV4 topk=256 decode dispatch, includes NUM_HEADS=128, tests, and the unsupported-shape guard. The release image will use #3817 together with prefill fix #3896.