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, 22:15 UTC
[CI/Build][AMD] Fix import errors in tests/kernels/attention (#29032)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Nov 20, 2025 at 03:48 UTC
322cb02872d806afcaaa7d0aac3fad7f304b7888
6 files changed
+49
-15
tests/kernels/attention/test_cascade_flash_attn.py
+13
-5
@@ -7,11 +7,19 @@ import torch
7
8
from vllm.platforms import current_platform
9
from vllm.v1.attention.backends.flash_attn import cascade_attention, merge_attn_states
10
-from vllm.vllm_flash_attn import (
11
- fa_version_unsupported_reason,
12
- flash_attn_varlen_func,
13
- is_fa_version_supported,
14
-)
10
+
11
+try:
12
+ from vllm.vllm_flash_attn import (
13
+ fa_version_unsupported_reason,
14
+ flash_attn_varlen_func,
15
+ is_fa_version_supported,
16
+ )
17
+except ImportError:
18
+ if current_platform.is_rocm():
19
+ pytest.skip(
20
+ "vllm_flash_attn is not supported for vLLM on ROCm.",
21
+ allow_module_level=True,
22
+ )
23
24
NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
25
HEAD_SIZES = [128, 192, 256]
tests/kernels/attention/test_flash_attn.py
+14
-5
@@ -6,11 +6,20 @@ import pytest
6
import torch
7
8
from vllm.platforms import current_platform
9
-from vllm.vllm_flash_attn import (
10
- fa_version_unsupported_reason,
11
- flash_attn_varlen_func,
12
- is_fa_version_supported,
13
-)
9
+
10
+try:
11
+ from vllm.vllm_flash_attn import (
12
+ fa_version_unsupported_reason,
13
+ flash_attn_varlen_func,
14
+ is_fa_version_supported,
15
+ )
16
+except ImportError:
17
+ if current_platform.is_rocm():
18
+ pytest.skip(
19
+ "vllm_flash_attn is not supported for vLLM on ROCm.",
20
+ allow_module_level=True,
21
+ )
22
+
23
24
NUM_HEADS = [(4, 4), (8, 2)]
25
HEAD_SIZES = [40, 72, 80, 128, 256]
tests/kernels/attention/test_flashinfer.py
+10
-2
@@ -2,12 +2,20 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
-import flashinfer
5
import pytest
7
-import torch
6
7
from vllm.platforms import current_platform
8
9
+try:
10
+ import flashinfer
11
+except ImportError:
12
+ if current_platform.is_rocm():
13
+ pytest.skip(
14
+ "flashinfer is not supported for vLLM on ROCm.", allow_module_level=True
15
+ )
16
+
17
+import torch
18
+
19
NUM_HEADS = [(32, 8), (6, 1)]
20
HEAD_SIZES = [128, 256]
21
BLOCK_SIZES = [16, 32]
tests/kernels/attention/test_flashinfer_mla_decode.py
+2
-1
@@ -3,7 +3,6 @@
3
import pytest
4
import torch
5
import torch.nn.functional as F
6
-from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
6
from torch import Tensor
7
8
from vllm.platforms import current_platform
@@ -15,6 +14,8 @@ if not current_platform.has_device_capability(100):
14
reason="FlashInfer MLA Requires compute capability of 10 or above.",
15
allow_module_level=True,
16
)
17
+else:
18
+ from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
19
20
21
def ref_mla(
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+2
-1
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import flashinfer
4
import pytest
5
import torch
6
@@ -16,6 +15,8 @@ if not current_platform.is_device_capability(100):
15
pytest.skip(
16
"This TRTLLM kernel requires NVIDIA Blackwell.", allow_module_level=True
17
)
18
+else:
19
+ import flashinfer
20
21
FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
22
FP8_DTYPE = current_platform.fp8_dtype()
tests/kernels/moe/test_flashinfer.py
+8
-1
@@ -22,7 +22,14 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22
from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
23
from vllm.model_executor.models.llama4 import Llama4MoE
24
from vllm.platforms import current_platform
25
-from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
25
+
26
+try:
27
+ from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
28
+except ImportError:
29
+ if current_platform.is_rocm():
30
+ pytest.skip(
31
+ "flashinfer not supported for vLLM on ROCm", allow_module_level=True
32
+ )
33
34
if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
35
90