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, 16:00 UTC
[Bugfix] [ROCm] [UX] Reorganize ROCm Backend Selection Logic (#26980)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Nov 24, 2025 at 22:24 UTC
e48b2e6848ac0084860c8a6bd73927d7535a2c61
2 files changed
+394
-23
tests/v1/attention/test_rocm_attention_backends_selection.py
new
+337
@@ -0,0 +1,337 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for attention backend selectors."""
4
+
5
+from unittest.mock import MagicMock, patch
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.attention.backends.registry import AttentionBackendEnum
11
+from vllm.platforms import current_platform
12
+
13
+# ROCm-specific attention backend selection tests
14
+pytestmark = pytest.mark.skipif(
15
+ not current_platform.is_rocm(), reason="ROCm-specific tests"
16
+)
17
+
18
+
19
+@pytest.fixture
20
+def mock_vllm_config():
21
+ """Create a mock VllmConfig for testing."""
22
+ config = MagicMock()
23
+ config.model_config.dtype = torch.float16
24
+ config.model_config.hf_config.architectures = ["LlamaForCausalLM"]
25
+ config.cache_config.block_size = 16
26
+ return config
27
+
28
+
29
+@pytest.fixture
30
+def mock_on_gfx9():
31
+ """Mock the on_gfx9 function to return True."""
32
+ with patch("vllm.platforms.rocm.on_gfx9", return_value=True):
33
+ yield
34
+
35
+
36
+@pytest.mark.parametrize(
37
+ "env_vars, selected_backend, expected_backend_path",
38
+ [
39
+ # Test Case 1: Default (no env vars, no explicit backend)
40
+ (
41
+ {},
42
+ None,
43
+ AttentionBackendEnum.TRITON_ATTN.get_path(),
44
+ ),
45
+ # Test Case 2: Explicit TRITON_ATTN backend
46
+ (
47
+ {},
48
+ "TRITON_ATTN",
49
+ AttentionBackendEnum.TRITON_ATTN.get_path(),
50
+ ),
51
+ # Test Case 3: Explicit ROCM_ATTN backend
52
+ (
53
+ {},
54
+ "ROCM_ATTN",
55
+ AttentionBackendEnum.ROCM_ATTN.get_path(),
56
+ ),
57
+ # Test Case 4: Explicit ROCM_AITER_FA backend
58
+ (
59
+ {},
60
+ "ROCM_AITER_FA",
61
+ AttentionBackendEnum.ROCM_AITER_FA.get_path(),
62
+ ),
63
+ # Test Case 5: Explicit ROCM_AITER_UNIFIED_ATTN backend
64
+ (
65
+ {},
66
+ "ROCM_AITER_UNIFIED_ATTN",
67
+ AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
68
+ ),
69
+ # Test Case 6: VLLM_ROCM_USE_AITER=1
70
+ # (defaults to AITER FA when MHA not explicitly disabled)
71
+ (
72
+ {"VLLM_ROCM_USE_AITER": "1"},
73
+ None,
74
+ AttentionBackendEnum.ROCM_AITER_FA.get_path(),
75
+ ),
76
+ # Test Case 7: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=1
77
+ (
78
+ {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "1"},
79
+ None,
80
+ AttentionBackendEnum.ROCM_AITER_FA.get_path(),
81
+ ),
82
+ # Test Case 8: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1
83
+ (
84
+ {
85
+ "VLLM_ROCM_USE_AITER": "1",
86
+ "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": "1",
87
+ },
88
+ None,
89
+ AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
90
+ ),
91
+ # Test Case 9: VLLM_V1_USE_PREFILL_DECODE_ATTENTION=1
92
+ (
93
+ {"VLLM_V1_USE_PREFILL_DECODE_ATTENTION": "1"},
94
+ None,
95
+ AttentionBackendEnum.ROCM_ATTN.get_path(),
96
+ ),
97
+ # Test Case 10: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
98
+ (
99
+ {"VLLM_ROCM_USE_AITER": "1"},
100
+ "TRITON_ATTN",
101
+ AttentionBackendEnum.TRITON_ATTN.get_path(),
102
+ ),
103
+ # Test Case 11: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0
104
+ # (explicitly disabled)
105
+ (
106
+ {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"},
107
+ None,
108
+ AttentionBackendEnum.TRITON_ATTN.get_path(),
109
+ ),
110
+ # Test Case 12: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
111
+ (
112
+ {"VLLM_ROCM_USE_AITER": "1"},
113
+ "ROCM_ATTN",
114
+ AttentionBackendEnum.ROCM_ATTN.get_path(),
115
+ ),
116
+ ],
117
+)
118
+def test_standard_attention_backend_selection(
119
+ env_vars,
120
+ selected_backend,
121
+ expected_backend_path,
122
+ mock_vllm_config,
123
+ mock_on_gfx9,
124
+ monkeypatch,
125
+):
126
+ """Test standard attention backend selection with various configurations."""
127
+ # Set environment variables
128
+ for key, value in env_vars.items():
129
+ monkeypatch.setenv(key, value)
130
+
131
+ # Import after setting env vars to ensure they're picked up
132
+ # Reload envs to pick up new environment variables
133
+ import importlib
134
+
135
+ import vllm.envs as envs
136
+ from vllm.attention.backends.registry import _Backend
137
+
138
+ importlib.reload(envs)
139
+
140
+ # Convert string backend to enum if provided
141
+ backend_enum = None
142
+ if selected_backend:
143
+ backend_enum = getattr(_Backend, selected_backend)
144
+
145
+ # Get the backend class path
146
+ from vllm.platforms.rocm import RocmPlatform
147
+
148
+ backend_path = RocmPlatform.get_attn_backend_cls(
149
+ selected_backend=backend_enum,
150
+ head_size=128,
151
+ dtype=torch.float16,
152
+ kv_cache_dtype="auto",
153
+ block_size=16,
154
+ use_mla=False,
155
+ has_sink=False,
156
+ use_sparse=False,
157
+ )
158
+ assert backend_path == expected_backend_path
159
+
160
+
161
+@pytest.mark.parametrize(
162
+ "env_vars, selected_backend, block_size, expected_backend_path, should_raise",
163
+ [
164
+ # Test Case 1: TRITON_MLA with block_size != 1
165
+ (
166
+ {},
167
+ "TRITON_MLA",
168
+ 16,
169
+ AttentionBackendEnum.TRITON_MLA.get_path(),
170
+ False,
171
+ ),
172
+ # Test Case 2: TRITON_MLA with block_size == 1 (should raise)
173
+ (
174
+ {},
175
+ "TRITON_MLA",
176
+ 1,
177
+ None,
178
+ True,
179
+ ),
180
+ # Test Case 3: ROCM_AITER_MLA with block_size == 1
181
+ (
182
+ {},
183
+ "ROCM_AITER_MLA",
184
+ 1,
185
+ AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
186
+ False,
187
+ ),
188
+ # Test Case 4: ROCM_AITER_MLA with block_size != 1 (should raise)
189
+ (
190
+ {},
191
+ "ROCM_AITER_MLA",
192
+ 16,
193
+ AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
194
+ False,
195
+ ),
196
+ # Test Case 5: VLLM_ROCM_USE_AITER=1 with block_size == 1
197
+ (
198
+ {"VLLM_ROCM_USE_AITER": "1"},
199
+ None,
200
+ 1,
201
+ AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
202
+ False,
203
+ ),
204
+ # Test Case 6: VLLM_ROCM_USE_AITER=1 with block_size == 16
205
+ # (should use ROCM_AITER_MLA now, as it supports block_size 16)
206
+ (
207
+ {"VLLM_ROCM_USE_AITER": "1"},
208
+ None,
209
+ 16,
210
+ AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
211
+ False,
212
+ ),
213
+ # Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_MLA
214
+ (
215
+ {"VLLM_ROCM_USE_AITER": "1"},
216
+ "TRITON_MLA",
217
+ 16,
218
+ AttentionBackendEnum.TRITON_MLA.get_path(),
219
+ False,
220
+ ),
221
+ # Test Case 8: Explicit ROCM_AITER_TRITON_MLA
222
+ (
223
+ {},
224
+ "ROCM_AITER_TRITON_MLA",
225
+ 16,
226
+ AttentionBackendEnum.ROCM_AITER_TRITON_MLA.get_path(),
227
+ False,
228
+ ),
229
+ ],
230
+)
231
+def test_mla_backend_selection(
232
+ env_vars,
233
+ selected_backend,
234
+ block_size,
235
+ expected_backend_path,
236
+ should_raise,
237
+ mock_vllm_config,
238
+ monkeypatch,
239
+):
240
+ """Test MLA backend selection with various configurations."""
241
+ # Set environment variables
242
+ for key, value in env_vars.items():
243
+ monkeypatch.setenv(key, value)
244
+
245
+ # Import after setting env vars
246
+ # Reload envs
247
+ import importlib
248
+
249
+ import vllm.envs as envs
250
+ from vllm.attention.backends.registry import _Backend
251
+
252
+ importlib.reload(envs)
253
+
254
+ # Mock is_aiter_mla_enabled based on env vars and block_size
255
+ aiter_enabled = env_vars.get("VLLM_ROCM_USE_AITER") == "1"
256
+
257
+ mock_rocm_ops = MagicMock()
258
+ mock_rocm_ops.is_mla_enabled.return_value = aiter_enabled
259
+ mock_aiter_module = MagicMock()
260
+ mock_aiter_module.rocm_aiter_ops = mock_rocm_ops
261
+
262
+ with patch.dict("sys.modules", {"vllm._aiter_ops": mock_aiter_module}):
263
+ # Convert string backend to enum if provided
264
+ backend_enum = None
265
+ if selected_backend:
266
+ backend_enum = getattr(_Backend, selected_backend)
267
+
268
+ from vllm.platforms.rocm import RocmPlatform
269
+
270
+ if should_raise:
271
+ with pytest.raises(ValueError):
272
+ RocmPlatform.get_attn_backend_cls(
273
+ selected_backend=backend_enum,
274
+ head_size=128,
275
+ dtype=torch.float16,
276
+ kv_cache_dtype="auto",
277
+ block_size=block_size,
278
+ use_mla=True,
279
+ has_sink=False,
280
+ use_sparse=False,
281
+ )
282
+ else:
283
+ backend_path = RocmPlatform.get_attn_backend_cls(
284
+ selected_backend=backend_enum,
285
+ head_size=128,
286
+ dtype=torch.float16,
287
+ kv_cache_dtype="auto",
288
+ block_size=block_size,
289
+ use_mla=True,
290
+ has_sink=False,
291
+ use_sparse=False,
292
+ )
293
+ assert backend_path == expected_backend_path
294
+
295
+
296
+def test_aiter_fa_requires_gfx9(mock_vllm_config):
297
+ """Test that ROCM_AITER_FA requires gfx9 architecture."""
298
+ from vllm.attention.backends.registry import _Backend
299
+ from vllm.platforms.rocm import RocmPlatform
300
+
301
+ # Mock on_gfx9 to return False
302
+ with (
303
+ patch("vllm.platforms.rocm.on_gfx9", return_value=False),
304
+ pytest.raises(
305
+ ValueError,
306
+ match="only supported on gfx9",
307
+ ),
308
+ ):
309
+ RocmPlatform.get_attn_backend_cls(
310
+ selected_backend=_Backend.ROCM_AITER_FA,
311
+ head_size=128,
312
+ dtype=torch.float16,
313
+ kv_cache_dtype="auto",
314
+ block_size=16,
315
+ use_mla=False,
316
+ has_sink=False,
317
+ use_sparse=False,
318
+ )
319
+
320
+
321
+def test_sparse_not_supported(mock_vllm_config):
322
+ """Test that sparse attention is not supported on ROCm."""
323
+ from vllm.platforms.rocm import RocmPlatform
324
+
325
+ with pytest.raises(
326
+ AssertionError, match="Sparse MLA backend on ROCm only supports block size 1"
327
+ ):
328
+ RocmPlatform.get_attn_backend_cls(
329
+ selected_backend=None,
330
+ head_size=128,
331
+ dtype=torch.float16,
332
+ kv_cache_dtype="auto",
333
+ block_size=16,
334
+ use_mla=False,
335
+ has_sink=False,
336
+ use_sparse=True,
337
+ )
vllm/platforms/rocm.py
+57
-23
@@ -262,30 +262,64 @@ class RocmPlatform(Platform):
262
f"is not MLA type while requested for MLA backend."
263
)
264
265
- if selected_backend == AttentionBackendEnum.FLEX_ATTENTION:
266
- logger.info("Using FlexAttention backend.")
267
- return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
268
- if (
269
- rocm_aiter_ops.is_mha_enabled()
270
- ) or selected_backend == AttentionBackendEnum.ROCM_AITER_FA:
271
- logger.info("Using Aiter Flash Attention backend.")
272
- return AttentionBackendEnum.ROCM_AITER_FA.get_path()
273
- if (
274
- rocm_aiter_ops.is_triton_unified_attn_enabled()
275
- ) or selected_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
276
- logger.info("Using Aiter Unified Attention backend.")
277
- return AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
278
- if (
279
- envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
280
- or selected_backend == AttentionBackendEnum.ROCM_ATTN
281
- ):
282
- # rocm specific backend, with aiter and/or
283
- # triton prefix-prefill
284
- logger.info("Using Rocm Attention backend.")
265
+ if selected_backend == AttentionBackendEnum.TRITON_ATTN:
266
+ logger.info("Using Triton Attention backend on V1 engine.")
267
+ return AttentionBackendEnum.TRITON_ATTN.get_path()
268
+
269
+ if selected_backend == AttentionBackendEnum.ROCM_ATTN:
270
+ logger.info("Using Rocm Attention backend on V1 engine.")
271
return AttentionBackendEnum.ROCM_ATTN.get_path()
286
- # default case, using triton unified attention
287
- logger.info("Using Triton Attention backend.")
288
- return AttentionBackendEnum.TRITON_ATTN.get_path()
272
+
273
+ if selected_backend == AttentionBackendEnum.ROCM_AITER_FA:
274
+ if on_gfx9():
275
+ logger.info("Using Aiter Flash Attention backend on V1 engine.")
276
+ return AttentionBackendEnum.ROCM_AITER_FA.get_path()
277
+ else:
278
+ raise ValueError(
279
+ f"The selected backend, {selected_backend.name}, "
280
+ "is only supported on gfx9 architectures."
281
+ )
282
+
283
+ if selected_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
284
+ logger.info("Using Aiter Unified Attention backend on V1 engine.")
285
+ return AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
286
+
287
+ # Handle automatic backend selection based on environment variables
288
+ if selected_backend is None:
289
+ # Priority 1: Check for AITER Unified Attention (must check before MHA)
290
+ if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION:
291
+ logger.info("Using Aiter Unified Attention backend on V1 engine.")
292
+ return AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
293
+
294
+ # Priority 2: Check for AITER MHA (Flash Attention)
295
+ # Only use if explicitly enabled (not just VLLM_ROCM_USE_AITER=1)
296
+ if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
297
+ logger.info("Using Aiter Flash Attention backend on V1 engine.")
298
+ return AttentionBackendEnum.ROCM_AITER_FA.get_path()
299
+
300
+ # Priority 3: Check for ROCM_ATTN (prefill-decode split)
301
+ if envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION:
302
+ logger.info("Using Rocm Attention backend on V1 engine.")
303
+ return AttentionBackendEnum.ROCM_ATTN.get_path()
304
+
305
+ # Priority 4: Check for AITER enabled without specific flags
306
+ # This defaults to AITER FA only if MHA is not explicitly disabled
307
+ if (
308
+ envs.VLLM_ROCM_USE_AITER
309
+ and on_gfx9()
310
+ and envs.VLLM_ROCM_USE_AITER_MHA is not False
311
+ ):
312
+ logger.info("Using Aiter Flash Attention backend on V1 engine.")
313
+ return AttentionBackendEnum.ROCM_AITER_FA.get_path()
314
+
315
+ # Default: Triton Unified Attention
316
+ logger.info("Using Triton Attention backend on V1 engine.")
317
+ return AttentionBackendEnum.TRITON_ATTN.get_path()
318
+
319
+ raise RuntimeError(
320
+ "V0 attention backends have been removed. Set VLLM_USE_V1=1 "
321
+ "to select a supported backend."
322
+ )
323
324
@classmethod
325
def set_device(cls, device: torch.device) -> None: