[Fix] SM120 DSV4 decode: read FlashInfer's dispatch envelope (0.7.0) as well as the 0.6.18 pair table; opt-in padded heads

[avifenesh]

Motivation

_flashinfer_dsv4_decode_capabilities (#36655) iterates flashinfer.mla._sparse_mla_sm120._DECODE_DSV4_DISPATCH as (num_heads, topk) pairs. That is the flashinfer_python 0.6.18 shape. FlashInfer main since flashinfer-ai/flashinfer#4802 (453aa7c7296e, 2026-09-03, in every 0.7.0 build) publishes it as a _DecodeDispatchEnvelope(min_topk) predicate with __contains__ only (topk became a runtime kernel argument, so the dispatch is no longer an enumerable pair set). With that FlashInfer installed, a DeepSeek-V4 launch with attention TP on SM120 dies at EagerRunner warm-up:

  File ".../sglang/srt/models/deepseek_v4.py", line 957, in _kernel_num_heads
    if flashinfer_dsv4_decode_supports_num_heads(
  File ".../sglang/kernels/ops/attention/flash_mla_sm120.py", line 314, in _flashinfer_dsv4_decode_capabilities
    return int(_DECODE_MAX_TOKENS), frozenset(
TypeError: '_DecodeDispatchEnvelope' object is not iterable

Seen on 4x RTX PRO 6000, TP4, DeepSeek-V4.1-Flash, flashinfer_python 0.7.0 (7629d218). main hits the same the day its flashinfer pin moves past 0.6.18; the capability check was written to be fail-closed for a locally installed FlashInfer and this is the case it did not cover.

Modifications

• _flashinfer_dsv4_decode_dispatch() (lru_cache, one read) returns (decode_max_tokens, head widths, shape): pairs when the dispatch is iterable (0.6.18: the tuples' first elements), envelope when it is a container only (heads 1..._DECODE_MAX_HEADS probed with the envelope's own min_topk), none for anything else (no flashinfer, no attribute, an envelope without _DECODE_MAX_HEADS, a predicate that raises, a table that does not unpack), which keeps every caller on the padded 64-head path. _flashinfer_dsv4_decode_capabilities and flashinfer_dsv4_decode_supports_num_heads keep their signatures (test_flash_mla_backends.py calls the latter).
• SGLANG_SM120_DSV4_DECODE_PADDED (environ.py, EnvBool(False)): opt-in. While measuring the fix, [SM120] Use exact query-head widths for DeepSeek-V4 sparse MLA decode #36655's routing put TP4 decode (16 local heads) on FlashInfer's native 16-head kernel for the first time on this box, and single-request decode TPOT went from 8.115 to 9.134 ms (+12.6%) against the padded 64-head path; batch 4 / 8 (-2.0% / -0.9%) and every prefill and prefix-cache row were within the bar. The exact-width routing stays the default (it is the measured win for batched decode); =1 keeps the padded path for a deployment tuned for c=1 latency. A width the installed FlashInfer does not cover takes the padded path either way.
• _kernel_num_heads asks the new flashinfer_dsv4_decode_native_heads(n_local_heads, num_tokens), which logs once per head count which path it took and which dispatch shape it read, so the routing is visible in the server log:
  SM120 DSV4 decode heads: path=native local_heads=16 flashinfer_supports=yes flashinfer_dispatch=envelope(max_tokens=64, heads=1-128) SGLANG_SM120_DSV4_DECODE_PADDED=0

Not a duplicate of the open SM120 PRs that touch the same files: #38646 (NVFP4 KV cache format), #35104 (drops the 64-head pad on the prefill path; this PR is about the decode dispatch read and leaves prefill alone), #38969 (per-call Triton fallback for prefill shapes). None of them changes _flashinfer_dsv4_decode_capabilities.

Accuracy Tests

No numerics change: the fix only decides which already-existing decode path runs. On the 4x RTX PRO 6000 box the fixed build passed the same DeepSeek-V4.1 gate set as the pre-#36655 reference (greedy decode, tool calls, long-prefix cache hits, speculative acceptance 2.644 vs 2.636 tokens/verify).

Benchmarking and Profiling

4x RTX PRO 6000 (SM120), DeepSeek-V4.1-Flash, TP4 (16 local heads), flashinfer_python 0.7.0, 5% bar against the padded 64-head reference:

cell	padded 64-head (reference)	native 16-head (this PR's default)
decode TPOT c=1	8.115 ms	9.134 ms (+12.6%)
decode TPOT c=4	11.26 ms	11.033 ms (-2.0%)
decode TPOT c=8	14.152 ms	14.031 ms (-0.9%)
prefill 8k / 32k / 128k / 256k tok/s	7448 / 7965 / 7632 / 6883	7947 / 8025 / 7684 / 6927
warm prefix-cache TTFT	0.148 s	0.141 s

The c=1 row is the reason the padded path stays reachable through the knob.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests: test/registered/unit/kernels/ops/attention/test_flash_mla_sm120_decode_dispatch.py (11, CPU, no torch: the module is loaded from its path with torch / triton / envs stubbed and a stub flashinfer module in each of the two shapes; the fail-closed cases; the cached read; default native, knob padded, uncovered width padded, > decode_max_tokens never native).
• Update documentation as needed: the env var is documented in environ.py next to SGLANG_SM120_FLASHMLA_BACKEND.
• Update benchmark and profiling as needed: numbers above.
• Add the run-ci label if you want the CI to run.

---

CI States

Latest PR Test (Base): ❌ Run #34911171997
Latest PR Test (Extra): ❌ Run #34911171792
Latest PR Test (AMD ROCm 10): ❌ Run #34911171966