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 11, 17:30 UTC
[CI/Build] Make test_mha_attn.py run on correct platform only and check for flash_attn_varlen_func in layer.py (#29145)
rasmith committed
Dec 9, 2025 at 14:18 UTC
7618dc973dd1e56a46162bc7bd6e7625143bead0
2 files changed
+13
-3
tests/kernels/attention/test_mha_attn.py
+9
-2
@@ -26,7 +26,14 @@ def clear_cache():
26
_cached_get_attn_backend.cache_clear()
27
28
29
-@pytest.mark.parametrize("device", ["cpu", "hip", "cuda"])
29
+devices = ["cpu"]
30
+if current_platform.is_cuda():
31
+ devices.append("cuda")
32
+if current_platform.is_rocm():
33
+ devices.append("hip")
34
+
35
+
36
+@pytest.mark.parametrize("device", devices)
37
def test_mha_attn_platform(device: str):
38
"""
39
Test the attention selector between different platform and device.
@@ -46,7 +53,7 @@ def test_mha_attn_platform(device: str):
53
patch("vllm.model_executor.models.vision.current_platform", RocmPlatform()),
54
):
55
attn = MultiHeadAttention(16, 64, scale=1)
49
- assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
56
+ assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
57
else:
58
# Test CUDA with head_size=64 (divisible by 32)
59
# - should use vLLM's FlashAttention
vllm/attention/layer.py
+4
-1
@@ -89,7 +89,10 @@ def maybe_get_vit_flash_attn_backend(
89
if attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
90
from aiter import flash_attn_varlen_func
91
else:
92
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
92
+ try:
93
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
94
+ except ImportError:
95
+ flash_attn_varlen_func = None
96
else:
97
flash_attn_varlen_func = None
98