[BugFix] Sparse MLA: fix req_id_per_token OOB write and fp8_ds_mla context-prefill routing after the dense-MHA split

[drakosha]

FIX #48611

Purpose

Two regressions from #47327 (dense-MHA path for sparse-MLA short sequences). Full root-cause analysis with file:line is in the linked issue; the short version:

1. OOB write in the FlashMLA sparse index conversion. After #47327, forward_mqa receives only the MQA subset (the leading decode tokens), but all three FlashMLA forward paths still passed the full-batch req_id_per_token to triton_convert_req_index_to_global_index, whose grid is sized by req_id while the output is allocated like token_indices — the extra rows were written past the end of the output buffer. Fixed by slicing req_id_per_token[: topk_indices.shape[0]] at the three call sites, exactly as the FlashInfer and FlashAttn sparse backends already do, plus a length-invariant assert in the shared converter so a future mismatch fails loudly instead of corrupting memory.

2. fp8_ds_mla context gather on the dense-MHA path. When a short prefill carries cached context, forward_mha gathers it with kernels that do not understand the 656B fp8_ds_mla entry layout: raw cp_gather_cache rejects the uint8/bf16 dtype mismatch under DCP, and gather_and_maybe_dequant_cache otherwise dequantizes the raw bytes as flat fp8 with a single global scale, silently corrupting the context K/V. Fixed by keeping such prefills on the top-k MQA path, gated on kv_cache_dtype == "fp8_ds_mla" and prefill.chunked_context is not None.

Short no-context prefills — the case #47327 optimizes and benchmarks (DeepSeek-V3.2 short prompts) — never read the cache, so they keep the new dense-MHA fast path. kv-cache-dtype aliases cannot bypass the gate: fp8/fp8_e4m3 are canonicalized to fp8_ds_mla before the layer stores kv_cache_dtype.

Not a duplicate: no open PR addresses either regression (searched req_id_per_token, sparse-MLA / fp8_ds_mla dense-MHA). #34744 is the masked-MHA follow-up to #47327; #45537 is an unrelated cp_gather_cache bounds fix.

Test plan

New regression tests in tests/v1/attention/test_sparse_mla_backends.py:

• test_triton_convert_rejects_req_id_longer_than_token_indices — the converter rejects the full-batch/subset length mismatch instead of writing OOB, and the sliced call matches the reference.
• test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens — call-site regression: _forward_bf16_kv with full-batch metadata and a subset q/topk_indices converts exactly the MQA rows (fails before the fix).
• test_fp8_ds_mla_context_prefill_stays_on_mqa_path — the dense-MHA path is blocked precisely for fp8_ds_mla + chunked context and stays available with no context; covers dtype-alias canonicalization.

pytest tests/v1/attention/test_sparse_mla_backends.py -v \
  -k "triton_convert or forward_bf16_kv_slices or fp8_ds_mla_context"

Test result

• ruff check / ruff format on all touched files: clean.
• GPU suite (4×H200, CUDA 12.9): pytest tests/v1/attention/test_sparse_mla_backends.py → 111 passed, 422 skipped, 0 failed (8 min). The three new regression tests pass by name:
  • test_triton_convert_rejects_req_id_longer_than_token_indices PASSED
  • test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens PASSED
  • test_fp8_ds_mla_context_prefill_stays_on_mqa_path PASSED
• End-to-end: both fixes are running in our production deployment (GLM-5.2-NVFP4, 4×H200, TP4 / DCP4 / EP, fp8_ds_mla, MTP, prefix caching + CPU KV offload). Multi-turn chat with prefix-cache hits — the exact case that hits the broken fp8_ds_mla context gather — returns correct output; a 4×551k-token concurrent stress run retrieved all 4 needles with no crash.
• No model-eval delta expected on supported paths: the change only restores pre-[1/N] Add dense MHA path for sparse MLA short sequences #47327 routing for the broken cases and leaves the new fast path in place for short no-context prefills.

AI assistance

AI assistance (Claude) was used for this change. The human submitter reviewed every changed line and ran the tests above.

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

[MatthewBonanni]

Thanks for catching these! I'd like to supersede this PR with #48642, though, which enables dense MHA in these circumstances instead of routing away from it.

[drakosha]

Closing in favour of #48642 — enabling dense MHA is the better fix than routing those prefills back to MQA, and it also covers _forward_fp8_kv_separate_prefill_decode, which this PR left alone (the "additional fallout" note in #48611). Thanks for picking up the tests.

Two notes:

• [Bug]: #47327 dense-MHA split breaks FlashMLA sparse: OOB write in top-k index conversion, corrupted fp8_ds_mla context gather #48611 stays open and also tracks that third issue — worth linking it from [Bugfix] Sparse MLA: enable fp8_ds_mla dense prefill #48642 so it closes on merge.
• We run GLM-5.2 on 4×H200 with TP4/DCP4, fp8_ds_mla, MTP and prefix caching (16 heads/rank → the mixed-batch path), which exercises exactly the DCP chunked-context gather your new seq_starts kernel now handles. Happy to run your branch e2e on that config and report back — it's a corner a DSV3.2 eval won't cover.