A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 17:15 UTC
[ROCm][perf] fix Aiter sparse MLA with MTP>1 (#37887)
Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Signed-off-by: Stig-Arne Grönroos <sgronroo@amd.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Stig-Arne Grönroos committed
Apr 1, 2026 at 02:22 UTC
31a719bcd37a195107711dc8b498288e49ef8576
2 files changed
+27
-25
vllm/v1/spec_decode/dflash.py
+5
-5
@@ -257,19 +257,19 @@ class DFlashProposer(SpecDecodeBaseProposer):
257
)
258
259
@override
260
- def build_per_layer_attn_metadata(
260
+ def build_per_group_and_layer_attn_metadata(
261
self, cad: CommonAttentionMetadata, draft_index: int = 0
262
- ) -> dict[str, object]:
263
- per_layer_attention_metadata = super().build_per_layer_attn_metadata(
262
+ ) -> tuple[list[object], dict[str, object]]:
263
+ per_group, per_layer = super().build_per_group_and_layer_attn_metadata(
264
cad, draft_index
265
)
266
- for layer_name, attn_metadata in per_layer_attention_metadata.items():
266
+ for layer_name, attn_metadata in per_layer.items():
267
assert getattr(attn_metadata, "causal", None) is False, (
268
f"Attention metadata for layer {layer_name} does not have"
269
" non-causal support, which is required for DFlash."
270
" Consider using a different attention backend, such as FlashAttention."
271
)
272
- return per_layer_attention_metadata
272
+ return per_group, per_layer
273
274
@override
275
def _get_eagle3_use_aux_hidden_state_from_config(self):
vllm/v1/spec_decode/eagle.py
+22
-20
@@ -225,6 +225,9 @@ class SpecDecodeBaseProposer:
225
# Determine allowed attention backends once during initialization.
226
self.allowed_attn_types: tuple | None = None
227
if current_platform.is_rocm():
228
+ from vllm.v1.attention.backends.mla.indexer import (
229
+ DeepseekV32IndexerMetadata,
230
+ )
231
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
232
ROCMAiterMLASparseMetadata,
233
)
@@ -234,6 +237,7 @@ class SpecDecodeBaseProposer:
237
TritonAttentionMetadata,
238
RocmAttentionMetadata,
239
ROCMAiterMLASparseMetadata,
240
+ DeepseekV32IndexerMetadata,
241
]
242
# ROCM_AITER_FA is an optional backend
243
# We check is_enabled() here to avoid importing the backend module during
@@ -444,8 +448,8 @@ class SpecDecodeBaseProposer:
448
)
449
)
450
447
- per_layer_attn_metadata = self.build_per_layer_attn_metadata(
448
- common_attn_metadata
451
+ per_group_attn_metadata, per_layer_attn_metadata = (
452
+ self.build_per_group_and_layer_attn_metadata(common_attn_metadata)
453
)
454
455
cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
@@ -486,10 +490,7 @@ class SpecDecodeBaseProposer:
490
positions = self.positions[token_indices_to_sample]
491
hidden_states = hidden_states[token_indices_to_sample]
492
489
- if any(
490
- isinstance(attn_metadata, TreeAttentionMetadata)
491
- for attn_metadata in per_layer_attn_metadata.values()
492
- ):
493
+ if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata):
494
# Draft using tree attention - requires full logits for top-k
495
logits = self.model.compute_logits(sample_hidden_states)
496
draft_token_ids_list = self.propose_tree(
@@ -505,16 +506,15 @@ class SpecDecodeBaseProposer:
506
507
draft_token_ids = self._greedy_sample(sample_hidden_states)
508
508
- for attn_metadata in per_layer_attn_metadata.values():
509
- if self.allowed_attn_types is not None and not isinstance(
510
- attn_metadata, self.allowed_attn_types
511
- ):
512
- raise ValueError(
513
- f"Unsupported attention metadata type for speculative "
514
- "decoding with num_speculative_tokens > 1: "
515
- f"{type(attn_metadata)}. Supported types are: "
516
- f"{self.allowed_attn_types}"
517
- )
509
+ if self.allowed_attn_types is not None:
510
+ for group_md in per_group_attn_metadata:
511
+ if not isinstance(group_md, self.allowed_attn_types):
512
+ raise ValueError(
513
+ f"Unsupported attention metadata type for speculative "
514
+ "decoding with num_speculative_tokens > 1: "
515
+ f"{type(group_md)}. Supported types are: "
516
+ f"{self.allowed_attn_types}"
517
+ )
518
519
# Generate the remaining draft tokens.
520
draft_token_ids_list = [draft_token_ids]
@@ -595,7 +595,7 @@ class SpecDecodeBaseProposer:
595
common_attn_metadata._num_computed_tokens_cpu += 1
596
597
# Rebuild attention metadata
598
- per_layer_attn_metadata = self.build_per_layer_attn_metadata(
598
+ _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(
599
common_attn_metadata, draft_index=token_index + 1
600
)
601
@@ -809,17 +809,19 @@ class SpecDecodeBaseProposer:
809
810
return model_kwargs, num_input_tokens
811
812
- def build_per_layer_attn_metadata(
812
+ def build_per_group_and_layer_attn_metadata(
813
self, common_attn_metadata: CommonAttentionMetadata, draft_index: int = 0
814
- ) -> dict[str, object]:
814
+ ) -> tuple[list[object], dict[str, object]]:
815
+ per_group_attn_metadata: list[object] = []
816
per_layer_attn_metadata: dict[str, object] = {}
817
for attn_group in self.draft_attn_groups:
818
attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
819
common_attn_metadata=common_attn_metadata, draft_index=draft_index
820
)
821
+ per_group_attn_metadata.append(attn_metadata)
822
for layer_name in attn_group.layer_names:
823
per_layer_attn_metadata[layer_name] = attn_metadata
822
- return per_layer_attn_metadata
824
+ return per_group_attn_metadata, per_layer_attn_metadata
825
826
def model_returns_tuple(self) -> bool:
827
return self.method not in ("mtp", "draft_model", "dflash")