From 590f07a3b56efe6ff8a6f8c63eaf11b4d1573502 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Mon, 30 Dec 2024 17:20:09 +0800 Subject: [PATCH] soft fix Signed-off-by: youkaichao --- tests/kernels/test_attention_selector.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/kernels/test_attention_selector.py b/tests/kernels/test_attention_selector.py index d37f95d48d5b2..916cc2efa3895 100644 --- a/tests/kernels/test_attention_selector.py +++ b/tests/kernels/test_attention_selector.py @@ -5,7 +5,10 @@ from tests.kernels.utils import override_backend_env_variable from vllm.attention.selector import which_attn_to_use -from vllm.platforms import cpu, cuda, openvino, rocm +from vllm.platforms.cpu import CpuPlatform +from vllm.platforms.cuda import CudaPlatform +from vllm.platforms.openvino import OpenVinoPlatform +from vllm.platforms.rocm import RocmPlatform from vllm.utils import STR_FLASH_ATTN_VAL, STR_INVALID_VAL @@ -20,26 +23,23 @@ def test_env(name: str, device: str, monkeypatch): override_backend_env_variable(monkeypatch, name) if device == "cpu": - with patch("vllm.attention.selector.current_platform", - cpu.CpuPlatform()): + with patch("vllm.attention.selector.current_platform", CpuPlatform()): backend = which_attn_to_use(16, torch.float16, torch.float16, 16, False) assert backend.name == "TORCH_SDPA" elif device == "hip": - with patch("vllm.attention.selector.current_platform", - rocm.RocmPlatform()): + with patch("vllm.attention.selector.current_platform", RocmPlatform()): backend = which_attn_to_use(16, torch.float16, torch.float16, 16, False) assert backend.name == "ROCM_FLASH" elif device == "openvino": with patch("vllm.attention.selector.current_platform", - openvino.OpenVinoPlatform()): + OpenVinoPlatform()): backend = which_attn_to_use(16, torch.float16, torch.float16, 16, False) assert backend.name == "OPENVINO" else: - with patch("vllm.attention.selector.current_platform", - cuda.CudaPlatform()): + with patch("vllm.attention.selector.current_platform", CudaPlatform()): backend = which_attn_to_use(16, torch.float16, torch.float16, 16, False) assert backend.name == name