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, 11:15 UTC
[Platform] Refactor Platform attention backend selection to avoid breakpoint for OOT platform (#30212)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Isotr0py committed
Dec 16, 2025 at 01:36 UTC
ec154c36ee74f35def28e4ddc1c16a0dc7a8c112
7 files changed
+73
-137
vllm/attention/selector.py
+36
-23
@@ -2,11 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from functools import cache
5
-from typing import cast, get_args
5
+from typing import NamedTuple, cast, get_args
6
7
import torch
8
9
-from vllm.attention.backends.abstract import AttentionBackend
9
+from vllm.attention.backends.abstract import AttentionBackend, AttentionType
10
from vllm.attention.backends.registry import (
11
MAMBA_TYPE_TO_BACKEND_MAP,
12
MambaAttentionBackendEnum,
@@ -18,6 +18,31 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
18
logger = init_logger(__name__)
19
20
21
+class AttentionSelectorConfig(NamedTuple):
22
+ head_size: int
23
+ dtype: torch.dtype
24
+ kv_cache_dtype: CacheDType | None
25
+ block_size: int | None
26
+ use_mla: bool = False
27
+ has_sink: bool = False
28
+ use_sparse: bool = False
29
+ use_mm_prefix: bool = False
30
+ attn_type: str = AttentionType.DECODER
31
+
32
+ def __repr__(self):
33
+ return (
34
+ f"AttentionSelectorConfig(head_size={self.head_size}, "
35
+ f"dtype={self.dtype}, "
36
+ f"kv_cache_dtype={self.kv_cache_dtype}, "
37
+ f"block_size={self.block_size}, "
38
+ f"use_mla={self.use_mla}, "
39
+ f"has_sink={self.has_sink}, "
40
+ f"use_sparse={self.use_sparse}, "
41
+ f"use_mm_prefix={self.use_mm_prefix}, "
42
+ f"attn_type={self.attn_type})"
43
+ )
44
+
45
+
46
def get_attn_backend(
47
head_size: int,
48
dtype: torch.dtype,
@@ -43,8 +68,7 @@ def get_attn_backend(
68
vllm_config = get_current_vllm_config()
69
backend_enum = vllm_config.attention_config.backend
70
46
- return _cached_get_attn_backend(
47
- backend=backend_enum,
71
+ attn_selector_config = AttentionSelectorConfig(
72
head_size=head_size,
73
dtype=dtype,
74
kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype),
@@ -53,36 +77,25 @@ def get_attn_backend(
77
has_sink=has_sink,
78
use_sparse=use_sparse,
79
use_mm_prefix=use_mm_prefix,
56
- attn_type=attn_type,
80
+ attn_type=attn_type or AttentionType.DECODER,
81
+ )
82
+
83
+ return _cached_get_attn_backend(
84
+ backend=backend_enum,
85
+ attn_selector_config=attn_selector_config,
86
)
87
88
89
@cache
90
def _cached_get_attn_backend(
91
backend,
63
- head_size: int,
64
- dtype: torch.dtype,
65
- kv_cache_dtype: CacheDType | None,
66
- block_size: int | None,
67
- use_mla: bool = False,
68
- has_sink: bool = False,
69
- use_sparse: bool = False,
70
- use_mm_prefix: bool = False,
71
- attn_type: str | None = None,
92
+ attn_selector_config: AttentionSelectorConfig,
93
) -> type[AttentionBackend]:
94
from vllm.platforms import current_platform
95
96
attention_cls = current_platform.get_attn_backend_cls(
97
backend,
77
- head_size,
78
- dtype,
79
- kv_cache_dtype,
80
- block_size,
81
- use_mla,
82
- has_sink,
83
- use_sparse,
84
- use_mm_prefix,
85
- attn_type,
98
+ attn_selector_config=attn_selector_config,
99
)
100
if not attention_cls:
101
raise ValueError(
vllm/platforms/cpu.py
+4
-11
@@ -23,6 +23,7 @@ from .interface import CpuArchEnum, Platform, PlatformEnum
23
logger = init_logger(__name__)
24
25
if TYPE_CHECKING:
26
+ from vllm.attention.selector import AttentionSelectorConfig
27
from vllm.config import VllmConfig
28
else:
29
VllmConfig = None
@@ -126,21 +127,13 @@ class CpuPlatform(Platform):
127
def get_attn_backend_cls(
128
cls,
129
selected_backend: "AttentionBackendEnum",
129
- head_size: int,
130
- dtype: torch.dtype,
131
- kv_cache_dtype: str | None,
132
- block_size: int,
133
- use_mla: bool,
134
- has_sink: bool,
135
- use_sparse: bool,
136
- use_mm_prefix: bool,
137
- attn_type: str | None = None,
130
+ attn_selector_config: "AttentionSelectorConfig",
131
) -> str:
132
if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:
133
logger.info("Cannot use %s backend on CPU.", selected_backend)
141
- if use_mla:
134
+ if attn_selector_config.use_mla:
135
raise NotImplementedError("MLA is not supported on CPU.")
143
- if use_sparse:
136
+ if attn_selector_config.use_sparse:
137
raise NotImplementedError("Sparse Attention is not supported on CPU.")
138
return AttentionBackendEnum.CPU_ATTN.get_path()
139
vllm/platforms/cuda.py
+15
-59
@@ -14,7 +14,6 @@ from typing_extensions import ParamSpec
14
15
# import custom ops, trigger op registration
16
import vllm._C # noqa
17
-from vllm.attention.backends.abstract import AttentionType
17
from vllm.attention.backends.registry import AttentionBackendEnum
18
from vllm.logger import init_logger
19
from vllm.utils.import_utils import import_pynvml
@@ -23,6 +22,7 @@ from vllm.utils.torch_utils import cuda_device_count_stateless
22
from .interface import DeviceCapability, Platform, PlatformEnum
23
24
if TYPE_CHECKING:
25
+ from vllm.attention.selector import AttentionSelectorConfig
26
from vllm.config import VllmConfig
27
from vllm.config.cache import CacheDType
28
else:
@@ -258,16 +258,8 @@ class CudaPlatformBase(Platform):
258
@classmethod
259
def get_valid_backends(
260
cls,
261
- head_size,
262
- dtype,
263
- kv_cache_dtype,
264
- block_size,
265
- use_mla,
266
- has_sink,
267
- use_sparse,
268
- use_mm_prefix,
269
- device_capability,
270
- attn_type,
261
+ device_capability: DeviceCapability,
262
+ attn_selector_config: "AttentionSelectorConfig",
263
) -> tuple[
264
list[tuple["AttentionBackendEnum", int]],
265
dict["AttentionBackendEnum", list[str]],
@@ -275,21 +267,15 @@ class CudaPlatformBase(Platform):
267
valid_backends_priorities = []
268
invalid_reasons = {}
269
278
- backend_priorities = _get_backend_priorities(use_mla, device_capability)
270
+ backend_priorities = _get_backend_priorities(
271
+ attn_selector_config.use_mla, device_capability
272
+ )
273
for priority, backend in enumerate(backend_priorities):
274
try:
275
backend_class = backend.get_class()
276
invalid_reasons_i = backend_class.validate_configuration(
283
- head_size,
284
- dtype,
285
- kv_cache_dtype,
286
- block_size,
287
- use_mla,
288
- has_sink,
289
- use_sparse,
290
- use_mm_prefix,
291
- device_capability,
292
- attn_type,
277
+ device_capability=device_capability,
278
+ **attn_selector_config._asdict(),
279
)
280
except ImportError:
281
invalid_reasons_i = ["ImportError"]
@@ -304,37 +290,19 @@ class CudaPlatformBase(Platform):
290
def get_attn_backend_cls(
291
cls,
292
selected_backend: "AttentionBackendEnum",
307
- head_size: int,
308
- dtype: torch.dtype,
309
- kv_cache_dtype: "CacheDType | None",
310
- block_size: int | None,
311
- use_mla: bool,
312
- has_sink: bool,
313
- use_sparse: bool,
314
- use_mm_prefix: bool,
315
- attn_type: str | None = None,
293
+ attn_selector_config: "AttentionSelectorConfig",
294
) -> str:
317
- if attn_type is None:
318
- attn_type = AttentionType.DECODER
319
-
295
device_capability = cls.get_device_capability()
296
assert device_capability is not None
297
298
+ attn_selector_config = attn_selector_config._replace(block_size=None)
299
# First try checking just the selected backend, if there is one.
300
if selected_backend is not None:
301
try:
302
backend_class = selected_backend.get_class()
303
invalid_reasons = backend_class.validate_configuration(
328
- head_size,
329
- dtype,
330
- kv_cache_dtype,
331
- None,
332
- use_mla,
333
- has_sink,
334
- use_sparse,
335
- use_mm_prefix,
336
- device_capability,
337
- attn_type,
304
+ device_capability=device_capability,
305
+ **attn_selector_config._asdict(),
306
)
307
except ImportError:
308
invalid_reasons = ["ImportError"]
@@ -350,16 +318,8 @@ class CudaPlatformBase(Platform):
318
# No selected backend or the selected backend is invalid,
319
# so we try finding a valid backend.
320
valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
353
- head_size,
354
- dtype,
355
- kv_cache_dtype,
356
- None,
357
- use_mla,
358
- has_sink,
359
- use_sparse,
360
- use_mm_prefix,
361
- device_capability,
362
- attn_type,
321
+ device_capability=device_capability,
322
+ attn_selector_config=attn_selector_config,
323
)
324
reasons_str = (
325
"{"
@@ -369,11 +329,7 @@ class CudaPlatformBase(Platform):
329
)
330
+ "}"
331
)
372
- config_str = (
373
- f"head_size: {head_size}, dtype: {dtype}, "
374
- f"kv_cache_dtype: {kv_cache_dtype}, block_size: {block_size}, "
375
- f"use_mla: {use_mla}, has_sink: {has_sink}, use_sparse: {use_sparse}"
376
- )
332
+ config_str = attn_selector_config.__repr__()
333
logger.debug_once(
334
f"Some attention backends are not valid for {cls.device_name} with "
335
f"{config_str}. Reasons: {reasons_str}."
vllm/platforms/interface.py
+2
-10
@@ -18,8 +18,8 @@ from vllm.logger import init_logger
18
if TYPE_CHECKING:
19
from torch.distributed import PrefixStore, ProcessGroup
20
21
+ from vllm.attention.selector import AttentionSelectorConfig
22
from vllm.config import VllmConfig
22
- from vllm.config.cache import CacheDType
23
from vllm.inputs import ProcessorInputs, PromptType
24
from vllm.pooling_params import PoolingParams
25
from vllm.sampling_params import SamplingParams
@@ -226,15 +226,7 @@ class Platform:
226
def get_attn_backend_cls(
227
cls,
228
selected_backend: "AttentionBackendEnum",
229
- head_size: int,
230
- dtype: torch.dtype,
231
- kv_cache_dtype: "CacheDType | None",
232
- block_size: int,
233
- use_mla: bool,
234
- has_sink: bool,
235
- use_sparse: bool,
236
- use_mm_prefix: bool,
237
- attn_type: str | None = None,
229
+ attn_selector_config: "AttentionSelectorConfig",
230
) -> str:
231
"""Get the attention backend class of a device."""
232
return ""
vllm/platforms/rocm.py
+9
-13
@@ -15,6 +15,7 @@ from vllm.utils.torch_utils import cuda_device_count_stateless
15
from .interface import DeviceCapability, Platform, PlatformEnum
16
17
if TYPE_CHECKING:
18
+ from vllm.attention.selector import AttentionSelectorConfig
19
from vllm.config import VllmConfig
20
21
logger = init_logger(__name__)
@@ -190,21 +191,16 @@ class RocmPlatform(Platform):
191
@classmethod
192
def get_attn_backend_cls(
193
cls,
193
- selected_backend,
194
- head_size,
195
- dtype,
196
- kv_cache_dtype,
197
- block_size,
198
- use_mla,
199
- has_sink,
200
- use_sparse,
201
- use_mm_prefix,
202
- attn_type: str | None = None,
194
+ selected_backend: "AttentionBackendEnum",
195
+ attn_selector_config: "AttentionSelectorConfig",
196
) -> str:
197
from vllm._aiter_ops import rocm_aiter_ops
198
206
- if use_sparse:
207
- if kv_cache_dtype.startswith("fp8"):
199
+ block_size = attn_selector_config.block_size
200
+ kv_cache_dtype = attn_selector_config.kv_cache_dtype
201
+
202
+ if attn_selector_config.use_sparse:
203
+ if kv_cache_dtype and kv_cache_dtype.startswith("fp8"):
204
raise ValueError(
205
"ROCMAiterMLASparseBackend doesn't support fp8 kv_cache_dtype."
206
)
@@ -214,7 +210,7 @@ class RocmPlatform(Platform):
210
logger.info_once("Using Sparse MLA backend on V1 engine.")
211
return AttentionBackendEnum.ROCM_AITER_MLA_SPARSE.get_path()
212
217
- if use_mla:
213
+ if attn_selector_config.use_mla:
214
if selected_backend is None:
215
selected_backend = (
216
AttentionBackendEnum.ROCM_AITER_MLA
vllm/platforms/tpu.py
+3
-10
@@ -16,6 +16,7 @@ from .interface import Platform, PlatformEnum
16
if TYPE_CHECKING:
17
from typing import TypeAlias
18
19
+ from vllm.attention.selector import AttentionSelectorConfig
20
from vllm.config import VllmConfig
21
from vllm.config.cache import BlockSize
22
from vllm.pooling_params import PoolingParams
@@ -57,17 +58,9 @@ class TpuPlatform(Platform):
58
def get_attn_backend_cls(
59
cls,
60
selected_backend: "AttentionBackendEnum",
60
- head_size: int,
61
- dtype: torch.dtype,
62
- kv_cache_dtype: str | None,
63
- block_size: int,
64
- use_mla: bool,
65
- has_sink: bool,
66
- use_sparse: bool,
67
- use_mm_prefix: bool,
68
- attn_type: str | None = None,
61
+ attn_selector_config: "AttentionSelectorConfig",
62
) -> str:
70
- if use_sparse:
63
+ if attn_selector_config.use_sparse:
64
raise NotImplementedError("Sparse Attention is not supported on TPU.")
65
if selected_backend != AttentionBackendEnum.PALLAS:
66
logger.info("Cannot use %s backend on TPU.", selected_backend)
vllm/platforms/xpu.py
+4
-11
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14
from .interface import DeviceCapability, Platform, PlatformEnum
15
16
if TYPE_CHECKING:
17
+ from vllm.attention.selector import AttentionSelectorConfig
18
from vllm.config import VllmConfig
19
else:
20
VllmConfig = None
@@ -42,15 +43,7 @@ class XPUPlatform(Platform):
43
def get_attn_backend_cls(
44
cls,
45
selected_backend: "AttentionBackendEnum",
45
- head_size: int,
46
- dtype: torch.dtype,
47
- kv_cache_dtype: str | None,
48
- block_size: int,
49
- use_mla: bool,
50
- has_sink: bool,
51
- use_sparse: bool,
52
- use_mm_prefix: bool,
53
- attn_type: str | None = None,
46
+ attn_selector_config: "AttentionSelectorConfig",
47
) -> str:
48
from vllm.v1.attention.backends.utils import set_kv_cache_layout
49
@@ -60,7 +53,7 @@ class XPUPlatform(Platform):
53
"only NHD layout is supported by XPU attention kernels."
54
)
55
63
- if use_sparse:
56
+ if attn_selector_config.use_sparse:
57
raise NotImplementedError("Sparse Attention is not supported on XPU.")
58
if selected_backend == AttentionBackendEnum.TRITON_ATTN:
59
logger.info_once("Using Triton backend.")
@@ -71,7 +64,7 @@ class XPUPlatform(Platform):
64
elif selected_backend:
65
raise ValueError(
66
f"Invalid attention backend for {cls.device_name}, "
74
- f"with use_mla: {use_mla}"
67
+ f"with use_mla: {attn_selector_config.use_mla}"
68
)
69
70
logger.info("Using Flash Attention backend.")