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
[ROCm] [V1] [SpecDec] Enable Speculative Decoding on ROCm V1 Engine (#21496)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Aug 7, 2025 at 19:13 UTC
1ee5ead5f8f1c3c77b73effcb230ee02952fbe1f
6 files changed
+128
-41
tests/utils.py
+16
@@ -986,3 +986,19 @@ def has_module_attribute(module_name, attribute_name):
986
return hasattr(module, attribute_name)
987
except ImportError:
988
return False
989
+
990
+
991
+def get_attn_backend_list_based_on_platform() -> list[str]:
992
+ if current_platform.is_cuda():
993
+ return ["FLASH_ATTN_VLLM_V1", "TRITON_ATTN_VLLM_V1", "TREE_ATTN"]
994
+ elif current_platform.is_rocm():
995
+ attn_backend_list = ["TRITON_ATTN_VLLM_V1"]
996
+ try:
997
+ import aiter # noqa: F401
998
+ attn_backend_list.append("FLASH_ATTN_VLLM_V1")
999
+ except Exception:
1000
+ print("Skip FLASH_ATTN_VLLM_V1 on ROCm as aiter is not installed")
1001
+
1002
+ return attn_backend_list
1003
+ else:
1004
+ raise ValueError("Unsupported platform")
tests/v1/attention/utils.py
+5
-2
@@ -11,7 +11,7 @@ import torch
11
from vllm.config import (CacheConfig, CompilationConfig, DeviceConfig,
12
LoadConfig, ModelConfig, ModelDType, ParallelConfig,
13
SchedulerConfig, VllmConfig)
14
-from vllm.platforms import _Backend
14
+from vllm.platforms import _Backend, current_platform
15
from vllm.utils import resolve_obj_by_qualname
16
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
17
from vllm.v1.kv_cache_interface import FullAttentionSpec
@@ -119,7 +119,10 @@ def get_attention_backend(backend_name: _Backend):
119
"""
120
backend_map = {
121
_Backend.FLASH_ATTN_VLLM_V1:
122
- "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend",
122
+ ("vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
123
+ if current_platform.is_cuda() else
124
+ "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
125
+ ),
126
_Backend.FLASHINFER_VLLM_V1:
127
"vllm.v1.attention.backends.flashinfer.FlashInferBackend",
128
_Backend.FLEX_ATTENTION:
tests/v1/e2e/test_spec_decode.py
+15
@@ -8,10 +8,12 @@ from typing import Any, Union
8
import pytest
9
import torch
10
11
+from tests.utils import get_attn_backend_list_based_on_platform
12
from vllm import LLM, SamplingParams
13
from vllm.assets.base import VLLM_S3_BUCKET_URL
14
from vllm.assets.image import VLM_IMAGES_DIR
15
from vllm.distributed import cleanup_dist_env_and_memory
16
+from vllm.platforms import current_platform
17
18
19
def get_test_prompts(mm_enabled: bool):
@@ -141,11 +143,14 @@ def test_ngram_correctness(
143
marks=pytest.mark.skip(reason="Skipping due to CI OOM issues")),
144
],
145
ids=["llama3_eagle", "llama3_eagle3", "llama4_eagle", "llama4_eagle_mm"])
146
+@pytest.mark.parametrize("attn_backend",
147
+ get_attn_backend_list_based_on_platform())
148
def test_eagle_correctness(
149
monkeypatch: pytest.MonkeyPatch,
150
sampling_config: SamplingParams,
151
model_setup: tuple[str, str, str, int],
152
mm_enabled: bool,
153
+ attn_backend: str,
154
):
155
# Generate test prompts inside the function instead of using fixture
156
test_prompts = get_test_prompts(mm_enabled)
@@ -156,6 +161,16 @@ def test_eagle_correctness(
161
'''
162
with monkeypatch.context() as m:
163
m.setenv("VLLM_USE_V1", "1")
164
+ m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
165
+
166
+ if (attn_backend == "TRITON_ATTN_VLLM_V1"
167
+ and not current_platform.is_rocm()):
168
+ pytest.skip("TRITON_ATTN_VLLM_V1 does not support "
169
+ "multi-token eagle spec decode on current platform")
170
+
171
+ if attn_backend == "FLASH_ATTN_VLLM_V1" and current_platform.is_rocm():
172
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
173
+
174
method, model_name, spec_model_name, tp_size = model_setup
175
176
ref_llm = LLM(model=model_name,
tests/v1/spec_decode/test_eagle.py
+45
-10
@@ -6,6 +6,7 @@ from unittest import mock
6
import pytest
7
import torch
8
9
+from tests.utils import get_attn_backend_list_based_on_platform
10
from tests.v1.attention.utils import (BatchSpec, _Backend,
11
create_common_attn_metadata,
12
create_standard_kv_cache_spec,
@@ -120,17 +121,28 @@ def test_prepare_inputs():
121
assert torch.equal(token_indices, expected_token_indices)
122
123
123
-@pytest.mark.parametrize("method,proposer_helper", [
124
- ("eagle", lambda k: _create_proposer("eagle", k)),
125
- ("eagle3", lambda k: _create_proposer("eagle3", k)),
126
-])
124
+@pytest.mark.parametrize("method", ["eagle", "eagle3"])
125
+@pytest.mark.parametrize("attn_backend",
126
+ get_attn_backend_list_based_on_platform())
127
@pytest.mark.parametrize("pp_size", [1, 2])
128
@pytest.mark.parametrize("use_distinct_embed_tokens", [True, False])
129
@mock.patch('vllm.v1.spec_decode.eagle.get_pp_group')
130
@mock.patch('vllm.v1.spec_decode.eagle.get_layers_from_vllm_config')
131
@mock.patch('vllm.v1.spec_decode.eagle.get_model')
132
def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
133
- proposer_helper, pp_size, use_distinct_embed_tokens):
133
+ attn_backend, pp_size, use_distinct_embed_tokens,
134
+ monkeypatch):
135
+
136
+ monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
137
+
138
+ if (attn_backend == "TRITON_ATTN_VLLM_V1"
139
+ and not current_platform.is_rocm()):
140
+ pytest.skip("TRITON_ATTN_VLLM_V1 does not support "
141
+ "multi-token eagle spec decode on current platform")
142
+
143
+ if attn_backend == "FLASH_ATTN_VLLM_V1" and current_platform.is_rocm():
144
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
145
+
146
# Setup draft model mock
147
mock_model = mock.MagicMock()
148
if use_distinct_embed_tokens:
@@ -177,7 +189,7 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
189
target_model.lm_head = mock.MagicMock()
190
191
# Create proposer using the helper function
180
- proposer = proposer_helper(k=8)
192
+ proposer = _create_proposer(method, k=8)
193
194
# Call the method under test
195
proposer.load_model(target_model)
@@ -201,10 +213,22 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
213
target_model.model.embed_tokens
214
215
216
+@pytest.mark.parametrize("method", ["eagle", "eagle3"])
217
+@pytest.mark.parametrize("attn_backend",
218
+ get_attn_backend_list_based_on_platform())
219
@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 8])
205
-@pytest.mark.parametrize("backend",
206
- [_Backend.FLASH_ATTN_VLLM_V1, _Backend.TREE_ATTN])
207
-def test_propose(num_speculative_tokens, backend):
220
+def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
221
+
222
+ monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
223
+
224
+ if (attn_backend == "TRITON_ATTN_VLLM_V1"
225
+ and not current_platform.is_rocm()):
226
+ pytest.skip("TRITON_ATTN_VLLM_V1 does not support "
227
+ "multi-token eagle spec decode on current platform")
228
+
229
+ if attn_backend == "FLASH_ATTN_VLLM_V1" and current_platform.is_rocm():
230
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
231
+
232
# Use GPU device
233
device = torch.device(current_platform.device_type)
234
@@ -303,7 +327,18 @@ def test_propose(num_speculative_tokens, backend):
327
device=device)
328
sampling_metadata = mock.MagicMock()
329
306
- attn_metadata_builder_cls, _ = get_attention_backend(backend)
330
+ if attn_backend == "FLASH_ATTN_VLLM_V1":
331
+ attn_metadata_builder_cls, _ = get_attention_backend(
332
+ _Backend.FLASH_ATTN_VLLM_V1)
333
+ elif attn_backend == "TRITON_ATTN_VLLM_V1":
334
+ attn_metadata_builder_cls, _ = get_attention_backend(
335
+ _Backend.TRITON_ATTN_VLLM_V1)
336
+ elif attn_backend == "TREE_ATTN":
337
+ attn_metadata_builder_cls, _ = get_attention_backend(
338
+ _Backend.TREE_ATTN)
339
+ else:
340
+ raise ValueError(f"Unsupported attention backend: {attn_backend}")
341
+
342
attn_metadata_builder = attn_metadata_builder_cls(
343
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
344
layer_names=proposer.attn_layer_names,
tests/v1/spec_decode/test_max_len.py
+30
-24
@@ -4,7 +4,9 @@
4
5
import pytest
6
7
+from tests.utils import get_attn_backend_list_based_on_platform
8
from vllm import LLM, SamplingParams
9
+from vllm.platforms import current_platform
10
11
_PROMPTS = [
12
"1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1",
@@ -14,36 +16,40 @@ _PROMPTS = [
16
17
18
@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
17
-def test_ngram_max_len(
18
- monkeypatch: pytest.MonkeyPatch,
19
- num_speculative_tokens: int,
20
-):
21
- with monkeypatch.context() as m:
22
- m.setenv("VLLM_USE_V1", "1")
23
-
24
- llm = LLM(
25
- model="facebook/opt-125m",
26
- max_model_len=100,
27
- enforce_eager=True, # For faster initialization.
28
- speculative_config={
29
- "method": "ngram",
30
- "prompt_lookup_max": 5,
31
- "prompt_lookup_min": 3,
32
- "num_speculative_tokens": num_speculative_tokens,
33
- },
34
- )
35
- sampling_params = SamplingParams(max_tokens=100, ignore_eos=True)
36
- llm.generate(_PROMPTS, sampling_params)
19
+def test_ngram_max_len(num_speculative_tokens: int):
20
+ llm = LLM(
21
+ model="facebook/opt-125m",
22
+ max_model_len=100,
23
+ enforce_eager=True, # For faster initialization.
24
+ speculative_config={
25
+ "method": "ngram",
26
+ "prompt_lookup_max": 5,
27
+ "prompt_lookup_min": 3,
28
+ "num_speculative_tokens": num_speculative_tokens,
29
+ },
30
+ )
31
+ sampling_params = SamplingParams(max_tokens=100, ignore_eos=True)
32
+ llm.generate(_PROMPTS, sampling_params)
33
34
35
@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
40
-def test_eagle_max_len(
41
- monkeypatch: pytest.MonkeyPatch,
42
- num_speculative_tokens: int,
43
-):
36
+@pytest.mark.parametrize("attn_backend",
37
+ get_attn_backend_list_based_on_platform())
38
+def test_eagle_max_len(monkeypatch: pytest.MonkeyPatch,
39
+ num_speculative_tokens: int, attn_backend: str):
40
with monkeypatch.context() as m:
41
m.setenv("VLLM_USE_V1", "1")
42
43
+ m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
44
+
45
+ if (attn_backend == "TRITON_ATTN_VLLM_V1"
46
+ and not current_platform.is_rocm()):
47
+ pytest.skip("TRITON_ATTN_VLLM_V1 does not support "
48
+ "multi-token eagle spec decode on current platform")
49
+
50
+ if attn_backend == "FLASH_ATTN_VLLM_V1" and current_platform.is_rocm():
51
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
52
+
53
llm = LLM(
54
model="meta-llama/Meta-Llama-3-8B-Instruct",
55
enforce_eager=True, # For faster initialization.
vllm/v1/spec_decode/eagle.py
+17
-5
@@ -17,10 +17,14 @@ from vllm.logger import init_logger
17
from vllm.model_executor.model_loader import get_model
18
from vllm.model_executor.models import supports_multimodal
19
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
20
+from vllm.platforms import current_platform
21
from vllm.utils import is_pin_memory_available
22
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
23
+from vllm.v1.attention.backends.rocm_aiter_fa import (
24
+ AiterFlashAttentionMetadata)
25
from vllm.v1.attention.backends.tree_attn import (TreeAttentionMetadata,
26
TreeAttentionMetadataBuilder)
27
+from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
28
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
29
from vllm.v1.kv_cache_interface import KVCacheConfig
30
from vllm.v1.sample.metadata import SamplingMetadata
@@ -230,11 +234,19 @@ class EagleProposer:
234
# one layer. Adapt this code to support multiple layers once
235
# there's a multi-layer MTP module.
236
233
- # Currently, only FlashAttention and TreeAttention support multi-token
234
- # eagle spec decode. This is because the code below
235
- # makes assumptions about attn_metadata attributes available.
236
- assert isinstance(attn_metadata,
237
- (FlashAttentionMetadata, TreeAttentionMetadata))
237
+ # On ROCm, both AiterFlashAttention and TritonAttention
238
+ # support multi-token eagle spec decode.
239
+ if current_platform.is_rocm():
240
+ assert isinstance(
241
+ attn_metadata,
242
+ (TritonAttentionMetadata, AiterFlashAttentionMetadata,
243
+ FlashAttentionMetadata))
244
+ else:
245
+ # Currently, only FlashAttention and TreeAttention support
246
+ # multi-token eagle spec decode. This is because the code below
247
+ # makes assumptions about attn_metadata attributes available.
248
+ assert isinstance(attn_metadata,
249
+ (FlashAttentionMetadata, TreeAttentionMetadata))
250
251
# Generate the remaining draft tokens.
252
draft_token_ids_list = [draft_token_ids]