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 12, 14:45 UTC
[ROCm][CI] Fix test_cudagraph_mode failure in AMD CI (#29367)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Nov 25, 2025 at 01:55 UTC
ef1f7030f016cc811236517e02fa51ee8876cc31
3 files changed
+51
-22
tests/v1/attention/utils.py
+7
@@ -340,4 +340,11 @@ full_cg_backend_configs = {
340
"cudagraph_mode": "FULL_AND_PIECEWISE",
341
},
342
),
343
+ "RocmAttn": BackendConfig(
344
+ name="RocmAttn",
345
+ env_vars={"VLLM_V1_USE_PREFILL_DECODE_ATTENTION": "1"},
346
+ comp_config={
347
+ "cudagraph_mode": "FULL",
348
+ },
349
+ ),
350
}
tests/v1/cudagraph/test_cudagraph_mode.py
+42
-20
@@ -35,14 +35,22 @@ def temporary_environ(env_vars):
35
36
# test attention backend and cudagraph_mode combo
37
# (backend_name, cudagraph_mode, supported)
38
-combo_cases_1 = [
39
- ("FA3", "FULL", True),
40
- ("FA3", "FULL_AND_PIECEWISE", True),
41
- ("FA2", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
42
- ("FA2", "FULL_AND_PIECEWISE", True),
43
- ("FlashInfer", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
44
- ("FlashInfer", "FULL_AND_PIECEWISE", True),
45
-]
38
+if current_platform.is_rocm():
39
+ combo_cases_1 = [
40
+ ("RocmAttn", "FULL", True),
41
+ ("RocmAttn", "FULL_AND_PIECEWISE", True),
42
+ ("TritonAttn", "FULL", True),
43
+ ("TritonAttn", "FULL_AND_PIECEWISE", True),
44
+ ]
45
+else:
46
+ combo_cases_1 = [
47
+ ("FA3", "FULL", True),
48
+ ("FA3", "FULL_AND_PIECEWISE", True),
49
+ ("FA2", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
50
+ ("FA2", "FULL_AND_PIECEWISE", True),
51
+ ("FlashInfer", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
52
+ ("FlashInfer", "FULL_AND_PIECEWISE", True),
53
+ ]
54
55
56
@pytest.mark.parametrize("backend_name, cudagraph_mode, supported", combo_cases_1)
@@ -92,18 +100,32 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte
100
101
# test cudagraph_mode with different compilation mode.
102
# (backend_name, cudagraph_mode, compilation_mode, supported)
95
-combo_cases_2 = [
96
- ("FA2", "FULL", CompilationMode.NONE, True),
97
- ("FA2", "FULL", CompilationMode.VLLM_COMPILE, True),
98
- ("FA2", "PIECEWISE", CompilationMode.NONE, False),
99
- ("FA2", "PIECEWISE", CompilationMode.VLLM_COMPILE, True),
100
- ("FA2", "FULL_AND_PIECEWISE", CompilationMode.NONE, False),
101
- ("FA2", "FULL_AND_PIECEWISE", CompilationMode.VLLM_COMPILE, True),
102
- ("FA2", "FULL_DECODE_ONLY", CompilationMode.NONE, True),
103
- ("FA2", "FULL_DECODE_ONLY", CompilationMode.VLLM_COMPILE, True),
104
- ("FA2", "NONE", CompilationMode.NONE, True),
105
- ("FA2", "NONE", CompilationMode.VLLM_COMPILE, True),
106
-]
103
+if current_platform.is_rocm():
104
+ combo_cases_2 = [
105
+ ("RocmAttn", "FULL", CompilationMode.NONE, True),
106
+ ("RocmAttn", "FULL", CompilationMode.VLLM_COMPILE, True),
107
+ ("RocmAttn", "PIECEWISE", CompilationMode.NONE, False),
108
+ ("RocmAttn", "PIECEWISE", CompilationMode.VLLM_COMPILE, True),
109
+ ("RocmAttn", "FULL_AND_PIECEWISE", CompilationMode.NONE, False),
110
+ ("RocmAttn", "FULL_AND_PIECEWISE", CompilationMode.VLLM_COMPILE, True),
111
+ ("RocmAttn", "FULL_DECODE_ONLY", CompilationMode.NONE, True),
112
+ ("RocmAttn", "FULL_DECODE_ONLY", CompilationMode.VLLM_COMPILE, True),
113
+ ("RocmAttn", "NONE", CompilationMode.NONE, True),
114
+ ("RocmAttn", "NONE", CompilationMode.VLLM_COMPILE, True),
115
+ ]
116
+else:
117
+ combo_cases_2 = [
118
+ ("FA2", "FULL", CompilationMode.NONE, True),
119
+ ("FA2", "FULL", CompilationMode.VLLM_COMPILE, True),
120
+ ("FA2", "PIECEWISE", CompilationMode.NONE, False),
121
+ ("FA2", "PIECEWISE", CompilationMode.VLLM_COMPILE, True),
122
+ ("FA2", "FULL_AND_PIECEWISE", CompilationMode.NONE, False),
123
+ ("FA2", "FULL_AND_PIECEWISE", CompilationMode.VLLM_COMPILE, True),
124
+ ("FA2", "FULL_DECODE_ONLY", CompilationMode.NONE, True),
125
+ ("FA2", "FULL_DECODE_ONLY", CompilationMode.VLLM_COMPILE, True),
126
+ ("FA2", "NONE", CompilationMode.NONE, True),
127
+ ("FA2", "NONE", CompilationMode.VLLM_COMPILE, True),
128
+ ]
129
130
131
@pytest.mark.parametrize(
vllm/platforms/rocm.py
+2
-2
@@ -321,8 +321,8 @@ class RocmPlatform(Platform):
321
return AttentionBackendEnum.TRITON_ATTN.get_path()
322
323
raise RuntimeError(
324
- "V0 attention backends have been removed. Set VLLM_USE_V1=1 "
325
- "to select a supported backend."
324
+ f"Attention backend {selected_backend.name} is not supported on "
325
+ "ROCm. Note that V0 attention backends have been removed."
326
)
327
328
@classmethod