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 7, 22:45 UTC
[ROCM] Optimize ROCM_AITER_FA spec decode eagle performance (#34541)
Signed-off-by: jennyyyyzhen <yzhen@hmc.edu>
jennyyyyzhen committed
Feb 20, 2026 at 20:32 UTC
2aab2bb54366c5a26add1b07f107b86f7fe28ff5
1 file changed
+50
-2
vllm/v1/attention/backends/rocm_aiter_fa.py
+50
-2
@@ -396,8 +396,7 @@ class AiterFlashAttentionMetadata:
396
class AiterFlashAttentionMetadataBuilder(
397
AttentionMetadataBuilder[AiterFlashAttentionMetadata]
398
):
399
- _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
400
- reorder_batch_threshold: int = 1
399
+ _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
400
401
def __init__(
402
self,
@@ -422,6 +421,7 @@ class AiterFlashAttentionMetadataBuilder(
421
# populated on first build() call.
422
self.aot_sliding_window: tuple[int, int] | None = None
423
self.total_tokens: int = 0
424
+ self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
425
426
sliding_window_configs: set[tuple[int, int] | None] = set()
427
layers = get_layers_from_vllm_config(self.vllm_config, Attention)
@@ -466,6 +466,7 @@ class AiterFlashAttentionMetadataBuilder(
466
common_attn_metadata: CommonAttentionMetadata,
467
fast_build: bool = False,
468
) -> "AiterFlashAttentionMetadata":
469
+ assert self.reorder_batch_threshold is not None
470
split_ret = split_decodes_prefills_and_extends(
471
common_attn_metadata,
472
decode_threshold=self.reorder_batch_threshold,
@@ -677,6 +678,53 @@ class AiterFlashAttentionMetadataBuilder(
678
)
679
return attn_metadata
680
681
+ def build_for_drafting(
682
+ self,
683
+ common_attn_metadata: CommonAttentionMetadata,
684
+ draft_index: int,
685
+ ) -> AiterFlashAttentionMetadata:
686
+ """
687
+ Build attention metadata for draft model without CPU-GPU sync.
688
+
689
+ During EAGLE drafting all requests are uniform decodes, so we can
690
+ skip split_decodes_prefills_and_extends() and avoid all .cpu() /
691
+ .item() calls that would otherwise break CUDA graph capture.
692
+ """
693
+ num_reqs = common_attn_metadata.num_reqs
694
+ num_tokens = common_attn_metadata.num_actual_tokens
695
+
696
+ decode_metadata = AiterFlashAttentionDecodeMetadata(
697
+ max_query_len=common_attn_metadata.max_query_len,
698
+ min_query_len=common_attn_metadata.max_query_len, # uniform batch
699
+ max_seq_len=common_attn_metadata.max_seq_len,
700
+ query_start_loc=common_attn_metadata.query_start_loc,
701
+ )
702
+
703
+ return AiterFlashAttentionMetadata(
704
+ num_actual_tokens=num_tokens,
705
+ num_actual_kv_tokens=0, # not used in unified_attention path
706
+ max_query_len=common_attn_metadata.max_query_len,
707
+ query_start_loc=common_attn_metadata.query_start_loc,
708
+ max_seq_len=common_attn_metadata.max_seq_len,
709
+ seq_lens=common_attn_metadata.seq_lens,
710
+ block_table=common_attn_metadata.block_table_tensor,
711
+ slot_mapping=common_attn_metadata.slot_mapping,
712
+ num_decodes=num_reqs,
713
+ num_decode_tokens=num_tokens,
714
+ num_prefills=0,
715
+ num_prefill_tokens=0,
716
+ num_extends=0,
717
+ num_extend_tokens=0,
718
+ decode_metadata=decode_metadata,
719
+ prefill_metadata=None,
720
+ extend_metadata=None,
721
+ use_cascade=False,
722
+ common_prefix_len=0,
723
+ total_tokens=self.total_tokens,
724
+ k_scale=self.scale,
725
+ v_scale=self.scale,
726
+ )
727
+
728
def use_cascade_attention(self, *args, **kwargs) -> bool:
729
return False
730