[Attention] Skip sparse indexer scoring for dense short prefills (#48407)

Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/model_executor/layers/test_mla_short_prefill_indexer.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
+from vllm.config import CUDAGraphMode
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
+
+INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache"
+MLA_LAYER = "model.layers.0.self_attn.attn"
+
+
+def make_indexer_metadata(
+    *,
+    num_decodes: int = 0,
+    num_decode_tokens: int = 0,
+    num_prefills: int = 1,
+    num_prefill_tokens: int = 1,
+    slot_mapping: torch.Tensor | None = None,
+) -> DeepseekV32IndexerMetadata:
+    if slot_mapping is None:
+        slot_mapping = torch.zeros(num_prefill_tokens, dtype=torch.long)
+    return DeepseekV32IndexerMetadata(
+        seq_lens=torch.empty(0, dtype=torch.int32),
+        max_seq_len=2048,
+        slot_mapping=slot_mapping,
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decode_tokens,
+        num_prefills=num_prefills,
+        num_prefill_tokens=num_prefill_tokens,
+        prefill=SimpleNamespace(chunks=[]) if num_prefills else None,
+    )
+
+
+def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0):
+    return SimpleNamespace(
+        num_decode_tokens=num_decode_tokens,
+        prefill=SimpleNamespace(use_dense_mha=use_dense_mha),
+    )
+
+
+@pytest.mark.parametrize(
+    "batch_kind",
+    ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"],
+)
+def test_short_prefill_updates_k_cache_before_scoring_decision(
+    monkeypatch: pytest.MonkeyPatch,
+    batch_kind: str,
+):
+    slot_mapping = torch.tensor([63, 64, 127, 128, -1])
+    mla_num_decode_tokens = 1 if batch_kind == "mla_decode" else 0
+    runtime_mode = (
+        CUDAGraphMode.FULL if batch_kind == "full" else CUDAGraphMode.PIECEWISE
+    )
+    should_skip = batch_kind in ("short", "threshold_mismatch")
+    num_decodes = int(batch_kind == "threshold_mismatch")
+    num_decode_tokens = 3 if batch_kind == "threshold_mismatch" else 0
+    num_prefills = 0 if batch_kind == "threshold_mismatch" else 2
+    num_prefill_tokens = 0 if batch_kind == "threshold_mismatch" else 5
+    if batch_kind == "threshold_mismatch":
+        # With MTP=3 the indexer threshold is four. A main MLA backend whose
+        # threshold is one (for example FlashMLA under DCP) still routes this
+        # three-token extend through dense prefill attention.
+        slot_mapping = slot_mapping[:3]
+    indexer_metadata = make_indexer_metadata(
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decode_tokens,
+        num_prefills=num_prefills,
+        num_prefill_tokens=num_prefill_tokens,
+        slot_mapping=slot_mapping,
+    )
+    if indexer_metadata.num_decodes:
+        indexer_metadata.decode = object()
+    mla_metadata = make_mla_metadata(
+        use_dense_mha=batch_kind != "force_mqa",
+        num_decode_tokens=mla_num_decode_tokens,
+    )
+
+    observed: dict[str, object] = {}
+
+    monkeypatch.setattr(
+        sparse_indexer,
+        "get_forward_context",
+        lambda: SimpleNamespace(
+            attn_metadata={
+                INDEXER_LAYER: indexer_metadata,
+                MLA_LAYER: mla_metadata,
+            },
+            cudagraph_runtime_mode=runtime_mode,
+        ),
+    )
+    monkeypatch.setattr(
+        sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16
+    )
+    monkeypatch.setattr(
+        torch.cuda,
+        "is_current_stream_capturing",
+        lambda: batch_kind == "capture",
+    )
+
+    def record_cache_update(k, kv_cache, slots, block_size, scale_fmt):
+        observed.update(k=k.clone(), slots=slots)
+
+    monkeypatch.setattr(
+        sparse_indexer.ops, "indexer_k_quant_and_cache", record_cache_update
+    )
+
+    class ScoringReached(Exception):
+        pass
+
+    def scoring_trigger():
+        if should_skip:
+            pytest.fail("short dense-MHA prefill must not enter indexer scoring")
+        raise ScoringReached
+
+    def scoring_decode(*args):
+        raise ScoringReached
+
+    monkeypatch.setattr(sparse_indexer, "current_workspace_manager", scoring_trigger)
+    monkeypatch.setattr(
+        sparse_indexer,
+        "kv_cache_as_quant_view",
+        scoring_decode,
+    )
+
+    hidden_states = torch.full((7, 1), float("inf"))
+    k = torch.arange(28, dtype=torch.float32).reshape(7, 4)
+    topk_indices = torch.full((7, 2048), 17, dtype=torch.int32)
+
+    def run_indexer():
+        return sparse_indexer.sparse_attn_indexer(
+            hidden_states,
+            INDEXER_LAYER,
+            torch.empty(1),
+            torch.full((7, 1), float("inf")),
+            None,
+            k,
+            torch.full((7, 1), float("inf")),
+            128,
+            "ue8m0",
+            2048,
+            4,
+            4096,
+            4096,
+            topk_indices,
+            False,
+            False,
+            MLA_LAYER,
+        )
+
+    if should_skip:
+        assert run_indexer() is topk_indices
+        assert torch.all(topk_indices == 17)
+    else:
+        with pytest.raises(ScoringReached):
+            run_indexer()
+        assert torch.all(topk_indices == -1)
+
+    # K cache is always updated before the scoring decision.
+    torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()])
+    assert observed["slots"] is slot_mapping

vllm/model_executor/layers/attention/mla_attention.py

@@ -769,13 +769,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         num_mha_tokens = q.size(0) - num_mqa_tokens
 
         if self.impl.is_sparse and num_mha_tokens > 0:
-            prefill_max_seq_len = attn_metadata.prefill_max_seq_len  # type: ignore[attr-defined]
-            use_mha = (
-                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
-            )
-            if not use_mha:
+            prefill_metadata = getattr(attn_metadata, "prefill", None)
+            if not getattr(prefill_metadata, "use_dense_mha", False):
                 num_mqa_tokens = q.size(0)
                 num_mha_tokens = 0
 
@@ -1409,6 +1404,10 @@ class MLACommonPrefillMetadata:
     q_data_type: torch.dtype | None = None
     output_dtype: torch.dtype | None = None
     prefill_backend: MLAPrefillBackend | None = None
+    # Whether the prefill suffix is routed through dense MHA.
+    # Indexer scoring may be skipped only for a pure-prefill batch,
+    # since decode tokens still consume top-k indices.
+    use_dense_mha: bool = False
 
 
 @dataclass

vllm/model_executor/layers/attention/sparse_mla_attention.py

@@ -203,6 +203,10 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
                 q_data_type=self.model_config.dtype,
                 output_dtype=self.model_config.dtype,
                 prefill_backend=self._prefill_backend,
+                use_dense_mha=(
+                    prefill_max_seq_len <= self.topk_tokens
+                    and not self.vllm_config.attention_config.sparse_mla_force_mqa
+                ),
             )
             self._prefill_backend.prepare_metadata(prefill)
 

vllm/model_executor/layers/mla.py

@@ -8,6 +8,7 @@ from vllm.config import CacheConfig
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.attention import MLAAttention
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.platforms import current_platform
 
 
 @dataclass
@@ -65,6 +66,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         skip_topk: bool = False,
+        allow_short_prefill_indexer_scoring_skip: bool = False,
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -119,7 +121,26 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
             indexer=self.indexer,
             topk_indices_buffer=mla_modules.topk_indices_buffer,
         )
-
+        indexer_op = getattr(self.indexer, "indexer_op", None)
+        if indexer_op is not None and hasattr(
+            indexer_op, "dense_mha_metadata_layer_name"
+        ):
+            enable_short_prefill_scoring_skip = (
+                allow_short_prefill_indexer_scoring_skip
+                and not self.skip_topk
+                and not getattr(indexer_op, "use_pcp", False)
+                and current_platform.is_cuda()
+            )
+            # The indexer and main MLA use independent decode thresholds and
+            # may classify the same short extend differently. Bind the main
+            # MLA layer name so the eager indexer op can check whether the
+            # batch's top-k indices will be consumed.
+            # PCP is excluded because indexer cache/scoring ownership differs
+            # across ranks and the no-consumer invariant has not been
+            # established there.
+            indexer_op.dense_mha_metadata_layer_name = (
+                self.mla_attn.layer_name if enable_short_prefill_scoring_skip else ""
+            )
         self.prefix = prefix
 
     def forward(

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -8,7 +8,7 @@ import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-from vllm.config import get_current_vllm_config
+from vllm.config import CUDAGraphMode, get_current_vllm_config
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
@@ -310,6 +310,7 @@ def sparse_attn_indexer(
     topk_indices_buffer: torch.Tensor,
     skip_k_cache_insert: bool,
     use_pcp: bool,
+    dense_mha_metadata_layer_name: LayerNameType,
     use_fp4_cache: bool = False,
     dcp_rank: int = 0,
     dcp_world_size: int = 1,
@@ -317,7 +318,8 @@ def sparse_attn_indexer(
     skip_topk_buffer_clear: bool = False,
 ) -> torch.Tensor:
     # careful! this will be None in dummy run
-    attn_metadata = get_forward_context().attn_metadata
+    forward_context = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
     fp8_dtype = current_platform.fp8_dtype()
     k_cache_prefix = _resolve_layer_name(k_cache_prefix)
 
@@ -357,6 +359,7 @@ def sparse_attn_indexer(
             topk_indices_buffer,
             skip_k_cache_insert,
             use_pcp,
+            dense_mha_metadata_layer_name,
             use_fp4_cache,
         )
     attn_metadata_narrowed = attn_metadata[k_cache_prefix]
@@ -402,6 +405,24 @@ def sparse_attn_indexer(
             scale_fmt,
         )
 
+    # The indexer and main MLA may classify the same short extend differently
+    # because they use independent decode thresholds. Only the main MLA route
+    # can determine whether the top-k indices will be consumed.
+    if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL:
+        dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name)
+        if dense_mha_layer:
+            mla_metadata = attn_metadata.get(dense_mha_layer)
+            prefill_metadata = getattr(mla_metadata, "prefill", None)
+            if (
+                getattr(prefill_metadata, "use_dense_mha", False)
+                and getattr(mla_metadata, "num_decode_tokens", -1) == 0
+                and not torch.cuda.is_current_stream_capturing()
+            ):
+                # Deliberately leave the buffer untouched. Dense MHA does not
+                # consume top-k indices for this batch; clearing it would be
+                # unnecessary work.
+                return topk_indices_buffer
+
     # The buffer must be pre-filled with -1 (the "no token" sentinel) before the
     # top-k kernels scatter valid indices into it. On the fused deepseek_v32
     # nvidia path, _fused_norm_rope_kernel already cleared the same
@@ -684,6 +705,7 @@ def sparse_attn_indexer_fake(
     topk_indices_buffer: torch.Tensor | None,
     skip_k_cache_insert: bool,
     use_pcp: bool,
+    dense_mha_metadata_layer_name: LayerNameType,
     use_fp4_cache: bool = False,
     dcp_rank: int = 0,
     dcp_world_size: int = 1,
@@ -739,6 +761,7 @@ class SparseAttnIndexer(CustomOp):
         self.topk_indices_buffer = topk_indices_buffer
         self.skip_k_cache_insert = skip_k_cache_insert
         self.use_fp4_cache = use_fp4_cache
+        self.dense_mha_metadata_layer_name = ""
         # DCP scalars are constant for the run; resolve them here (config is set
         # during model construction) and pass them into the custom op, rather
         # than threading them through per-step metadata.
@@ -800,6 +823,7 @@ class SparseAttnIndexer(CustomOp):
             self.topk_indices_buffer,
             self.skip_k_cache_insert,
             self.use_pcp,
+            _encode_layer_name(self.dense_mha_metadata_layer_name),
             self.use_fp4_cache,
             self.dcp_rank,
             self.dcp_world_size,

vllm/model_executor/models/deepseek_v2.py

@@ -1178,6 +1178,9 @@ class DeepseekV2MLAAttention(nn.Module):
             # the V1 proposer. A frozen True would leave the draft reading a
             # never-written topk buffer.
             skip_topk=_skip_topk and not is_mtp_layer,
+            # Do not skip scoring for MTP layers: their top-k buffer may be
+            # reused by later draft iterations through index sharing.
+            allow_short_prefill_indexer_scoring_skip=not is_mtp_layer,
         )
 
     def forward(

vllm/models/deepseek_v32/nvidia/attention.py

@@ -494,8 +494,10 @@ class DeepseekV32Attention(MLAAttention):
                 self.indexer.max_model_len,
                 self.indexer.max_total_seq_len,
                 self.topk_indices_buffer,
-                True,  # skip_k_cache_insert
-                False,  # use_fp4_cache
+                skip_k_cache_insert=True,
+                use_pcp=False,
+                dense_mha_metadata_layer_name="",
+                use_fp4_cache=False,
                 # fused_norm_rope already cleared the topk buffer this forward.
                 skip_topk_buffer_clear=True,
             )