[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>
Yiliu Dong committed
Jul 29, 2026 at 00:17 UTC
ba702e978e3bc6af3a601cee10fefdeb49e7e8b5
7 files changed
+230
-12
tests/model_executor/layers/test_mla_short_prefill_indexer.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
10
+from vllm.config import CUDAGraphMode
11
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
12
+
13
+INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache"
14
+MLA_LAYER = "model.layers.0.self_attn.attn"
15
+
16
+
17
+def make_indexer_metadata(
18
+ *,
19
+ num_decodes: int = 0,
20
+ num_decode_tokens: int = 0,
21
+ num_prefills: int = 1,
22
+ num_prefill_tokens: int = 1,
23
+ slot_mapping: torch.Tensor | None = None,
24
+) -> DeepseekV32IndexerMetadata:
25
+ if slot_mapping is None:
26
+ slot_mapping = torch.zeros(num_prefill_tokens, dtype=torch.long)
27
+ return DeepseekV32IndexerMetadata(
28
+ seq_lens=torch.empty(0, dtype=torch.int32),
29
+ max_seq_len=2048,
30
+ slot_mapping=slot_mapping,
31
+ num_decodes=num_decodes,
32
+ num_decode_tokens=num_decode_tokens,
33
+ num_prefills=num_prefills,
34
+ num_prefill_tokens=num_prefill_tokens,
35
+ prefill=SimpleNamespace(chunks=[]) if num_prefills else None,
36
+ )
37
+
38
+
39
+def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0):
40
+ return SimpleNamespace(
41
+ num_decode_tokens=num_decode_tokens,
42
+ prefill=SimpleNamespace(use_dense_mha=use_dense_mha),
43
+ )
44
+
45
+
46
+@pytest.mark.parametrize(
47
+ "batch_kind",
48
+ ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"],
49
+)
50
+def test_short_prefill_updates_k_cache_before_scoring_decision(
51
+ monkeypatch: pytest.MonkeyPatch,
52
+ batch_kind: str,
53
+):
54
+ slot_mapping = torch.tensor([63, 64, 127, 128, -1])
55
+ mla_num_decode_tokens = 1 if batch_kind == "mla_decode" else 0
56
+ runtime_mode = (
57
+ CUDAGraphMode.FULL if batch_kind == "full" else CUDAGraphMode.PIECEWISE
58
+ )
59
+ should_skip = batch_kind in ("short", "threshold_mismatch")
60
+ num_decodes = int(batch_kind == "threshold_mismatch")
61
+ num_decode_tokens = 3 if batch_kind == "threshold_mismatch" else 0
62
+ num_prefills = 0 if batch_kind == "threshold_mismatch" else 2
63
+ num_prefill_tokens = 0 if batch_kind == "threshold_mismatch" else 5
64
+ if batch_kind == "threshold_mismatch":
65
+ # With MTP=3 the indexer threshold is four. A main MLA backend whose
66
+ # threshold is one (for example FlashMLA under DCP) still routes this
67
+ # three-token extend through dense prefill attention.
68
+ slot_mapping = slot_mapping[:3]
69
+ indexer_metadata = make_indexer_metadata(
70
+ num_decodes=num_decodes,
71
+ num_decode_tokens=num_decode_tokens,
72
+ num_prefills=num_prefills,
73
+ num_prefill_tokens=num_prefill_tokens,
74
+ slot_mapping=slot_mapping,
75
+ )
76
+ if indexer_metadata.num_decodes:
77
+ indexer_metadata.decode = object()
78
+ mla_metadata = make_mla_metadata(
79
+ use_dense_mha=batch_kind != "force_mqa",
80
+ num_decode_tokens=mla_num_decode_tokens,
81
+ )
82
+
83
+ observed: dict[str, object] = {}
84
+
85
+ monkeypatch.setattr(
86
+ sparse_indexer,
87
+ "get_forward_context",
88
+ lambda: SimpleNamespace(
89
+ attn_metadata={
90
+ INDEXER_LAYER: indexer_metadata,
91
+ MLA_LAYER: mla_metadata,
92
+ },
93
+ cudagraph_runtime_mode=runtime_mode,
94
+ ),
95
+ )
96
+ monkeypatch.setattr(
97
+ sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16
98
+ )
99
+ monkeypatch.setattr(
100
+ torch.cuda,
101
+ "is_current_stream_capturing",
102
+ lambda: batch_kind == "capture",
103
+ )
104
+
105
+ def record_cache_update(k, kv_cache, slots, block_size, scale_fmt):
106
+ observed.update(k=k.clone(), slots=slots)
107
+
108
+ monkeypatch.setattr(
109
+ sparse_indexer.ops, "indexer_k_quant_and_cache", record_cache_update
110
+ )
111
+
112
+ class ScoringReached(Exception):
113
+ pass
114
+
115
+ def scoring_trigger():
116
+ if should_skip:
117
+ pytest.fail("short dense-MHA prefill must not enter indexer scoring")
118
+ raise ScoringReached
119
+
120
+ def scoring_decode(*args):
121
+ raise ScoringReached
122
+
123
+ monkeypatch.setattr(sparse_indexer, "current_workspace_manager", scoring_trigger)
124
+ monkeypatch.setattr(
125
+ sparse_indexer,
126
+ "kv_cache_as_quant_view",
127
+ scoring_decode,
128
+ )
129
+
130
+ hidden_states = torch.full((7, 1), float("inf"))
131
+ k = torch.arange(28, dtype=torch.float32).reshape(7, 4)
132
+ topk_indices = torch.full((7, 2048), 17, dtype=torch.int32)
133
+
134
+ def run_indexer():
135
+ return sparse_indexer.sparse_attn_indexer(
136
+ hidden_states,
137
+ INDEXER_LAYER,
138
+ torch.empty(1),
139
+ torch.full((7, 1), float("inf")),
140
+ None,
141
+ k,
142
+ torch.full((7, 1), float("inf")),
143
+ 128,
144
+ "ue8m0",
145
+ 2048,
146
+ 4,
147
+ 4096,
148
+ 4096,
149
+ topk_indices,
150
+ False,
151
+ False,
152
+ MLA_LAYER,
153
+ )
154
+
155
+ if should_skip:
156
+ assert run_indexer() is topk_indices
157
+ assert torch.all(topk_indices == 17)
158
+ else:
159
+ with pytest.raises(ScoringReached):
160
+ run_indexer()
161
+ assert torch.all(topk_indices == -1)
162
+
163
+ # K cache is always updated before the scoring decision.
164
+ torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()])
165
+ assert observed["slots"] is slot_mapping
vllm/model_executor/layers/attention/mla_attention.py
+6
-7
@@ -769,13 +769,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
769
num_mha_tokens = q.size(0) - num_mqa_tokens
770
771
if self.impl.is_sparse and num_mha_tokens > 0:
772
- prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]
773
- use_mha = (
774
- self.prefill_backend is not None
775
- and prefill_max_seq_len <= attn_metadata.topk_tokens # type: ignore[attr-defined]
776
- and not self._vllm_config.attention_config.sparse_mla_force_mqa
777
- )
778
- if not use_mha:
772
+ prefill_metadata = getattr(attn_metadata, "prefill", None)
773
+ if not getattr(prefill_metadata, "use_dense_mha", False):
774
num_mqa_tokens = q.size(0)
775
num_mha_tokens = 0
776
@@ -1409,6 +1404,10 @@ class MLACommonPrefillMetadata:
1404
q_data_type: torch.dtype | None = None
1405
output_dtype: torch.dtype | None = None
1406
prefill_backend: MLAPrefillBackend | None = None
1407
+ # Whether the prefill suffix is routed through dense MHA.
1408
+ # Indexer scoring may be skipped only for a pure-prefill batch,
1409
+ # since decode tokens still consume top-k indices.
1410
+ use_dense_mha: bool = False
1411
1412
1413
@dataclass
vllm/model_executor/layers/attention/sparse_mla_attention.py
+4
@@ -203,6 +203,10 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
203
q_data_type=self.model_config.dtype,
204
output_dtype=self.model_config.dtype,
205
prefill_backend=self._prefill_backend,
206
+ use_dense_mha=(
207
+ prefill_max_seq_len <= self.topk_tokens
208
+ and not self.vllm_config.attention_config.sparse_mla_force_mqa
209
+ ),
210
)
211
self._prefill_backend.prepare_metadata(prefill)
212
vllm/model_executor/layers/mla.py
+22
-1
@@ -8,6 +8,7 @@ from vllm.config import CacheConfig
8
from vllm.model_executor.custom_op import PluggableLayer
9
from vllm.model_executor.layers.attention import MLAAttention
10
from vllm.model_executor.layers.quantization import QuantizationConfig
11
+from vllm.platforms import current_platform
12
13
14
@dataclass
@@ -65,6 +66,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
66
quant_config: QuantizationConfig | None = None,
67
prefix: str = "",
68
skip_topk: bool = False,
69
+ allow_short_prefill_indexer_scoring_skip: bool = False,
70
) -> None:
71
super().__init__()
72
self.hidden_size = hidden_size
@@ -119,7 +121,26 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
121
indexer=self.indexer,
122
topk_indices_buffer=mla_modules.topk_indices_buffer,
123
)
122
-
124
+ indexer_op = getattr(self.indexer, "indexer_op", None)
125
+ if indexer_op is not None and hasattr(
126
+ indexer_op, "dense_mha_metadata_layer_name"
127
+ ):
128
+ enable_short_prefill_scoring_skip = (
129
+ allow_short_prefill_indexer_scoring_skip
130
+ and not self.skip_topk
131
+ and not getattr(indexer_op, "use_pcp", False)
132
+ and current_platform.is_cuda()
133
+ )
134
+ # The indexer and main MLA use independent decode thresholds and
135
+ # may classify the same short extend differently. Bind the main
136
+ # MLA layer name so the eager indexer op can check whether the
137
+ # batch's top-k indices will be consumed.
138
+ # PCP is excluded because indexer cache/scoring ownership differs
139
+ # across ranks and the no-consumer invariant has not been
140
+ # established there.
141
+ indexer_op.dense_mha_metadata_layer_name = (
142
+ self.mla_attn.layer_name if enable_short_prefill_scoring_skip else ""
143
+ )
144
self.prefix = prefix
145
146
def forward(
vllm/model_executor/layers/sparse_attn_indexer.py
+26
-2
@@ -8,7 +8,7 @@ import vllm.envs as envs
8
from vllm import _custom_ops as ops
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11
-from vllm.config import get_current_vllm_config
11
+from vllm.config import CUDAGraphMode, get_current_vllm_config
12
from vllm.distributed import get_dcp_group, get_pcp_group
13
from vllm.forward_context import get_forward_context
14
from vllm.logger import init_logger
@@ -310,6 +310,7 @@ def sparse_attn_indexer(
310
topk_indices_buffer: torch.Tensor,
311
skip_k_cache_insert: bool,
312
use_pcp: bool,
313
+ dense_mha_metadata_layer_name: LayerNameType,
314
use_fp4_cache: bool = False,
315
dcp_rank: int = 0,
316
dcp_world_size: int = 1,
@@ -317,7 +318,8 @@ def sparse_attn_indexer(
318
skip_topk_buffer_clear: bool = False,
319
) -> torch.Tensor:
320
# careful! this will be None in dummy run
320
- attn_metadata = get_forward_context().attn_metadata
321
+ forward_context = get_forward_context()
322
+ attn_metadata = forward_context.attn_metadata
323
fp8_dtype = current_platform.fp8_dtype()
324
k_cache_prefix = _resolve_layer_name(k_cache_prefix)
325
@@ -357,6 +359,7 @@ def sparse_attn_indexer(
359
topk_indices_buffer,
360
skip_k_cache_insert,
361
use_pcp,
362
+ dense_mha_metadata_layer_name,
363
use_fp4_cache,
364
)
365
attn_metadata_narrowed = attn_metadata[k_cache_prefix]
@@ -402,6 +405,24 @@ def sparse_attn_indexer(
405
scale_fmt,
406
)
407
408
+ # The indexer and main MLA may classify the same short extend differently
409
+ # because they use independent decode thresholds. Only the main MLA route
410
+ # can determine whether the top-k indices will be consumed.
411
+ if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL:
412
+ dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name)
413
+ if dense_mha_layer:
414
+ mla_metadata = attn_metadata.get(dense_mha_layer)
415
+ prefill_metadata = getattr(mla_metadata, "prefill", None)
416
+ if (
417
+ getattr(prefill_metadata, "use_dense_mha", False)
418
+ and getattr(mla_metadata, "num_decode_tokens", -1) == 0
419
+ and not torch.cuda.is_current_stream_capturing()
420
+ ):
421
+ # Deliberately leave the buffer untouched. Dense MHA does not
422
+ # consume top-k indices for this batch; clearing it would be
423
+ # unnecessary work.
424
+ return topk_indices_buffer
425
+
426
# The buffer must be pre-filled with -1 (the "no token" sentinel) before the
427
# top-k kernels scatter valid indices into it. On the fused deepseek_v32
428
# nvidia path, _fused_norm_rope_kernel already cleared the same
@@ -684,6 +705,7 @@ def sparse_attn_indexer_fake(
705
topk_indices_buffer: torch.Tensor | None,
706
skip_k_cache_insert: bool,
707
use_pcp: bool,
708
+ dense_mha_metadata_layer_name: LayerNameType,
709
use_fp4_cache: bool = False,
710
dcp_rank: int = 0,
711
dcp_world_size: int = 1,
@@ -739,6 +761,7 @@ class SparseAttnIndexer(CustomOp):
761
self.topk_indices_buffer = topk_indices_buffer
762
self.skip_k_cache_insert = skip_k_cache_insert
763
self.use_fp4_cache = use_fp4_cache
764
+ self.dense_mha_metadata_layer_name = ""
765
# DCP scalars are constant for the run; resolve them here (config is set
766
# during model construction) and pass them into the custom op, rather
767
# than threading them through per-step metadata.
@@ -800,6 +823,7 @@ class SparseAttnIndexer(CustomOp):
823
self.topk_indices_buffer,
824
self.skip_k_cache_insert,
825
self.use_pcp,
826
+ _encode_layer_name(self.dense_mha_metadata_layer_name),
827
self.use_fp4_cache,
828
self.dcp_rank,
829
self.dcp_world_size,
vllm/model_executor/models/deepseek_v2.py
+3
@@ -1178,6 +1178,9 @@ class DeepseekV2MLAAttention(nn.Module):
1178
# the V1 proposer. A frozen True would leave the draft reading a
1179
# never-written topk buffer.
1180
skip_topk=_skip_topk and not is_mtp_layer,
1181
+ # Do not skip scoring for MTP layers: their top-k buffer may be
1182
+ # reused by later draft iterations through index sharing.
1183
+ allow_short_prefill_indexer_scoring_skip=not is_mtp_layer,
1184
)
1185
1186
def forward(
vllm/models/deepseek_v32/nvidia/attention.py
+4
-2
@@ -494,8 +494,10 @@ class DeepseekV32Attention(MLAAttention):
494
self.indexer.max_model_len,
495
self.indexer.max_total_seq_len,
496
self.topk_indices_buffer,
497
- True, # skip_k_cache_insert
498
- False, # use_fp4_cache
497
+ skip_k_cache_insert=True,
498
+ use_pcp=False,
499
+ dense_mha_metadata_layer_name="",
500
+ use_fp4_cache=False,
501
# fused_norm_rope already cleared the topk buffer this forward.
502
skip_topk_buffer_clear=True,
503
)