diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py new file mode 100644 index 00000000000..6e1e10e8b45 --- /dev/null +++ b/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 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1dea276d2d2..16fb961e563 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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 diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 19cad7986bf..1463f9fb35b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/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) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index ca4a20874a0..ac956461415 100644 --- a/vllm/model_executor/layers/mla.py +++ b/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( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 5b8e2bf008e..9a671be5563 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/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, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 4b92e351caa..bf67e040a15 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/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( diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index dcf955ad59b..0e604da87ba 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/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, )