[Bugfix] Sparse MLA: enable fp8_ds_mla dense prefill

[MatthewBonanni]

Purpose

PR #47327 added dense-MHA prefill routing for sparse MLA, exposing two fp8_ds_mla mixed-batch bugs:

• req_id_per_token could exceed the top-k tensor length and cause out-of-bounds writes.
• Dense MHA could not gather the packed 656-byte FP8 cache, and FlashMLA assumed its metadata covered the full batch.

This PR bounds the converter input, adds arbitrary sequence starts to the packed FP8 gather/upconvert kernel, and supports decode-only MQA metadata while dense MHA handles prefills. It supersedes #48612.

AI assistance was used. The human submitter reviewed the diff.

Tests

.venv/bin/python -m pytest tests/kernels/test_cp_gather_fp8.py -q
.venv/bin/python -m pytest tests/v1/attention/test_sparse_mla_backends.py -q -k 'flashmla_fp8_paths or cache_dtype_aliases or triton_convert or forward_bf16_kv_slices'
.venv/bin/python -m pytest 'tests/v1/attention/test_sparse_mla_backends.py::test_sparse_backend_decode_correctness[1.0-1.0-64-1-fp8_ds_mla-mixed_small-FlashMLA]' -q
.venv/bin/pre-commit run --files <changed files>

• FP8 gather: 16 passed
• Sparse-MLA regressions: 16 passed, 515 deselected
• Mixed decode/dense-prefill correctness: 1 passed
• Pre-commit: passed

GSM8K

RedHatAI/DeepSeek-V3.2-NVFP4-FP8-BLOCK, 4×B200, TP2/DP2/EP, FLASHMLA_SPARSE, fp8_ds_mla, official 5-shot lm_eval, all 1,319 samples:

Routing	Flexible	Strict
Dense MHA	89.69%	89.54%
Forced sparse MQA	89.92%	89.69%

The difference is below the ~0.84-point standard error. No measurable accuracy regression from dense-MHA routing was observed. Throughput is not compared because scheduler batch sizes differed.

[drakosha]

Ran your branch with the #46514 DCP change on top, on 4×H200 GLM-5.2-NVFP4 (TP4/DCP4/EP, fp8_ds_mla, MTP=3, prefix caching + offload). Your dense-MHA seq_starts gather is correct under DCP: a prefix-cache-hit test (short questions reading a cached ~1.4k briefing back through the new gather) recalls every planted fact, with prefix-cache hits confirmed — the DCP path a DSV3.2 eval won't reach. Decode throughput is at parity with our gated build. The DCP merge was clean (only conflict: the req_id slice vs our DCP branch in _forward_fp8_kv_mixed_batch).

Two heads-ups:

• Blocks a clean fp8_ds_mla eval of this PR: on current main, MLAAttention.get_kv_cache_spec omits kv_quant_mode, so the Support nvfp4 kv with kv-cache-dtype-skip-layers sliding_window #42890 reshape sizes the KV view by head_size (576) vs the 656-byte page → crash at engine init. Not caused by this PR; [Bugfix] Set kv_quant_mode on the generic MLA KV-cache spec #48379 fixes it.
• When [Attention][MLA] FlashMLA sparse: DCP on the fp8_ds_mla mixed-batch path + MTP #46514 rebases here, your test_flashmla_fp8_paths_accept_decode_subset[True] mock needs dcp_world_size=1 + need_to_return_lse_for_decode=False (the DCP branch reads both off self). Separately, DCP+batch needs [Bugfix] Size the sparse-indexer expanded block table from the runner's block-table width #48404 (indexer expanded-block-table width) — DCP-only, harmless at DCP=1.

(AI assistance used; every claim checked against source.)

[gaby]

@MatthewBonanni PR #48379 was merge, this can be closed ?

[MatthewBonanni]

@gaby, this PR is still necessary to resolve the listed issues

[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.

[mgoin]

Looks reasonable to me, will rely on CI and manually unblocked tests