diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index cc9ac6b8d7..aa60362641 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -88,6 +88,81 @@ def _make_common_metadata(): ) +def _make_mixed_common_metadata(): + # req0: query_len 1 (decode), req1: query_len 4 (prefill) -> decode-first + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32, device="cpu") + seq_lens = torch.tensor([16, 10], dtype=torch.int32, device="cpu") + return CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + _seq_lens_cpu=seq_lens, + num_reqs=2, + num_actual_tokens=5, + max_query_len=4, + max_seq_len=16, + block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8), + slot_mapping=torch.arange(5, dtype=torch.int64, device="cpu"), + ) + + +def _patch_build_deps(monkeypatch, events=None): + """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs + on CPU.""" + + def fake_generate_sparse_seqlen_triton( + query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len + ): + return torch.zeros(num_tokens, dtype=torch.int32, device="cpu") + + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = Mock(side_effect=lambda *a, **k: None) + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton + ) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace( + synchronize=lambda: events.append("sync") if events is not None else None + ), + ) + + +def test_build_populates_decode_only_split_fields(monkeypatch): + """Decode-only batch: all reqs count as decodes, prefill fields default.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_common_metadata() + ) + + assert md.num_decodes == 2 + assert md.num_prefills == 0 + assert md.num_decode_tokens == 2 + assert md.prefill_max_seq_len == 0 + assert md.prefill is None + + +def test_build_populates_mixed_split_fields(monkeypatch): + """Mixed decode+prefill batch: split is reported, prefill fields stay + default because this impl always runs the MQA path.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_mixed_common_metadata() + ) + + assert md.num_decodes == 1 + assert md.num_prefills == 1 + assert md.num_decode_tokens == 1 + assert md.prefill_max_seq_len == 0 + assert md.prefill is None + + def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch): builder = _make_builder() common_metadata = _make_common_metadata() diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 884c4ca4e6..1dea276d2d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -515,29 +515,36 @@ class MLAAttention(nn.Module, AttentionLayerBase): compilation_config.static_forward_context[prefix] = self self.prefill_backend: MLAPrefillBackend | None - try: - prefill_backend_cls = get_mla_prefill_backend(vllm_config) - except ValueError: - if ( - not self.impl.is_sparse - or vllm_config.attention_config.mla_prefill_backend is not None - ): - raise + if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill: logger.warning_once( - "No MLA prefill backend supports this model; sparse MLA will use the " - "top-k MQA path only (no dense-MHA prefill)." + "Sparse MLA impl has no dense-MHA prefill path; using the top-k " + "MQA path only." ) self.prefill_backend = None else: - self.prefill_backend = prefill_backend_cls( - num_heads=self.num_heads, - scale=self.scale, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - vllm_config=vllm_config, - ) + try: + prefill_backend_cls = get_mla_prefill_backend(vllm_config) + except ValueError: + if ( + not self.impl.is_sparse + or vllm_config.attention_config.mla_prefill_backend is not None + ): + raise + logger.warning_once( + "No MLA prefill backend supports this model; sparse MLA will " + "use the top-k MQA path only (no dense-MHA prefill)." + ) + self.prefill_backend = None + else: + self.prefill_backend = prefill_backend_cls( + num_heads=self.num_heads, + scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + vllm_config=vllm_config, + ) self.kv_cache = torch.tensor([]) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 2e30fd4aa2..4b92e351ca 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -851,11 +851,16 @@ def _try_load_fp8_indexer_wk( # We have both weight and scale: dequantize FP8 to BF16. weight_fp8, scale_inv = entry["weight"], entry["scale"] del buf[layer_prefix] - block_size = weight_fp8.shape[1] // scale_inv.shape[1] + if scale_inv.ndim == 1: + # Per-channel scale: one scale per row of [out, in] + group_shape = GroupShape(1, weight_fp8.shape[1]) + else: + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + group_shape = GroupShape(block_size, block_size) weight_bf16 = scaled_dequantize( weight_fp8, scale_inv, - group_shape=GroupShape(block_size, block_size), + group_shape=group_shape, out_dtype=torch.bfloat16, ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index b479aab581..b5d8962fd4 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -803,6 +803,10 @@ class AttentionImplBase(ABC, Generic[T]): # route between the dense-MHA prefill and sparse-MQA paths. is_sparse: ClassVar[bool] = False + # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse + # impls without one run the top-k MQA path for all requests. + supports_dense_mha_prefill: ClassVar[bool] = True + # Required attributes that all impls should have num_heads: int head_size: int diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 2bbaaad765..d5db42cb7f 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -30,6 +30,7 @@ from vllm.v1.attention.backend import ( from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( AiterMLAHelper, ) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.kv_cache_interface import AttentionSpec from vllm.v1.worker.workspace import current_workspace_manager @@ -333,6 +334,14 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata): block_size: int = 1 topk_tokens: int = 2048 + # Fields read by the shared MLA forward. This impl has no dense-MHA prefill + # path (supports_dense_mha_prefill=False), so it always runs the MQA path; + num_decodes: int = 0 + num_prefills: int = 0 + num_decode_tokens: int = 0 + prefill_max_seq_len: int = 0 + prefill: object = None + # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). work_meta_data: torch.Tensor | None = None @@ -488,6 +497,10 @@ class ROCMAiterMLASparseMetadataBuilder( fast_build: bool = False, ) -> ROCMAiterMLASparseMetadata: num_tokens = common_attn_metadata.num_actual_tokens + (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold or 1, + ) starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) seg_lengths = np.diff(starts) req_id_per_token = np.repeat( @@ -604,6 +617,9 @@ class ROCMAiterMLASparseMetadataBuilder( block_size=self.kv_cache_spec.block_size, attn_out_dtype=self.model_dtype, topk_tokens=self.topk_tokens, + num_decodes=num_decodes, + num_prefills=num_prefills, + num_decode_tokens=num_decode_tokens, qo_indptr=qo_indptr, paged_kv_last_page_len=paged_kv_last_page_len, paged_kv_indices=paged_kv_indices, @@ -649,6 +665,7 @@ def reference_mla_sparse_prefill( class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]): is_sparse = True + supports_dense_mha_prefill = False def __init__( self,