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 14, 05:30 UTC
[Bugfix] Move current_platform import to avoid python import cache. (#16601)
Signed-off-by: iwzbi <wzbi@zju.edu.cn>
Wenzheng Bi committed
Oct 9, 2025 at 18:46 UTC
ec10fd0abcc7f5553c4f9792ce92ad93a6c7bfa5
2 files changed
+8
-7
tests/kernels/attention/test_attention_selector.py
+6
-6
@@ -84,12 +84,12 @@ def test_env(
84
m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")
85
86
if device == "cpu":
87
- with patch("vllm.attention.selector.current_platform", CpuPlatform()):
87
+ with patch("vllm.platforms.current_platform", CpuPlatform()):
88
backend = get_attn_backend(16, torch.float16, None, block_size)
89
assert backend.get_name() == "TORCH_SDPA"
90
91
elif device == "hip":
92
- with patch("vllm.attention.selector.current_platform", RocmPlatform()):
92
+ with patch("vllm.platforms.current_platform", RocmPlatform()):
93
if use_mla:
94
# ROCm MLA backend logic:
95
# - TRITON_MLA: supported when block_size != 1
@@ -126,7 +126,7 @@ def test_env(
126
assert backend.get_name() == expected
127
128
elif device == "cuda":
129
- with patch("vllm.attention.selector.current_platform", CudaPlatform()):
129
+ with patch("vllm.platforms.current_platform", CudaPlatform()):
130
if use_mla:
131
# CUDA MLA backend logic:
132
# - CUTLASS_MLA: only supported with block_size == 128
@@ -214,12 +214,12 @@ def test_env(
214
def test_fp32_fallback(device: str):
215
"""Test attention backend selection with fp32."""
216
if device == "cpu":
217
- with patch("vllm.attention.selector.current_platform", CpuPlatform()):
217
+ with patch("vllm.platforms.current_platform", CpuPlatform()):
218
backend = get_attn_backend(16, torch.float32, None, 16)
219
assert backend.get_name() == "TORCH_SDPA"
220
221
elif device == "cuda":
222
- with patch("vllm.attention.selector.current_platform", CudaPlatform()):
222
+ with patch("vllm.platforms.current_platform", CudaPlatform()):
223
backend = get_attn_backend(16, torch.float32, None, 16)
224
assert backend.get_name() == "FLEX_ATTENTION"
225
@@ -277,7 +277,7 @@ def test_invalid_env(monkeypatch: pytest.MonkeyPatch):
277
"""Test that invalid attention backend names raise ValueError."""
278
with (
279
monkeypatch.context() as m,
280
- patch("vllm.attention.selector.current_platform", CudaPlatform()),
280
+ patch("vllm.platforms.current_platform", CudaPlatform()),
281
):
282
m.setenv(STR_BACKEND_ENV_VAR, STR_INVALID_VAL)
283
vllm/attention/selector.py
+2
-1
@@ -14,7 +14,6 @@ import vllm.envs as envs
14
from vllm.attention.backends.abstract import AttentionBackend
15
from vllm.attention.backends.registry import _Backend, backend_name_to_enum
16
from vllm.logger import init_logger
17
-from vllm.platforms import current_platform
17
from vllm.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
18
19
logger = init_logger(__name__)
@@ -192,6 +191,8 @@ def _cached_get_attn_backend(
191
)
192
193
# get device-specific attn_backend
194
+ from vllm.platforms import current_platform
195
+
196
attention_cls = current_platform.get_attn_backend_cls(
197
selected_backend,
198
head_size,