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 7, 00:15 UTC
[ROCm][CI] Fix ROCm attention backend validation for head sizes, block sizes, and compute capability checks (#36292)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 9, 2026 at 12:02 UTC
c174d54f86aa10e63ae236dc09f05f821134d469
7 files changed
+55
-9
docs/design/attention_backends.md
+1
-1
@@ -213,4 +213,4 @@ configuration.
213
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
215
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
216
-| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
216
+| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
tests/v1/attention/test_rocm_attention_backends_selection.py
+17
-8
@@ -29,11 +29,18 @@ def mock_vllm_config():
29
30
@pytest.fixture
31
def mock_on_gfx9():
32
- """Mock the on_gfx9 function to return True."""
32
+ """Mock gfx9 arch detection to return True."""
33
with patch("vllm.platforms.rocm.on_gfx9", return_value=True):
34
yield
35
36
37
+@pytest.fixture
38
+def mock_on_mi3xx():
39
+ """Mock mi3xx arch detection to return True."""
40
+ with patch("vllm.platforms.rocm.on_mi3xx", return_value=True):
41
+ yield
42
+
43
+
44
@pytest.mark.parametrize(
45
"env_vars, selected_backend, expected_backend_path",
46
[
@@ -122,6 +129,7 @@ def test_standard_attention_backend_selection(
129
expected_backend_path,
130
mock_vllm_config,
131
mock_on_gfx9,
132
+ mock_on_mi3xx,
133
monkeypatch,
134
):
135
"""Test standard attention backend selection with various configurations."""
@@ -313,16 +321,16 @@ def test_mla_backend_selection(
321
assert backend_path == expected_backend_path
322
323
316
-def test_aiter_fa_requires_gfx9(mock_vllm_config):
317
- """Test that ROCM_AITER_FA requires gfx9 architecture."""
324
+def test_aiter_fa_requires_mi3xx(mock_vllm_config):
325
+ """Test that ROCM_AITER_FA requires mi3xx architecture."""
326
from vllm.platforms.rocm import RocmPlatform
327
320
- # Mock on_gfx9 to return False
328
+ # Mock on_mi3xx to return False (used by supports_compute_capability)
329
with (
322
- patch("vllm.platforms.rocm.on_gfx9", return_value=False),
330
+ patch("vllm.platforms.rocm.on_mi3xx", return_value=False),
331
pytest.raises(
332
ValueError,
325
- match="only supported on gfx9",
333
+ match="compute capability not supported",
334
),
335
):
336
attn_selector_config = AttentionSelectorConfig(
@@ -342,11 +350,12 @@ def test_aiter_fa_requires_gfx9(mock_vllm_config):
350
351
352
def test_sparse_not_supported(mock_vllm_config):
345
- """Test that sparse attention is not supported on ROCm."""
353
+ """Test that sparse MLA without use_mla flag raises an error."""
354
from vllm.platforms.rocm import RocmPlatform
355
356
with pytest.raises(
349
- AssertionError, match="Sparse MLA backend on ROCm only supports block size 1"
357
+ ValueError,
358
+ match="No valid attention backend found",
359
):
360
attn_selector_config = AttentionSelectorConfig(
361
head_size=128,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+4
@@ -31,6 +31,10 @@ class AiterMLABackend(MLACommonBackend):
31
"fp8_e5m2",
32
]
33
34
+ @classmethod
35
+ def get_supported_head_sizes(cls) -> list[int]:
36
+ return []
37
+
38
@staticmethod
39
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
40
return [1]
vllm/v1/attention/backends/mla/triton_mla.py
+15
@@ -19,6 +19,7 @@ from vllm.platforms.interface import DeviceCapability
19
from vllm.v1.attention.backend import (
20
AttentionLayer,
21
AttentionType,
22
+ MultipleOf,
23
is_quantized_kv_cache,
24
)
25
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
@@ -33,6 +34,20 @@ class TritonMLABackend(MLACommonBackend):
34
"bfloat16",
35
]
36
37
+ @classmethod
38
+ def get_supported_head_sizes(cls) -> list[int]:
39
+ return []
40
+
41
+ @staticmethod
42
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
43
+ return [MultipleOf(16)]
44
+
45
+ @classmethod
46
+ def supports_block_size(cls, block_size: int | None) -> bool:
47
+ if block_size is None:
48
+ return True
49
+ return block_size % 16 == 0
50
+
51
@staticmethod
52
def get_name() -> str:
53
return "TRITON_MLA"
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+6
@@ -29,6 +29,12 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
29
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
30
return [MultipleOf(16)]
31
32
+ @classmethod
33
+ def supports_block_size(cls, block_size: int | None) -> bool:
34
+ if block_size is None:
35
+ return True
36
+ return block_size % 16 == 0
37
+
38
@classmethod
39
def supports_head_size(cls, head_size: int) -> bool:
40
return head_size >= 32
vllm/v1/attention/backends/rocm_attn.py
+6
@@ -188,6 +188,12 @@ class RocmAttentionBackend(AttentionBackend):
188
# uses our optimized kernel logic.
189
return [16, 32, 544]
190
191
+ @classmethod
192
+ def supports_block_size(cls, block_size: int | None) -> bool:
193
+ if block_size is None:
194
+ return True
195
+ return block_size in (16, 32, 544)
196
+
197
@classmethod
198
def get_supported_head_sizes(cls) -> list[int]:
199
return [32, 64, 80, 96, 128, 160, 192, 224, 256]
vllm/v1/attention/backends/triton_attn.py
+6
@@ -273,6 +273,12 @@ class TritonAttentionBackend(AttentionBackend):
273
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
274
return [MultipleOf(16)]
275
276
+ @classmethod
277
+ def supports_block_size(cls, block_size: int | None) -> bool:
278
+ if block_size is None:
279
+ return True
280
+ return block_size % 16 == 0
281
+
282
forward_includes_kv_cache_update: bool = False
283
284
@staticmethod