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

tests/v1/attention/test_sparse_mla_backends.py

@@ -40,6 +40,7 @@
 )
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     FlashMLASparseBackend,
+    FlashMLASparseImpl,
     triton_convert_req_index_to_global_index,
 )
 from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks
@@ -721,6 +722,120 @@ def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_s
     torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
 
 
+@pytest.mark.skipif(
+    torch.cuda.get_device_capability() < (9, 0),
+    reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
+)
+def test_triton_convert_rejects_req_id_longer_than_token_indices():
+    """Guard against the #47327 regression: the kernel grid is sized by
+    req_id but the output is allocated like token_indices, so a full-batch
+    req_id combined with an MQA-subset token_indices wrote past the end of
+    the output buffer. The wrapper must reject the length mismatch instead
+    of corrupting memory."""
+    device = torch.device(DEVICE_TYPE)
+    num_topk_tokens = 128
+    block_size = 64
+    block_table = torch.arange(40, dtype=torch.int32, device=device).view(4, 10)
+
+    # Full batch: 2 decode tokens + 10 prefill tokens
+    req_id_full = torch.tensor(
+        [0, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3], dtype=torch.int32, device=device
+    )
+    num_mqa_tokens = 2
+    token_indices = torch.randint(
+        0,
+        block_size * 10,
+        (num_mqa_tokens, num_topk_tokens),
+        dtype=torch.int32,
+        device=device,
+    )
+
+    with pytest.raises(AssertionError, match="must cover the same tokens"):
+        triton_convert_req_index_to_global_index(
+            req_id_full,
+            block_table,
+            token_indices,
+            BLOCK_SIZE=block_size,
+            NUM_TOPK_TOKENS=num_topk_tokens,
+        )
+
+    # The sliced call is the intended usage and must match the reference.
+    result = triton_convert_req_index_to_global_index(
+        req_id_full[:num_mqa_tokens],
+        block_table,
+        token_indices,
+        BLOCK_SIZE=block_size,
+        NUM_TOPK_TOKENS=num_topk_tokens,
+    )
+    reference = _triton_convert_reference_impl(
+        req_id_full[:num_mqa_tokens],
+        block_table,
+        token_indices,
+        block_size,
+        num_topk_tokens,
+    )
+    torch.testing.assert_close(result, reference, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    torch.cuda.get_device_capability() < (9, 0),
+    reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
+)
+def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
+    """Guard against the #47327 regression: when the dense-MHA prefill split
+    is active, forward_mqa only receives the leading decode tokens, but
+    _forward_bf16_kv passed the full-batch req_id_per_token to the index
+    conversion, making it write past the end of its output buffer. The call
+    site must slice req_id_per_token to the MQA tokens."""
+    device = torch.device(DEVICE_TYPE)
+    num_topk_tokens = 128
+    block_size = 64
+    num_batch_tokens = 12
+    num_mqa_tokens = 2
+
+    attn_metadata = SimpleNamespace(
+        req_id_per_token=torch.tensor(
+            [0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32, device=device
+        ),
+        block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10),
+        block_size=block_size,
+    )
+    assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens
+
+    q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device)
+    kv_cache = torch.zeros(40 * block_size, 576, dtype=torch.bfloat16, device=device)
+    topk_indices = torch.randint(
+        0,
+        block_size * 10,
+        (num_mqa_tokens, num_topk_tokens),
+        dtype=torch.int32,
+        device=device,
+    )
+
+    captured = {}
+
+    def _stub_kernel(q, kv, indices, lengths):
+        captured["indices"] = indices
+        return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device)
+
+    stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel)
+
+    out = FlashMLASparseImpl._forward_bf16_kv(
+        stub_impl, q, kv_cache, topk_indices, attn_metadata
+    )
+
+    assert out.shape[0] == num_mqa_tokens
+    assert captured["indices"].shape[0] == num_mqa_tokens
+    reference = _triton_convert_reference_impl(
+        attn_metadata.req_id_per_token[:num_mqa_tokens],
+        attn_metadata.block_table,
+        topk_indices,
+        block_size,
+        num_topk_tokens,
+    )
+    torch.testing.assert_close(captured["indices"], reference, rtol=0, atol=0)
+
+
 @pytest.mark.parametrize(
     "seq_lens,max_buf,expected",
     [
@@ -1124,3 +1239,33 @@ def test_triton_convert_returns_valid_counts():
     )
     assert isinstance(result_only, torch.Tensor)
     torch.testing.assert_close(result_only, result, rtol=0, atol=0)
+
+
+def test_fp8_ds_mla_context_prefill_stays_on_mqa_path():
+    """Guard against the #47327 regression: the dense-MHA prefill path cannot
+    gather cached context from an fp8_ds_mla cache (cp_gather_cache rejects
+    the uint8/bf16 dtype mismatch under DCP; gather_and_maybe_dequant_cache
+    silently dequantizes the 656B entries as flat fp8 otherwise), so use_mha
+    must be blocked exactly when kv cache is fp8_ds_mla and the prefill has
+    chunked context — and stay available when there is no context, keeping
+    the fast path for short no-context prefills."""
+    from vllm.model_executor.layers.attention.mla_attention import (
+        _canonicalize_sparse_mla_kv_cache_dtype,
+        _dense_mha_context_gather_unsupported,
+    )
+
+    with_context = SimpleNamespace(chunked_context=object())
+    no_context = SimpleNamespace(chunked_context=None)
+
+    assert _dense_mha_context_gather_unsupported("fp8_ds_mla", with_context)
+    assert not _dense_mha_context_gather_unsupported("fp8_ds_mla", no_context)
+    assert not _dense_mha_context_gather_unsupported("fp8_ds_mla", None)
+    assert not _dense_mha_context_gather_unsupported("auto", with_context)
+
+    # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
+    # stores kv_cache_dtype, so they cannot bypass the gate.
+    for alias in ("fp8", "fp8_e4m3"):
+        assert (
+            _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)
+            == "fp8_ds_mla"
+        )

vllm/model_executor/layers/attention/mla_attention.py

@@ -335,6 +335,24 @@ def _canonicalize_sparse_mla_kv_cache_dtype(
     return kv_cache_dtype
 
 
+def _dense_mha_context_gather_unsupported(
+    kv_cache_dtype: str,
+    prefill: "MLACommonPrefillMetadata | None",
+) -> bool:
+    """Whether the dense-MHA prefill path would have to gather cached context
+    it cannot read. The chunked-context gather kernels do not understand the
+    fp8_ds_mla cache layout (656B entries: fp8 NoPE + inline tile scales +
+    bf16 RoPE): cp_gather_cache rejects the uint8/bf16 dtype mismatch and
+    gather_and_maybe_dequant_cache dequantizes the raw bytes as flat fp8,
+    producing corrupted K/V. Such prefills must stay on the top-k MQA path.
+    """
+    return (
+        kv_cache_dtype == "fp8_ds_mla"
+        and prefill is not None
+        and prefill.chunked_context is not None
+    )
+
+
 class MLAAttention(nn.Module, AttentionLayerBase):
     """Multi-Head Latent Attention layer.
 
@@ -720,6 +738,9 @@ def forward_impl(
                 self.prefill_backend is not None
                 and prefill_max_seq_len <= attn_metadata.topk_tokens  # type: ignore[attr-defined]
                 and not self._vllm_config.attention_config.sparse_mla_force_mqa
+                and not _dense_mha_context_gather_unsupported(
+                    self.kv_cache_dtype, attn_metadata.prefill
+                )
             )
             if not use_mha:
                 num_mqa_tokens = q.size(0)

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -597,9 +597,10 @@ def _forward_bf16_kv(
         attn_metadata: FlashMLASparseMetadata,
     ) -> torch.Tensor:
         # Convert per-request indices to global slots (decode) or workspace
-        # offsets (prefill).
+        # offsets (prefill). req_id_per_token covers the whole batch; slice it
+        # to the MQA tokens (q may exclude prefill tokens routed to dense MHA).
         topk_indices, topk_length = triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token,
+            attn_metadata.req_id_per_token[: topk_indices.shape[0]],
             attn_metadata.block_table,
             topk_indices,
             BLOCK_SIZE=attn_metadata.block_size,
@@ -640,7 +641,7 @@ def _forward_fp8_kv_separate_prefill_decode(
         # prefill_workspace_starts has been adjusted in-place per chunk so
         # prefill indices automatically come out chunk-local
         topk_indices, topk_length = triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token,
+            attn_metadata.req_id_per_token[: topk_indices.shape[0]],
             attn_metadata.block_table,
             topk_indices,
             BLOCK_SIZE=attn_metadata.block_size,
@@ -738,7 +739,7 @@ def _forward_fp8_kv_mixed_batch(
         # Convert per-request indices to global slots (decode) or workspace
         # offsets (prefill).
         topk_indices = triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token,
+            attn_metadata.req_id_per_token[: topk_indices.shape[0]],
             attn_metadata.block_table,
             topk_indices,
             BLOCK_SIZE=attn_metadata.block_size,

vllm/v1/attention/backends/mla/sparse_utils.py

@@ -154,6 +154,11 @@ def triton_convert_req_index_to_global_index(
     assert req_id.dtype == torch.int32
     assert block_table.dtype == torch.int32
     assert token_indices.dtype == torch.int32
+    assert req_id.shape[0] == token_indices.shape[0], (
+        f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) "
+        "must cover the same tokens; the grid is sized by req_id but the output "
+        "is allocated like token_indices, so a longer req_id writes out of bounds"
+    )
     assert token_indices.shape[1] == NUM_TOPK_TOKENS
     assert NUM_TOPK_TOKENS % BLOCK_N == 0, (
         f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"