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 10, 16:30 UTC
[ROCm][CI] Fix tests/compile unit tests (#28895)
Signed-off-by: charlifu <charlifu@amd.com> Signed-off-by: Micah Williamson <micah.williamson@amd.com> Signed-off-by: Charlie Fu <Charlie.Fu@amd.com> Co-authored-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Charlie Fu committed
Jan 6, 2026 at 12:50 UTC
c07163663d0a5ab6db1e4833c44305545f847c85
3 files changed
+30
-13
tests/compile/fullgraph/test_basic_correctness.py
+20
-11
@@ -5,10 +5,13 @@ import dataclasses
5
import pytest
6
7
from vllm.config import CompilationMode
8
+from vllm.platforms import current_platform
9
from vllm.utils.torch_utils import cuda_device_count_stateless
10
11
from ...utils import compare_all_settings
12
13
+ATTN_BACKEND = "FLASH_ATTN" if not current_platform.is_rocm() else "ROCM_ATTN"
14
+
15
16
@dataclasses.dataclass
17
class TestSetting:
@@ -31,7 +34,7 @@ class TestSetting:
34
model_args=["--max-model-len", "2048"],
35
pp_size=2,
36
tp_size=2,
34
- attn_backend="FLASH_ATTN",
37
+ attn_backend=ATTN_BACKEND,
38
method="generate",
39
),
40
# llama model with quantization
@@ -40,7 +43,7 @@ class TestSetting:
43
model_args=["--quantization", "gptq", "--max-model-len", "2048"],
44
pp_size=1,
45
tp_size=1,
43
- attn_backend="FLASH_ATTN",
46
+ attn_backend=ATTN_BACKEND,
47
method="generate",
48
),
49
# MoE model
@@ -49,7 +52,7 @@ class TestSetting:
52
model_args=["--max-model-len", "2048"],
53
pp_size=1,
54
tp_size=2,
52
- attn_backend="FLASH_ATTN",
55
+ attn_backend=ATTN_BACKEND,
56
method="generate",
57
),
58
# embedding model
@@ -65,16 +68,22 @@ class TestSetting:
68
],
69
pp_size=1,
70
tp_size=1,
68
- attn_backend="FLASH_ATTN",
71
+ attn_backend=ATTN_BACKEND,
72
method="encode",
73
),
71
- TestSetting(
72
- model="BAAI/bge-base-en-v1.5",
73
- model_args=["--runner", "pooling"],
74
- pp_size=1,
75
- tp_size=1,
76
- attn_backend="FLASH_ATTN",
77
- method="encode",
74
+ pytest.param(
75
+ TestSetting(
76
+ model="BAAI/bge-base-en-v1.5",
77
+ model_args=["--runner", "pooling"],
78
+ pp_size=1,
79
+ tp_size=1,
80
+ attn_backend="FLASH_ATTN",
81
+ method="encode",
82
+ ),
83
+ marks=pytest.mark.skipif(
84
+ current_platform.is_rocm(),
85
+ reason="Encoder self-attention is not implemented for ROCm",
86
+ ),
87
),
88
# vision language model
89
# See https://github.com/vllm-project/vllm/issues/26716.
tests/compile/fullgraph/test_full_cudagraph.py
+5
@@ -9,6 +9,7 @@ import pytest
9
from tests.utils import wait_for_gpu_memory_to_clear
10
from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
11
from vllm import LLM, SamplingParams
12
+from vllm.attention.backends.registry import AttentionBackendEnum
13
from vllm.config import CompilationConfig
14
from vllm.platforms import current_platform
15
from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -70,6 +71,10 @@ def llm_pair(request):
71
elif backend_config.specific_gpu_arch == (10, 0):
72
pytest.skip("Only Blackwell GPUs support Cutlass MLA")
73
74
+ # FlashInfer is not supported on ROCm
75
+ if backend_config == AttentionBackendEnum.FLASHINFER and current_platform.is_rocm():
76
+ pytest.skip("FlashInfer is not supported on ROCm")
77
+
78
env_vars = {
79
# Force native sampler to avoid potential nondeterminism in FlashInfer
80
# when per-request generators are not used in V1.
tests/compile/test_noop_elimination.py
+5
-2
@@ -25,10 +25,13 @@ def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
25
class Model(torch.nn.Module):
26
def __init__(self) -> None:
27
super().__init__()
28
- self.pos_embed = torch.empty(buffer_size, hidden_size, dtype=dtype)
28
+ # Avoid using empty, since on rocm torch.empty
29
+ # does not initialize the memory.
30
+ self.pos_embed = torch.randn(buffer_size, hidden_size, dtype=dtype)
31
32
def forward(self, x):
31
- x += self.pos_embed[: x.shape[0]]
33
+ # Avoid += to prevent inplace addition.
34
+ x = x + self.pos_embed[: x.shape[0]]
35
# Chain of reshapes
36
y = x.reshape(-1, 128, 32)
37
z = y.reshape(-1, 4096)