[SM12x] Drop the 64-head TP pad on the DSv4 MLA prefill path

[zhendonghua]

Problem

With attn_tp_size=2, MQALayer builds a 64-row query tensor, writes only [0:n_local_heads] (= 32), runs attention on all 64 rows, and discards half the output via o = o[:, tp_slice, :].

# FlashMLA's fp8 sparse decode kernel only specializes h_q for {64, 128}.
padded_num_heads = 64 if self.n_local_heads <= 64 else self.n_heads

That comment is true of sgl_kernel.flash_mla — the else branch at deepseek_v4_backend.py:1774, which SM12x never takes. On SM12x the kernel is FlashInfer's, and the decode/prefill split happens by row count inside the FlashInfer entry (B <= _DECODE_MAX_TOKENS, flash_mla_sm120.py:558), not by forward mode. So a wide prefill chunk is dispatched to the paged/prefill kernel, which takes arbitrary h_q, and the pad buys it nothing.

MLA is the one place where this is pure waste rather than a wash: 64 query heads share a single latent KV (num_key_value_heads = 1), so head-sharding saves compute without saving KV bandwidth. Prefill (many query rows, compute-bound) scales with head count; decode (few rows, KV-bandwidth-bound) does not.

Fix

Gate the pad off for wide prefill chunks only, behind SGLANG_OPT_DSV4_UNPAD_PREFILL_MLA_HEADS (default on, SM12x only).

• row term — this is the one that selects the kernel. Strictly above _DECODE_MAX_TOKENS, read from FlashInfer (flashinfer/mla/_sparse_mla_sm120.py; 64 in every flashinfer release to date) rather than hardcoded, so the gate and the dispatch fork at flash_mla_sm120.py:558 are provably the same number. If FlashInfer ever raised that cutoff, a hardcoded 64 here would unpad rows that still route to the decode kernel — a silent kernel-family swap, not a crash. That is what the import buys.
• mode terms — scope, not legality. is_extend_or_draft_extend_or_mixed() (an allowlist, not is_extend(), which is also True for TARGET_VERIFY), plus _original_forward_mode because DP padding can rewrite a verify batch into EXTEND. To be explicit about what these do and do not do: FlashInfer's SM120 decode kernel is instantiated for num_heads ∈ {8,16,32,64,128} as well (_DECODE_DSV4_DISPATCH), so unpadding is legal on both sides of the row fork. The mode terms exist to keep TARGET_VERIFY and large-bs DECODE — both CUDA-graph-captured — on exactly the kernel instantiation they run today. I would rather change one path at a time here than two.
• Because the gate reads forward_mode, which is not part of a captured graph's shape key, it bails under is_in_breakable_cuda_graph(). That is a consequence of the mode terms, not an independent requirement; a row-count-only gate would not need it (and would also work under --cuda-graph-backend-prefill=breakable, where this PR is a no-op).

Supporting changes:

• sm120_flashmla_decode_max_tokens() exposes the threshold instead of duplicating the constant. On an import failure it returns a value no batch can exceed, so the gate fails closed (everything keeps the pad). It returns 0 when the backend is not FlashInfer — there the triton and torch fallbacks read h_q off the query shape, so there is no floor to respect and the gate is free to fire.
• _local_attn_sink() takes an optional num_heads. The narrowed sink is a prefix view of the cached padded tensor, not a second allocation, so a device pointer baked into a captured decode graph stays valid; and when num_heads is None or already equals the stored width, the same object is returned. DSparkAttention (deepseek_v4_dspark.py:246) calls it with no argument and is unchanged — the DSPARK draft block stays padded.

Measurements

2× DGX Spark (GB10, sm_121), DeepSeek-V4-Flash-0731, TP=2 across two nodes.

• MLA prefill lane: −49.7%
• Kernel instantiation goes from sparse_mla_prefill_mg_dual_kernel<(ModelType)1, (ComputeMode)1, 64, ...> to <..., 32, ...>, verified from an in-server profile rather than inferred from the launch flags. That profile is also the evidence that the gate actually fired.
• That matches the geometry the reference vLLM deployment of this model runs on the same box (<..., 32, 128, ...>, grid [4072,1,1] vs ours [4096,1,1], block [384,1,1]).

Decode, TARGET_VERIFY, and the DSPARK draft block are untouched and stay padded.

Correctness — what is and is not established

• GSM8K 500 questions: 0.944 → 0.948, Invalid: 0.000. This is the wrong instrument and I am not claiming it as evidence of equivalence. Attention is per-(token, head) independent and the dropped heads were fed new_empty garbage and then sliced away, so the surviving heads should be bit-identical and the score should have been unchanged. A 2-in-500 delta means either the two runs differed in more than this flag, or the NUM_HEADS=32 and NUM_HEADS=64 template instantiations are not bit-identical on the heads that survive. A gate-ON/gate-OFF greedy text-exact comparison is owed and I will post it.
• Shape safety on the path this PR changes is enforced by FlashInfer, loudly: TVM_FFI_ICHECK_EQ(attn_sink.size(0), num_heads) and a compiled switch (num_heads) whose default: returns false into TVM_FFI_ICHECK. Nothing on the prefill path fails silently.

Known gaps

• No head-count legality guard, and the head envelope is flashinfer-version-dependent. The FlashInfer prefill dispatch is a closed switch (num_heads). The case 8: instantiation landed with feat(sm120): consolidate DSV4 sparse MLA top-k 192/256 support flashinfer-ai/flashinfer#4380 (merged 2026-08-08) — after the v0.6.17 release branch was cut, so the version this repo pins (flashinfer_python==0.6.17) ships without it; the first releases carrying it are the v0.6.18 nightlies. On the pinned wheel the prefill head set is {16, 32, 64, 128}: at attn_tp_size = 8 (n_local_heads = 8) the unpadded path would abort where the padded path worked. attn_tp_size ≥ 16 is already non-functional for this model for an unrelated reason (o_groups = 8 ⇒ n_local_groups = 0). A n_local_heads in (16, 32, 64, 128) term in the gate would close this (16 as the floor while 0.6.17 is pinned); I would like a view on whether the allowlist or a capability probe is preferred. (Same-breath note, independent of this PR: feat(sm120): consolidate DSV4 sparse MLA top-k 192/256 support flashinfer-ai/flashinfer#4380 is also what adds the topk ∈ {192, 256} decode instantiations that SM12x DSPARK — topk fixed at 192 — needs, so stock v0.6.17 fails that path regardless of head count.)
• No test. The gate is a pure predicate over (rows, forward_mode, _original_forward_mode, env, graph flag) and is testable on CPU with _is_sm120 and the threshold monkeypatched; happy to add it.
• The pre-existing comment at deepseek_v4.py:1362 attributes the {64,128} restriction to "FlashMLA's fp8 sparse decode kernel" without saying which kernel. It is correct for sgl_kernel.flash_mla and stale for the FlashInfer kernel SM12x calls. Say the word and I will amend it in this PR.

Notes

Only sm120/sm121 hardware was available; the gate is is_sm120_supported()-scoped and defaults off everywhere else. Note that is_sm120_supported() is cc_major == 12, so it also covers (12,0) parts I could not test — the row-threshold argument transfers because they take the same FlashInfer entry.

CI: the red on base-b-test-1-gpu-large (4) is test/registered/quant/test_awq.py::TestAWQMarlinBfloat16::test_mmlu (0.82421875 not greater than 0.83), which is unrelated to this change and is currently flaky at its threshold on unrelated branches too — #34477 raised num_examples 64 → 256 without re-calibrating the bound. The CPU build-test and AMD failures reproduce on main.

🤖 Generated with Claude Code

[zhendonghua]

https://github.com/flashinfer-ai/flashinfer/blob/9cd1fcffcb0f801c8337585fa0182a35fc1a4c91/flashinfer/mla/_sparse_mla_sm120.py#L74

[zhendonghua]

/tag-and-rerun-ci