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 14, 05:30 UTC
[FIXBUG ] Allow disabling rocm_aiter_fa backend for ROCm GPUs not compatible with AITER (#22795)
Signed-off-by: JartX <sagformas@epdcenter.es> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
JartX committed
Aug 20, 2025 at 18:08 UTC
3b11b26b5069718a6bde11b9041681bc17369f96
1 file changed
+45
-35
vllm/v1/spec_decode/eagle.py
+45
-35
@@ -2,7 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import ast
4
from dataclasses import replace
5
-from typing import Optional
5
+from importlib.util import find_spec
6
+from typing import Optional, Protocol
7
8
import numpy as np
9
import torch
@@ -20,8 +21,6 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
21
from vllm.platforms import current_platform
22
from vllm.utils import is_pin_memory_available
23
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
23
-from vllm.v1.attention.backends.rocm_aiter_fa import (
24
- AiterFlashAttentionMetadata)
24
from vllm.v1.attention.backends.tree_attn import (TreeAttentionMetadata,
25
TreeAttentionMetadataBuilder)
26
from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
@@ -34,6 +33,17 @@ logger = init_logger(__name__)
33
PADDING_SLOT_ID = -1
34
35
36
+class EagleAttentionMetadata(Protocol):
37
+ # Required attributes
38
+ num_actual_tokens: int
39
+ max_query_len: int
40
+ query_start_loc: torch.Tensor
41
+ max_seq_len: int
42
+ seq_lens: torch.Tensor
43
+ block_table: torch.Tensor
44
+ slot_mapping: torch.Tensor
45
+
46
+
47
class EagleProposer:
48
49
def __init__(
@@ -97,6 +107,20 @@ class EagleProposer:
107
dtype=self.dtype,
108
device=device)
109
110
+ # Determine allowed attention backends once during initialization.
111
+ self.allowed_attn_types: tuple[type[EagleAttentionMetadata], ...]
112
+ if current_platform.is_rocm():
113
+ rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]
114
+ # vllm.v1.attention.backends.rocm_aiter_fa is an optional backend
115
+ if find_spec("vllm.v1.attention.backends.rocm_aiter_fa"):
116
+ from vllm.v1.attention.backends.rocm_aiter_fa import (
117
+ AiterFlashAttentionMetadata)
118
+ rocm_types.append(AiterFlashAttentionMetadata)
119
+ self.allowed_attn_types = tuple(rocm_types)
120
+ else:
121
+ self.allowed_attn_types = (FlashAttentionMetadata,
122
+ TreeAttentionMetadata)
123
+
124
# Parse the speculative token tree.
125
spec_token_tree = self.speculative_config.speculative_token_tree
126
self.tree_choices: list[tuple[int,
@@ -165,7 +189,7 @@ class EagleProposer:
189
for layer_name in self.attn_layer_names:
190
per_layer_attn_metadata[layer_name] = attn_metadata
191
if self.use_cuda_graph and \
168
- num_tokens <= self.cudagraph_batch_sizes[-1]:
192
+ num_tokens <= self.cudagraph_batch_sizes[-1]:
193
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
194
else:
195
num_input_tokens = num_tokens
@@ -225,25 +249,13 @@ class EagleProposer:
249
# TODO: Currently, MTP module released by deepseek only has
250
# one layer. Adapt this code to support multiple layers once
251
# there's a multi-layer MTP module.
228
-
229
- # On ROCm, both AiterFlashAttention and TritonAttention
230
- # support multi-token eagle spec decode.
231
- if current_platform.is_rocm():
232
- assert isinstance(
233
- attn_metadata,
234
- (TritonAttentionMetadata, AiterFlashAttentionMetadata,
235
- FlashAttentionMetadata))
236
- else:
237
- # Currently, only FlashAttention supports multi-token eagle spec
238
- # decode. This is because the code below makes assumptions about
239
- # attn_metadata attributes available.
240
- assert isinstance(attn_metadata, FlashAttentionMetadata)
252
+ assert isinstance(attn_metadata, self.allowed_attn_types)
253
254
# Generate the remaining draft tokens.
255
draft_token_ids_list = [draft_token_ids]
256
257
if self.use_cuda_graph and \
246
- batch_size <= self.cudagraph_batch_sizes[-1]:
258
+ batch_size <= self.cudagraph_batch_sizes[-1]:
259
input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
260
else:
261
input_batch_size = batch_size
@@ -449,7 +461,7 @@ class EagleProposer:
461
num_tokens, -1)
462
463
if self.use_cuda_graph and \
452
- num_tokens <= self.cudagraph_batch_sizes[-1]:
464
+ num_tokens <= self.cudagraph_batch_sizes[-1]:
465
num_input_tokens = self.vllm_config.pad_for_cudagraph(
466
num_tokens)
467
else:
@@ -508,19 +520,19 @@ class EagleProposer:
520
"""
521
# E.g.
522
# common_attn_metadata.query_start_loc{_cpu}:
511
- # [0, q1, q1 + q2, q1 + q2 + q3]
523
+ # [0, q1, q1 + q2, q1 + q2 + q3]
524
# common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3]
525
# num_rejected_tokens: [n1, n2, n3]
526
# This function computes the intermediate values:
527
# num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3]
528
# And returns:
529
# common_attn_metadata.query_start_loc{_cpu}:
518
- # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
530
+ # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
531
# common_attn_metadata.seq_lens{_cpu}:
520
- # [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]
532
+ # [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]
533
# token_indices: [0, 1, ..., q1 - n1 - 1,
522
- # q1, q1 + 1, ..., q1 + q2 - n2 - 1,
523
- # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
534
+ # q1, q1 + 1, ..., q1 + q2 - n2 - 1,
535
+ # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
536
537
device = common_attn_metadata.query_start_loc.device
538
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
@@ -564,9 +576,9 @@ class EagleProposer:
576
old_query_start_locs_expanded = np.repeat(
577
query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np)
578
# Final token indices are:
567
- # [0, 1, // req 1
568
- # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2
569
- # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3
579
+ # [0, 1, // req 1
580
+ # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2
581
+ # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3
582
token_indices_np = token_offests + old_query_start_locs_expanded
583
token_indices = torch.from_numpy(token_indices_np).to(
584
device, non_blocking=True)
@@ -616,20 +628,18 @@ class EagleProposer:
628
target_language_model = target_model
629
# share embed_tokens with the target model if needed
630
if get_pp_group().world_size == 1 \
619
- and self.model.model.embed_tokens.weight.shape \
620
- == target_language_model.model.embed_tokens.weight.shape:
631
+ and self.model.model.embed_tokens.weight.shape \
632
+ == target_language_model.model.embed_tokens.weight.shape:
633
logger.info(
622
- "Assuming the EAGLE head shares the same vocab embedding" \
623
- " with the target model."
624
- )
634
+ "Assuming the EAGLE head shares the same vocab embedding"
635
+ " with the target model.")
636
del self.model.model.embed_tokens
637
self.model.model.embed_tokens = (
638
target_language_model.model.embed_tokens)
639
else:
640
logger.info(
630
- "The EAGLE head's vocab embedding will be loaded separately" \
631
- " from the target model."
632
- )
641
+ "The EAGLE head's vocab embedding will be loaded separately"
642
+ " from the target model.")
643
644
# share lm_head with the target model if needed
645
# some model definition do not define lm_head explicitly