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, 20:15 UTC
[CI/ROCm] Fixing "V1 Test attention (H100)" test group. (#31187)
Signed-off-by: DCCS-4560 <alivanov@chi-mi325x-pod1-108.ord.vultr.cpe.ice.amd.com> Signed-off-by: <> Co-authored-by: DCCS-4560 <alivanov@chi-mi325x-pod1-108.ord.vultr.cpe.ice.amd.com> Co-authored-by: root <root@chi-mi325x-pod1-108.ord.vultr.cpe.ice.amd.com>
Alexei-V-Ivanov-AMD committed
Dec 29, 2025 at 15:53 UTC
d63b969675f1e655ad318b629d9eb1516705a8c4
3 files changed
+72
-19
tests/v1/attention/test_attention_backends.py
+29
-9
@@ -557,9 +557,21 @@ def test_causal_backend_correctness(
557
if is_torch_equal_or_newer("2.9.0.dev0")
558
else []
559
)
560
- SMALL_BLOCK_BACKENDS = [
561
- x for x in BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
562
- ]
560
+
561
+ if current_platform.is_rocm():
562
+ SMALL_BLOCK_BACKENDS = [
563
+ x
564
+ for x in BACKENDS_TO_TEST
565
+ if (
566
+ x not in LARGE_BLOCK_BACKENDS
567
+ and x is not AttentionBackendEnum.FLASH_ATTN
568
+ )
569
+ ]
570
+ else:
571
+ SMALL_BLOCK_BACKENDS = [
572
+ x for x in BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
573
+ ]
574
+
575
_test_backend_correctness(
576
batch_spec,
577
model,
@@ -580,12 +592,20 @@ def test_causal_backend_correctness(
592
)
593
594
583
-SLIDING_WINDOW_BACKENDS_TO_TEST = [
584
- AttentionBackendEnum.FLASH_ATTN,
585
- AttentionBackendEnum.FLEX_ATTENTION,
586
- AttentionBackendEnum.TRITON_ATTN,
587
- "FLEX_ATTENTION_SLOW",
588
-]
595
+if current_platform.is_rocm():
596
+ # FLASH_ATTN is not supported on ROCm
597
+ SLIDING_WINDOW_BACKENDS_TO_TEST = [
598
+ AttentionBackendEnum.FLEX_ATTENTION,
599
+ AttentionBackendEnum.TRITON_ATTN,
600
+ "FLEX_ATTENTION_SLOW",
601
+ ]
602
+else:
603
+ SLIDING_WINDOW_BACKENDS_TO_TEST = [
604
+ AttentionBackendEnum.FLASH_ATTN,
605
+ AttentionBackendEnum.FLEX_ATTENTION,
606
+ AttentionBackendEnum.TRITON_ATTN,
607
+ "FLEX_ATTENTION_SLOW",
608
+ ]
609
610
611
@pytest.mark.parametrize(
tests/v1/attention/test_rocm_attention_backends_selection.py
+39
-10
@@ -8,6 +8,7 @@ import pytest
8
import torch
9
10
from vllm.attention.backends.registry import AttentionBackendEnum
11
+from vllm.attention.selector import AttentionSelectorConfig
12
from vllm.platforms import current_platform
13
14
# ROCm-specific attention backend selection tests
@@ -144,8 +145,7 @@ def test_standard_attention_backend_selection(
145
# Get the backend class path
146
from vllm.platforms.rocm import RocmPlatform
147
147
- backend_path = RocmPlatform.get_attn_backend_cls(
148
- selected_backend=backend_enum,
148
+ attn_selector_config = AttentionSelectorConfig(
149
head_size=128,
150
dtype=torch.float16,
151
kv_cache_dtype="auto",
@@ -154,6 +154,11 @@ def test_standard_attention_backend_selection(
154
has_sink=False,
155
use_sparse=False,
156
)
157
+
158
+ backend_path = RocmPlatform.get_attn_backend_cls(
159
+ selected_backend=backend_enum, attn_selector_config=attn_selector_config
160
+ )
161
+
162
assert backend_path == expected_backend_path
163
164
@@ -267,8 +272,16 @@ def test_mla_backend_selection(
272
273
if should_raise:
274
with pytest.raises(ValueError):
270
- RocmPlatform.get_attn_backend_cls(
271
- selected_backend=backend_enum,
275
+ attn_selector_config = AttentionSelectorConfig(
276
+ head_size=128,
277
+ dtype=torch.float16,
278
+ kv_cache_dtype="auto",
279
+ block_size=block_size,
280
+ use_mla=True,
281
+ has_sink=False,
282
+ use_sparse=False,
283
+ )
284
+ attn_selector_config = AttentionSelectorConfig(
285
head_size=128,
286
dtype=torch.float16,
287
kv_cache_dtype="auto",
@@ -277,9 +290,13 @@ def test_mla_backend_selection(
290
has_sink=False,
291
use_sparse=False,
292
)
293
+ backend_path = RocmPlatform.get_attn_backend_cls(
294
+ selected_backend=backend_enum,
295
+ attn_selector_config=attn_selector_config,
296
+ )
297
+
298
else:
281
- backend_path = RocmPlatform.get_attn_backend_cls(
282
- selected_backend=backend_enum,
299
+ attn_selector_config = AttentionSelectorConfig(
300
head_size=128,
301
dtype=torch.float16,
302
kv_cache_dtype="auto",
@@ -288,6 +305,11 @@ def test_mla_backend_selection(
305
has_sink=False,
306
use_sparse=False,
307
)
308
+
309
+ backend_path = RocmPlatform.get_attn_backend_cls(
310
+ selected_backend=backend_enum, attn_selector_config=attn_selector_config
311
+ )
312
+
313
assert backend_path == expected_backend_path
314
315
@@ -303,8 +325,7 @@ def test_aiter_fa_requires_gfx9(mock_vllm_config):
325
match="only supported on gfx9",
326
),
327
):
306
- RocmPlatform.get_attn_backend_cls(
307
- selected_backend=AttentionBackendEnum.ROCM_AITER_FA,
328
+ attn_selector_config = AttentionSelectorConfig(
329
head_size=128,
330
dtype=torch.float16,
331
kv_cache_dtype="auto",
@@ -314,6 +335,11 @@ def test_aiter_fa_requires_gfx9(mock_vllm_config):
335
use_sparse=False,
336
)
337
338
+ RocmPlatform.get_attn_backend_cls(
339
+ selected_backend=AttentionBackendEnum.ROCM_AITER_FA,
340
+ attn_selector_config=attn_selector_config,
341
+ )
342
+
343
344
def test_sparse_not_supported(mock_vllm_config):
345
"""Test that sparse attention is not supported on ROCm."""
@@ -322,8 +348,7 @@ def test_sparse_not_supported(mock_vllm_config):
348
with pytest.raises(
349
AssertionError, match="Sparse MLA backend on ROCm only supports block size 1"
350
):
325
- RocmPlatform.get_attn_backend_cls(
326
- selected_backend=None,
351
+ attn_selector_config = AttentionSelectorConfig(
352
head_size=128,
353
dtype=torch.float16,
354
kv_cache_dtype="auto",
@@ -332,3 +357,7 @@ def test_sparse_not_supported(mock_vllm_config):
357
has_sink=False,
358
use_sparse=True,
359
)
360
+
361
+ RocmPlatform.get_attn_backend_cls(
362
+ selected_backend=None, attn_selector_config=attn_selector_config
363
+ )
tests/v1/attention/test_sparse_mla_backends.py
+4
@@ -24,6 +24,7 @@ from vllm import _custom_ops as ops
24
from vllm.attention.ops import flashmla
25
from vllm.config import set_current_vllm_config
26
from vllm.model_executor.layers.linear import ColumnParallelLinear
27
+from vllm.platforms import current_platform
28
from vllm.utils.math_utils import cdiv
29
from vllm.v1.attention.backends.mla.flashmla_sparse import (
30
FlashMLASparseBackend,
@@ -125,6 +126,9 @@ def _quantize_dequantize_fp8_ds_mla(
126
def test_sparse_backend_decode_correctness(
127
dist_init, batch_name, kv_cache_dtype, tensor_parallel_size, workspace_init
128
):
129
+ if current_platform.is_rocm():
130
+ pytest.skip("ROCm does not support fp8_ds_mla data type for kv cache.")
131
+
132
if not torch.cuda.is_available():
133
pytest.skip("CUDA is required for sparse MLA decode test")
134