[Bugfix][SM120][MLA] Read fp8_ds_mla tile scales as arbitrary fp32

[oops-oom]

Purpose

The read format has to match the write format. The vLLM writers used with this backend store each tile scale as raw fp32 amax / 448 (concat_and_cache_ds_mla and the DeepSeek-V3.2 fused Triton path). However, FLASHINFER_MLA_SPARSE_SM120 asked FlashInfer for pow2_fp32, which extracts the exponent and reconstructs 2^floor(log2(scale)). For a non-power-of-two scale written by vLLM, that is smaller by a factor in (0.5, 1].

The format was chosen from a model_type allowlist that only sent glm* to the correct reader, so every other sparse model landing here reads its scales wrong: DeepSeek-V3.2 today, and HY V4 as soon as it can select this backend (it needs
the sink support in the stacked follow-up, since learnable_sink defaults to on). #54434

FlashInfer supports both scale conventions for the same 656-byte physical layout, but vLLM's writer contract currently always stores arbitrary fp32 scales. Kimi-K3 demonstrates that this convention is not tied to the model label: its dense MLA path writes the same layout even though it never reaches this sparse backend. This backend accepts no other layout (it raises unless
kv_cache_dtype == "fp8_ds_mla"), so its reader format becomes a constant.

DeepSeek-V4 is not a counterexample. It reuses the fp8_ds_mla string for a different footer-scale layout: 576 bytes of per-token data plus 8 bytes of UE8M0 2^(u8-127) scales in the page footer, or 584 logical bytes per token.
That layout is written by DSv4's own CuteDSL/Triton insert rather than concat_and_cache_ds_mla, and DSv4 is rejected here: FLASHINFER_MLA_SPARSE_SM120 is not a DSv4 backend; those models must use FLASHINFER_MLA_SPARSE_DSV4.

Test Plan

Run on real SM120 hardware (AWS g7.2xlarge: RTX PRO 4500 Blackwell, sm_120, driver 595.91.07, CUDA 13.2, torch 2.13.0+cu132, flashinfer 0.6.17).

.venv/bin/python -m pytest tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py -v
.venv/bin/python -m pytest tests/v1/attention/test_sparse_mla_backends.py -q

Because this patch changes how bytes are dequantized, the decisive test is a kernel-level A/B: pack a fp8_ds_mla cache exactly the way vLLM's writer does (amax / 448 stored as raw fp32), call the same public FlashInfer decode API
this backend calls, once per kv_scale_format, and compare both against an fp32 dequant-then-attention reference computed from the same bytes.

Test Result

Unit tests: 18 passed.

Neighbouring suite test_sparse_mla_backends.py: 46 passed, 504 skipped, 4 failed — the same 4 masked_mha prefill cases fail on the unpatched base commit (d3d79ffc1e) on this GPU, so they are pre-existing and unrelated.

Scale A/B, 1024-token cache, 32 heads, top-k 2048. 4083 of 4096 stored tile scales are not powers of two (mean pow2/true factor 0.776, min 0.505), so the two readers genuinely disagree:

reader	rel L2 vs writer reference	cos
arbitrary_fp32 (this PR)	0.0030	1.0000
pow2_fp32 (before)	0.0213	0.9998

The gap widens with per-tile scale spread, which is what real K vectors look like (outlier channels in the nope tiles):

log2 scale spread	arbitrary_fp32	pow2_fp32
1.4	0.0030	0.0213 (7.0x)
4.7	0.0032	0.0224 (7.0x)
10.1	0.0062	0.0737 (11.9x)
18.0	0.0196	0.1985 (10.2x)

FlashInfer's own docstring confirms the two semantics: "auto"/"pow2_fp32" "select DSv3.2 power-of-2 FP32 inline scales", "arbitrary_fp32" "selects GLM-style arbitrary FP32 inline scales".

No end-to-end eval: every model that can reach this backend (DeepSeek-V3.2, GLM, HY V4) is far larger than the biggest single SM120 (96 GB), so a serving eval is not reachable on this hardware. The A/B above measures exactly the quantity this patch changes, on the kernel this patch configures.

A/B harness (abridged)

# pack the way vLLM's concat_and_cache_ds_mla does: raw fp32 amax/448
for ti in range(4):
    tile = kv[:, ti * 128 : (ti + 1) * 128].float()
    scale = (tile.abs().amax(-1).clamp(min=1e-4) / 448.0).to(torch.float32)
    packed[:, ti * 128 : (ti + 1) * 128] = (tile / scale[:, None]).to(
        torch.float8_e4m3fn
    ).view(torch.uint8)
    packed[:, 512 + ti * 4 : 512 + (ti + 1) * 4] = scale.view(torch.uint8).view(-1, 4)

for fmt in ("pow2_fp32", "arbitrary_fp32"):
    out = trtllm_batch_decode_with_kv_cache_mla(
        query=q.unsqueeze(1), kv_cache=kv_hnd, workspace_buffer=ws,
        qk_nope_head_dim=512, kv_lora_rank=512, qk_rope_head_dim=64,
        block_tables=indices.unsqueeze(1), seq_lens=None,
        max_seq_len=indices.shape[-1], sparse_mla_top_k=indices.shape[-1],
        bmm1_scale=sm_scale, bmm2_scale=1.0, backend="sparse",
        kv_scale_format=fmt,
    ).squeeze(1)
    report(fmt, out, reference_from_the_same_bytes)

Why this is not duplicating an existing PR

Checked per AGENTS.md:

gh pr list --repo vllm-project/vllm --state open --search "pow2_fp32"
gh pr list --repo vllm-project/vllm --state open --search "arbitrary_fp32"
gh pr list --repo vllm-project/vllm --state open --search "kv_scale_format"
gh pr list --repo vllm-project/vllm --state open --search "SM120 sparse MLA"

No open PR changes the model-dependent kv_scale_format selection fixed here (both explicit format searches return nothing). #47527 wires the packed-cache arguments, including kv_scale_format, into the released FlashInfer API but
preserves the existing selection; this PR fixes the value passed through that API. #53969 adds NoPE support and validates the effective top-k buffer width. Both touch the same files, so textual conflicts are possible, but neither duplicates this change.

AI assistance disclosure (per AGENTS.md)

AI assistance was used to cross-check the vLLM writer kernels against FlashInfer's reader paths, to run the SM120 hardware A/B above, and to write the test. The submitter reviewed every changed line.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[oops-oom]

Hi @lucifer1004, could you help sanity-check the intended scale semantics here?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @oops-oom.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork