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 11, 20:00 UTC
[ROCm][CI] Fix test_max_len.py for Rocm (#29916)
Signed-off-by: charlifu <charlifu@amd.com> Signed-off-by: Charlie Fu <Charlie.Fu@amd.com>
Charlie Fu committed
Dec 8, 2025 at 15:58 UTC
6af70e11a0a3cb7109ce904e23ff90c73f573ef0
5 files changed
+15
-8
tests/basic_correctness/test_basic_correctness.py
+4
-1
@@ -13,12 +13,15 @@ import pytest
13
import torch
14
15
from vllm import LLM
16
+from vllm.platforms import current_platform
17
from vllm.v1.engine.llm_engine import LLMEngine
18
19
from ..conftest import HfRunner, VllmRunner
20
from ..models.utils import check_outputs_equal
21
from ..utils import multi_gpu_test
22
23
+ATTN_BACKEND = ["ROCM_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"]
24
+
25
MODELS = [
26
"hmellor/tiny-random-Gemma2ForCausalLM",
27
"meta-llama/Llama-3.2-1B-Instruct",
@@ -57,7 +60,7 @@ def _fix_prompt_embed_outputs(
60
61
62
@pytest.mark.parametrize("model", MODELS)
60
-@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
63
+@pytest.mark.parametrize("backend", ATTN_BACKEND)
64
@pytest.mark.parametrize("max_tokens", [5])
65
@pytest.mark.parametrize("enforce_eager", [False])
66
@pytest.mark.parametrize("async_scheduling", [True, False])
tests/utils.py
+2
-2
@@ -1225,9 +1225,9 @@ def get_attn_backend_list_based_on_platform() -> list[str]:
1225
try:
1226
import aiter # noqa: F401
1227
1228
- attn_backend_list.append("FLASH_ATTN")
1228
+ attn_backend_list.append("ROCM_AITER_FA")
1229
except Exception:
1230
- print("Skip FLASH_ATTN on ROCm as aiter is not installed")
1230
+ print("Skip ROCM_AITER_FA on ROCm as aiter is not installed")
1231
1232
return attn_backend_list
1233
elif current_platform.is_xpu():
tests/v1/e2e/test_spec_decode.py
+2
-2
@@ -417,9 +417,9 @@ def test_eagle_correctness(
417
"multi-token eagle spec decode on current platform"
418
)
419
420
- if attn_backend == "FLASH_ATTN" and current_platform.is_rocm():
420
+ if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
421
if "deepseek" in model_setup[1].lower():
422
- pytest.skip("FLASH_ATTN for deepseek not supported on ROCm platform")
422
+ pytest.skip("ROCM_AITER_FA for deepseek not supported on ROCm platform")
423
else:
424
m.setenv("VLLM_ROCM_USE_AITER", "1")
425
tests/v1/spec_decode/test_eagle.py
+6
-2
@@ -339,7 +339,7 @@ def test_load_model(
339
"multi-token eagle spec decode on current platform"
340
)
341
342
- if attn_backend == "FLASH_ATTN" and current_platform.is_rocm():
342
+ if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
343
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
344
345
# Setup draft model mock
@@ -434,7 +434,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
434
"because it requires special input mocking."
435
)
436
437
- if attn_backend == "FLASH_ATTN" and current_platform.is_rocm():
437
+ if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
438
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
439
440
# Use GPU device
@@ -541,6 +541,10 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
541
attn_metadata_builder_cls, _ = try_get_attention_backend(
542
AttentionBackendEnum.TREE_ATTN
543
)
544
+ elif attn_backend == "ROCM_AITER_FA":
545
+ attn_metadata_builder_cls, _ = try_get_attention_backend(
546
+ AttentionBackendEnum.ROCM_AITER_FA
547
+ )
548
else:
549
raise ValueError(f"Unsupported attention backend: {attn_backend}")
550
tests/v1/spec_decode/test_max_len.py
+1
-1
@@ -47,7 +47,7 @@ def test_eagle_max_len(
47
"multi-token eagle spec decode on current platform"
48
)
49
50
- if attn_backend == "FLASH_ATTN" and current_platform.is_rocm():
50
+ if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
51
m.setenv("VLLM_ROCM_USE_AITER", "1")
52
53
llm = LLM(