[Bug]: #47327 dense-MHA split breaks FlashMLA sparse: OOB write in top-k index conversion, corrupted fp8_ds_mla context gather

[drakosha]

Summary

#47327 (c4f5cd60d, "[1/N] Add dense MHA path for sparse MLA short sequences") routes short prefills (prefill_max_seq_len <= topk_tokens) to a dense-MHA path, so forward_mqa now receives only the leading num_decode_tokens tokens of the batch instead of all of them. Two parts of the FlashMLA sparse backend were not updated for that new contract:

1. an out-of-bounds write in the top-k index conversion whenever the split is active (any kv-cache dtype);
2. the dense-MHA path gathers cached context in a way that does not understand the fp8_ds_mla cache layout — silently wrong K/V without DCP, hard error with DCP.

Both are regressions from #47327: before it, every token went through the MQA path. cc @MatthewBonanni

Repro

vllm serve deepseek-ai/DeepSeek-V3.2 \
  --tensor-parallel-size 8 \
  --kv-cache-dtype fp8_ds_mla

Bug 1 (OOB). Send concurrent traffic so that a scheduler step holds at least one running decode plus a new short prompt (total seq len ≤ index_topk = 2048). The dense-MHA split activates, forward_mqa gets only the decode tokens, and the index-conversion kernel writes past the end of its output buffer: non-deterministic corruption of neighbouring GPU allocations, or CUDA error: an illegal memory access. Also reproduces with --kv-cache-dtype auto (bf16 cache).

Bug 2 (context gather). Multi-turn chat or any prefix-cache hit, i.e. a prefill that arrives with computed context, total sequence ≤ 2048. Without DCP the request completes but its context K/V is garbage, so the output is silently wrong. With DCP it raises RuntimeError: src_cache and dst must have the same dtype.

Root cause 1 — unsliced req_id_per_token

FlashMLASparseImpl.forward_mqa slices topk_indices = self.topk_indices_buffer[:num_actual_toks] with num_actual_toks = q.shape[0] (the MQA subset), but all three forward paths pass the full-batch attn_metadata.req_id_per_token to triton_convert_req_index_to_global_index:

• vllm/v1/attention/backends/mla/flashmla_sparse.py:602 (_forward_bf16_kv), :643 (_forward_fp8_kv_separate_prefill_decode), :741 (_forward_fp8_kv_mixed_batch)

The converter sizes its grid by req_id.shape[0] but allocates out = torch.empty_like(token_indices) (vllm/v1/attention/backends/mla/sparse_utils.py:168,176), so rows [num_mqa_tokens, num_batch_tokens) are written past the end of out.

The MQA subset is exactly the leading num_decode_tokens tokens (the runner reorders decodes first, reorder_batch_to_split_decodes_and_prefills), so the fix is the same slice the FlashInfer and FlashAttn sparse backends already do (flashinfer_mla_sparse.py:397, flashattn_mla_sparse.py:227).

Root cause 2 — fp8_ds_mla context gather on the dense-MHA path

When the prefill has context (prefill.chunked_context is not None), forward_mha gathers cached KV (vllm/model_executor/layers/attention/mla_attention.py:2349,2377-2393). The fp8_ds_mla entry is 656B — 512B fp8 NoPE quantized with per-128-tile scales, 4×fp32 inline scales, 64×bf16 RoPE (see the writer, concat_and_cache_ds_mla_kernel, csrc/libtorch_stable/cache_kernels.cu:447) — and neither gather understands it:

• DCP (_context_parallel_compute_prefill_context, mla_attention.py:2216-2241): the dequant branch explicitly excludes fp8_ds_mla, so it falls through to raw ops.cp_gather_cache, which requires equal src/dst dtypes (cache_kernels.cu:1363-1364) — uint8 cache vs bf16 workspace → RuntimeError.
• non-DCP (_compute_prefill_context, mla_attention.py:2102-2112): gather_and_maybe_dequant_cache(kv_cache_dtype="fp8_ds_mla") dispatches fp8_ds_mla as plain E4M3 (csrc/attention/dtype_fp8.cuh:28-30) and dequantizes the first 576 bytes of each 656B entry element-wise with a single global scale (cache_kernels.cu:1032-1051): the NoPE bytes get the wrong scale, the inline scale words and part of the bf16 RoPE bytes are decoded as fp8 values, and the RoPE tail is never read. Nothing raises — the context K/V is silently corrupted.

Short no-context prefills — the case #47327 optimizes and benchmarks — never read the cache, so they are unaffected and should keep the fast path.

Additional fallout found while triaging

With fp8_ds_mla and ≥32 heads per rank (e.g. DSV3.2 at TP1/2/4), FlashMLA takes _forward_fp8_kv_separate_prefill_decode, whose FP8 metadata is still built for the full batch (_build_fp8_separate_prefill_decode, flashmla_sparse.py:353-494). Under the dense-MHA split it allocates a full-batch output and iterates prefill chunks whose tokens_slice lies beyond the decode-only q → shape-mismatch RuntimeError, even without context. Making that path subset-aware (or gating use_mha for it) needs its own change; noting it here so it is not lost.

Environment

Found on vLLM main (793cf79c8) while rebasing #46514. Analysis is static plus code reading; a GPU repro of the OOB with compute-sanitizer is the obvious confirmation and we are running it tonight.

AI assistance (Claude) was used for this analysis; the submitter reviewed every claim against the source.