@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b30dfa03c5
Fork 0

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 13, 13:15 UTC

[Attention] Refactor CUDA attention backend selection logic (#24794)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Matthew Bonanni committed Nov 11, 2025 at 06:40 UTC b30dfa03c564ce51c56bf2dd16283f074253c27c
61 files changed +1333 -997
.buildkite/test-pipeline.yaml
+5
@@ -890,11 +890,16 @@ steps:
890 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
891 - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
892 - vllm/v1/attention/backends/flashinfer.py
893 + - vllm/v1/attention/backends/mla/cutlass_mla.py
894 + - vllm/v1/attention/backends/mla/flashinfer_mla.py
895 + - vllm/platforms/cuda.py
896 + - vllm/attention/selector.py
897 commands:
898 - nvidia-smi
899 - python3 examples/offline_inference/basic/chat.py
900 # Attention
901 # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
902 + - pytest -v -s tests/kernels/attention/test_attention_selector.py
903 - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
904 - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
905 - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
tests/compile/test_fusion_attn.py
+16 -15
@@ -10,7 +10,7 @@ from tests.utils import flat_product
10 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
11 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12 from vllm.attention import Attention, AttentionMetadata
13 -from vllm.attention.backends.registry import _Backend
13 +from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.attention.selector import global_force_attn_backend_context_manager
15 from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
16 from vllm.compilation.fx_utils import find_op_nodes
@@ -104,7 +104,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
104
105 # TODO(luka) use get_kv_cache_stride_order
106 # Create dummy KV cache for the selected backend
107 - if backend == _Backend.ROCM_ATTN:
107 + if backend == AttentionBackendEnum.ROCM_ATTN:
108 # k/v as 1st dimention
109 # HND: [num_blocks, num_kv_heads, block_size, head_size]
110 kv_cache = torch.zeros(
@@ -116,7 +116,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
116 dtype=self.kv_cache_dtype,
117 device=self.device,
118 )
119 - elif backend == _Backend.ROCM_AITER_UNIFIED_ATTN:
119 + elif backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
120 # k/v as 1st dimention
121 # NHD: [num_blocks, block_size, num_kv_heads, head_size]
122 kv_cache = torch.zeros(
@@ -128,7 +128,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
128 dtype=self.kv_cache_dtype,
129 device=self.device,
130 )
131 - elif backend == _Backend.TRITON_ATTN:
131 + elif backend == AttentionBackendEnum.TRITON_ATTN:
132 # k/v as 2nd dimention
133 # NHD: [num_blocks, block_size, num_kv_heads, head_size]
134 kv_cache = torch.zeros(
@@ -140,7 +140,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
140 dtype=self.kv_cache_dtype,
141 device=self.device,
142 )
143 - elif backend == _Backend.FLASHINFER:
143 + elif backend == AttentionBackendEnum.FLASHINFER:
144 kv_cache = torch.zeros(
145 num_blocks,
146 2,
@@ -244,8 +244,8 @@ MODELS_FP8: list[tuple[str, type]] = []
244 MODELS_FP4: list[tuple[str, type]] = []
245 HEADS: list[tuple[int, int]] = []
246 SPLIT_ATTENTION: list[bool] = []
247 -BACKENDS_FP8: list[_Backend] = []
248 -BACKENDS_FP4: list[_Backend] = []
247 +BACKENDS_FP8: list[AttentionBackendEnum] = []
248 +BACKENDS_FP4: list[AttentionBackendEnum] = []
249
250 if current_platform.is_cuda():
251 HEADS = [(64, 8), (40, 8)]
@@ -261,8 +261,8 @@ if current_platform.is_cuda():
261 TestAttentionNvfp4QuantPatternModel,
262 )
263 ]
264 - BACKENDS_FP8 = [_Backend.TRITON_ATTN, _Backend.FLASHINFER]
265 - BACKENDS_FP4 = [_Backend.FLASHINFER]
264 + BACKENDS_FP8 = [AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.FLASHINFER]
265 + BACKENDS_FP4 = [AttentionBackendEnum.FLASHINFER]
266
267 elif current_platform.is_rocm():
268 HEADS = [(32, 8), (40, 8)]
@@ -270,9 +270,9 @@ elif current_platform.is_rocm():
270 ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)
271 ]
272 BACKENDS = [
273 - _Backend.ROCM_AITER_UNIFIED_ATTN,
274 - _Backend.ROCM_ATTN,
275 - _Backend.TRITON_ATTN,
273 + AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
274 + AttentionBackendEnum.ROCM_ATTN,
275 + AttentionBackendEnum.TRITON_ATTN,
276 ]
277
278
@@ -302,11 +302,11 @@ def test_attention_quant_pattern(
302 custom_ops: str,
303 model_name: str,
304 model_class: type[AttentionQuantPatternModel],
305 - backend: _Backend,
305 + backend: AttentionBackendEnum,
306 dist_init,
307 ):
308 """Test AttentionStaticQuantPattern fusion pass"""
309 - if backend == _Backend.FLASHINFER and (
309 + if backend == AttentionBackendEnum.FLASHINFER and (
310 not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
311 ):
312 pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
@@ -314,6 +314,7 @@ def test_attention_quant_pattern(
314 custom_ops_list = custom_ops.split(",") if custom_ops else []
315
316 device = torch.device("cuda:0")
317 + torch.set_default_dtype(dtype)
318 torch.manual_seed(42)
319
320 vllm_config = VllmConfig(
@@ -402,7 +403,7 @@ def test_attention_quant_pattern(
403
404 result_fused_1 = model_compiled(q, k, v)
405
405 - if backend == _Backend.FLASHINFER:
406 + if backend == AttentionBackendEnum.FLASHINFER:
407 # With the Flashinfer backend after the 1st round of the forward
408 # pass, output quant scale should be loaded into the attn layer's
409 # _o_scale_float, the 2nd round should reuse the loaded
tests/compile/test_fusions_e2e.py
+12 -12
@@ -11,7 +11,7 @@ from typing import Any, NamedTuple
11 import pytest
12 import regex as re
13
14 -from tests.v1.attention.utils import _Backend
14 +from tests.v1.attention.utils import AttentionBackendEnum
15 from vllm import LLM, SamplingParams
16 from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
17 from vllm.platforms import current_platform
@@ -24,7 +24,7 @@ from ..utils import flat_product, multi_gpu_test
24 class ModelBackendTestCase(NamedTuple):
25 model_name: str
26 model_kwargs: dict[str, Any]
27 - backend: _Backend
27 + backend: AttentionBackendEnum
28 attention_fusions: int
29 allreduce_fusions: int | None = None
30
@@ -39,14 +39,14 @@ if current_platform.is_cuda():
39 # Use smaller model for L40s in CI
40 model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
41 model_kwargs=dict(max_model_len=1024),
42 - backend=_Backend.TRITON_ATTN,
42 + backend=AttentionBackendEnum.TRITON_ATTN,
43 attention_fusions=32,
44 allreduce_fusions=65,
45 ),
46 ModelBackendTestCase(
47 model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
48 model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
49 - backend=_Backend.FLASHINFER,
49 + backend=AttentionBackendEnum.FLASHINFER,
50 attention_fusions=48,
51 allreduce_fusions=96,
52 ),
@@ -56,7 +56,7 @@ if current_platform.is_cuda():
56 ModelBackendTestCase(
57 model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
58 model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
59 - backend=_Backend.FLASHINFER,
59 + backend=AttentionBackendEnum.FLASHINFER,
60 attention_fusions=32,
61 allreduce_fusions=65,
62 ),
@@ -67,7 +67,7 @@ if current_platform.is_cuda():
67 ModelBackendTestCase(
68 model_name="meta-llama/Llama-3.1-8B-Instruct",
69 model_kwargs=dict(max_model_len=1024),
70 - backend=_Backend.TRITON_ATTN,
70 + backend=AttentionBackendEnum.TRITON_ATTN,
71 attention_fusions=0,
72 allreduce_fusions=65,
73 ),
@@ -85,19 +85,19 @@ elif current_platform.is_rocm():
85 ModelBackendTestCase(
86 model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
87 model_kwargs=dict(max_model_len=1024),
88 - backend=_Backend.TRITON_ATTN,
88 + backend=AttentionBackendEnum.TRITON_ATTN,
89 attention_fusions=32,
90 ),
91 ModelBackendTestCase(
92 model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
93 model_kwargs=dict(max_model_len=1024),
94 - backend=_Backend.ROCM_ATTN,
94 + backend=AttentionBackendEnum.ROCM_ATTN,
95 attention_fusions=32,
96 ),
97 ModelBackendTestCase(
98 model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
99 model_kwargs=dict(max_model_len=1024),
100 - backend=_Backend.ROCM_AITER_UNIFIED_ATTN,
100 + backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
101 attention_fusions=32,
102 ),
103 ]
@@ -117,7 +117,7 @@ CUSTOM_OPS_FP8 = ["-quant_fp8", "+quant_fp8"]
117 def test_attn_quant(
118 model_name: str,
119 model_kwargs: dict[str, Any],
120 - backend: _Backend,
120 + backend: AttentionBackendEnum,
121 attention_fusions: int,
122 allreduce_fusions: int,
123 custom_ops: str,
@@ -125,7 +125,7 @@ def test_attn_quant(
125 caplog_mp_spawn,
126 monkeypatch,
127 ):
128 - if backend == _Backend.FLASHINFER and (
128 + if backend == AttentionBackendEnum.FLASHINFER and (
129 not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
130 ):
131 pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
@@ -208,7 +208,7 @@ def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]:
208 def test_tp2_attn_quant_allreduce_rmsnorm(
209 model_name: str,
210 model_kwargs: dict,
211 - backend: _Backend,
211 + backend: AttentionBackendEnum,
212 attention_fusions: int,
213 allreduce_fusions: int,
214 custom_ops: str,
tests/config/test_multimodal_config.py
+3 -3
@@ -3,13 +3,13 @@
3
4 import pytest
5
6 -from vllm.attention.backends.registry import _Backend
6 +from vllm.attention.backends.registry import AttentionBackendEnum
7 from vllm.config.multimodal import MultiModalConfig
8
9
10 def test_mm_encoder_attn_backend_str_conversion():
11 config = MultiModalConfig(mm_encoder_attn_backend="FLASH_ATTN")
12 - assert config.mm_encoder_attn_backend == _Backend.FLASH_ATTN
12 + assert config.mm_encoder_attn_backend == AttentionBackendEnum.FLASH_ATTN
13
14
15 def test_mm_encoder_attn_backend_invalid():
@@ -20,6 +20,6 @@ def test_mm_encoder_attn_backend_invalid():
20 def test_mm_encoder_attn_backend_hash_updates():
21 base_hash = MultiModalConfig().compute_hash()
22 overridden_hash = MultiModalConfig(
23 - mm_encoder_attn_backend=_Backend.FLASH_ATTN
23 + mm_encoder_attn_backend=AttentionBackendEnum.FLASH_ATTN
24 ).compute_hash()
25 assert base_hash != overridden_hash
tests/kernels/attention/test_attention_selector.py
+45 -30
@@ -120,12 +120,13 @@ def test_env(
120
121 elif device == "cuda":
122 with patch("vllm.platforms.current_platform", CudaPlatform()):
123 + capability = torch.cuda.get_device_capability()
124 if use_mla:
125 # CUDA MLA backend logic:
126 # - CUTLASS_MLA: only supported with block_size == 128
126 - # and Blackwell GPUs (SM 10.0), V1 only
127 + # and Blackwell GPUs (SM 10.x), V1 only
128 # - FLASHINFER_MLA: only supported on Blackwell GPUs
128 - # (SM 10.0+), V1 only
129 + # (SM 10.x), V1 only
130 # - FLASHMLA: only supported with block_size == 64
131 # - FLASH_ATTN_MLA: V1 only
132 # - TRITON_MLA: fallback for other cases
@@ -134,58 +135,72 @@ def test_env(
135 if block_size != 128:
136 # CUTLASS_MLA only supports block_size == 128
137 pytest.skip("CUTLASS_MLA only supports block_size 128")
137 - else:
138 - backend = get_attn_backend(
139 - 16, torch.float16, None, block_size, use_mla=use_mla
140 - )
141 - expected = "CUTLASS_MLA"
142 - assert backend.get_name() == expected
138 + if capability[0] != 10:
139 + pytest.skip("CUTLASS MLA is not supported on this platform")
140 + backend = get_attn_backend(
141 + 576, torch.float16, None, block_size, use_mla=use_mla
142 + )
143 + expected = "CUTLASS_MLA"
144 + assert backend.get_name() == expected
145 elif name == "FLASHINFER_MLA":
146 + if capability[0] != 10:
147 + pytest.skip(
148 + "FlashInfer MLA is not supported on this platform"
149 + )
150 if block_size not in [32, 64]:
151 # FlashInfer MLA only supports block_size 32 or 64
152 pytest.skip(
153 "FlashInfer MLA only supports block_size 32 or 64"
154 )
149 - else:
150 - backend = get_attn_backend(
151 - 16, torch.float16, None, block_size, use_mla=use_mla
152 - )
153 - expected = "FLASHINFER_MLA"
154 - assert backend.get_name() == expected
155 + backend = get_attn_backend(
156 + 576, torch.float16, None, block_size, use_mla=use_mla
157 + )
158 + expected = "FLASHINFER_MLA"
159 + assert backend.get_name() == expected
160 elif name == "FLASHMLA":
161 if block_size != 64:
162 # FlashMLA only supports block_size == 64
163 pytest.skip("FlashMLA only supports block_size 64")
159 - else:
160 - from vllm.v1.attention.backends.mla.flashmla import (
161 - is_flashmla_dense_supported,
162 - )
164 + from vllm.v1.attention.backends.mla.flashmla import (
165 + is_flashmla_dense_supported,
166 + )
167
164 - is_supported, _ = is_flashmla_dense_supported()
165 - if not is_supported:
166 - pytest.skip("FlashMLA not supported on this platform")
167 - else:
168 - backend = get_attn_backend(
169 - 16, torch.float16, None, block_size, use_mla=use_mla
170 - )
171 - expected = name
172 - assert backend.get_name() == expected
168 + is_supported, _ = is_flashmla_dense_supported()
169 + if not is_supported:
170 + pytest.skip("FlashMLA not supported on this platform")
171 + backend = get_attn_backend(
172 + 576,
173 + torch.float16,
174 + None,
175 + block_size,
176 + use_mla=use_mla,
177 + )
178 + expected = name
179 + assert backend.get_name() == expected
180 elif name == "FLASH_ATTN_MLA":
181 + from vllm.attention.utils.fa_utils import (
182 + flash_attn_supports_mla,
183 + )
184 +
185 + if not flash_attn_supports_mla():
186 + pytest.skip(
187 + "FlashAttention MLA not supported on this platform"
188 + )
189 backend = get_attn_backend(
175 - 16, torch.float16, None, block_size, use_mla=use_mla
190 + 576, torch.float16, None, block_size, use_mla=use_mla
191 )
192 expected = "FLASH_ATTN_MLA"
193 assert backend.get_name() == expected
194 else:
195 # TRITON_MLA or other fallback
196 backend = get_attn_backend(
182 - 16, torch.float16, None, block_size, use_mla=use_mla
197 + 576, torch.float16, None, block_size, use_mla=use_mla
198 )
199 expected = "TRITON_MLA"
200 assert backend.get_name() == expected
201 elif name == "FLASHINFER":
202 backend = get_attn_backend(
188 - 16, torch.float16, None, block_size, use_mla=use_mla
203 + 64, torch.float16, None, block_size, use_mla=use_mla
204 )
205 expected = "FLASHINFER"
206 assert backend.get_name() == expected
tests/kernels/attention/test_mha_attn.py
+6 -6
@@ -11,7 +11,7 @@ from unittest.mock import patch
11 import pytest
12 import torch
13
14 -from vllm.attention.backends.registry import _Backend
14 +from vllm.attention.backends.registry import AttentionBackendEnum
15 from vllm.attention.layer import MultiHeadAttention
16 from vllm.attention.selector import _cached_get_attn_backend
17 from vllm.platforms import current_platform
@@ -43,14 +43,14 @@ def test_mha_attn_platform(device: str):
43 patch("vllm.model_executor.models.vision.current_platform", CpuPlatform()),
44 ):
45 attn = MultiHeadAttention(16, 64, scale=1)
46 - assert attn.attn_backend == _Backend.TORCH_SDPA
46 + assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
47 elif device == "hip":
48 with (
49 patch("vllm.attention.layer.current_platform", RocmPlatform()),
50 patch("vllm.model_executor.models.vision.current_platform", RocmPlatform()),
51 ):
52 attn = MultiHeadAttention(16, 64, scale=1)
53 - assert attn.attn_backend == _Backend.TORCH_SDPA
53 + assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
54 else:
55 # Test CUDA with head_size=64 (divisible by 32)
56 # - should use vLLM's FlashAttention
@@ -59,7 +59,7 @@ def test_mha_attn_platform(device: str):
59 patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
60 ):
61 attn = MultiHeadAttention(16, 64, scale=1)
62 - assert attn.attn_backend == _Backend.FLASH_ATTN
62 + assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
63
64 # Test CUDA with head_size=72 (not divisible by 32)
65 # - with upstream FA not available
@@ -73,7 +73,7 @@ def test_mha_attn_platform(device: str):
73 ),
74 ):
75 attn = MultiHeadAttention(16, 72, scale=1)
76 - assert attn.attn_backend == _Backend.XFORMERS
76 + assert attn.attn_backend == AttentionBackendEnum.XFORMERS
77
78 # Test CUDA with head_size=72 (not divisible by 32)
79 # - with upstream FA available
@@ -96,7 +96,7 @@ def test_mha_attn_platform(device: str):
96 ),
97 ):
98 attn = MultiHeadAttention(16, 72, scale=1)
99 - assert attn.attn_backend == _Backend.FLASH_ATTN
99 + assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
100
101
102 def ref_attention(
tests/models/test_initialization.py
+11
@@ -93,6 +93,17 @@ def can_initialize(
93 "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
94 )
95
96 + if model_arch == "DeepseekV32ForCausalLM":
97 + from vllm.platforms import current_platform
98 +
99 + capability = current_platform.get_device_capability()
100 + if capability and capability.major < 9:
101 + pytest.skip(
102 + f"DeepseekV32 requires Hopper (9.0+) or Blackwell (10.0+) "
103 + f"for FLASHMLA_SPARSE backend. Current device has compute "
104 + f"capability {capability.major}.{capability.minor}"
105 + )
106 +
107 with (
108 patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1),
109 monkeypatch.context() as m,
tests/v1/attention/test_attention_backends.py
+27 -20
@@ -15,7 +15,7 @@ from tests.v1.attention.utils import (
15 create_vllm_config,
16 try_get_attention_backend,
17 )
18 -from vllm.attention.backends.registry import _Backend
18 +from vllm.attention.backends.registry import AttentionBackendEnum
19 from vllm.config import ModelConfig
20 from vllm.platforms import current_platform
21 from vllm.utils.math_utils import cdiv
@@ -27,11 +27,11 @@ from vllm.v1.attention.backends.utils import (
27 from vllm.v1.kv_cache_interface import FullAttentionSpec
28
29 BACKENDS_TO_TEST = [
30 - _Backend.FLASH_ATTN,
31 - _Backend.FLASHINFER,
32 - _Backend.FLEX_ATTENTION,
33 - _Backend.TRITON_ATTN,
34 - _Backend.TREE_ATTN,
30 + AttentionBackendEnum.FLASH_ATTN,
31 + AttentionBackendEnum.FLASHINFER,
32 + AttentionBackendEnum.FLEX_ATTENTION,
33 + AttentionBackendEnum.TRITON_ATTN,
34 + AttentionBackendEnum.TREE_ATTN,
35 "FLEX_ATTENTION_SLOW",
36 ]
37
@@ -39,7 +39,7 @@ BACKENDS_TO_TEST = [
39 try:
40 import flashinfer # noqa: F401
41 except ImportError:
42 - BACKENDS_TO_TEST.remove(_Backend.FLASHINFER)
42 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER)
43
44
45 def _convert_dtype_to_torch(dtype):
@@ -192,7 +192,7 @@ class MockAttentionLayer:
192
193
194 def run_attention_backend(
195 - backend: _Backend,
195 + backend: AttentionBackendEnum,
196 kv_cache_spec: FullAttentionSpec,
197 layer_names: list[str],
198 vllm_config,
@@ -211,13 +211,13 @@ def run_attention_backend(
211
212 use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
213 if backend == "FLEX_ATTENTION_SLOW":
214 - actual_backend = _Backend.FLEX_ATTENTION
214 + actual_backend = AttentionBackendEnum.FLEX_ATTENTION
215 use_direct_block_mask = False
216
217 builder_cls, impl_cls = try_get_attention_backend(actual_backend)
218
219 # Mock flashinfer's get_per_layer_parameters if needed
220 - if actual_backend == _Backend.FLASHINFER:
220 + if actual_backend == AttentionBackendEnum.FLASHINFER:
221 import unittest.mock
222
223 from vllm.v1.attention.backends.utils import PerLayerParameters
@@ -246,7 +246,7 @@ def run_attention_backend(
246 else:
247 # Build metadata
248 builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
249 - if actual_backend == _Backend.FLEX_ATTENTION:
249 + if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
250 builder.direct_build = use_direct_block_mask
251 attn_metadata = builder.build(
252 common_prefix_len=0,
@@ -289,7 +289,7 @@ def run_attention_backend(
289 def _test_backend_correctness(
290 batch_spec: BatchSpec,
291 model: str,
292 - backend_to_test: list[_Backend | str],
292 + backend_to_test: list[AttentionBackendEnum | str],
293 mask_mod,
294 *,
295 block_size: int = 16,
@@ -455,17 +455,20 @@ def _test_backend_correctness(
455 # Select the appropriate KV cache format for each backend
456 kv_cache_for_backend = kv_cache
457 reset_kv_cache_layout = False
458 - if backend_name in (_Backend.FLASHINFER, _Backend.TRITON_ATTN):
458 + if backend_name in (
459 + AttentionBackendEnum.FLASHINFER,
460 + AttentionBackendEnum.TRITON_ATTN,
461 + ):
462 kv_cache_for_backend = kv_cache.transpose(0, 1)
463
461 - if backend_name == _Backend.FLASHINFER:
464 + if backend_name == AttentionBackendEnum.FLASHINFER:
465 # For FlashInfer default to HND layout and
466 kv_cache_for_backend = (
467 kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
468 )
469 set_kv_cache_layout("HND")
470 reset_kv_cache_layout = True
468 - elif backend_name == _Backend.TRITON_ATTN:
471 + elif backend_name == AttentionBackendEnum.TRITON_ATTN:
472 kv_cache_for_backend = kv_cache_for_backend.contiguous()
473
474 try:
@@ -547,7 +550,9 @@ def test_causal_backend_correctness(
550
551 batch_spec = BATCH_SPECS[batch_spec_name]
552 LARGE_BLOCK_BACKENDS = (
550 - [_Backend.FLEX_ATTENTION] if is_torch_equal_or_newer("2.9.0.dev0") else []
553 + [AttentionBackendEnum.FLEX_ATTENTION]
554 + if is_torch_equal_or_newer("2.9.0.dev0")
555 + else []
556 )
557 SMALL_BLOCK_BACKENDS = [
558 x for x in BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
@@ -573,9 +578,9 @@ def test_causal_backend_correctness(
578
579
580 SLIDING_WINDOW_BACKENDS_TO_TEST = [
576 - _Backend.FLASH_ATTN,
577 - _Backend.FLEX_ATTENTION,
578 - _Backend.TRITON_ATTN,
581 + AttentionBackendEnum.FLASH_ATTN,
582 + AttentionBackendEnum.FLEX_ATTENTION,
583 + AttentionBackendEnum.TRITON_ATTN,
584 "FLEX_ATTENTION_SLOW",
585 ]
586
@@ -612,7 +617,9 @@ def test_sliding_window_backend_correctness(
617 )
618
619 LARGE_BLOCK_BACKENDS = (
615 - [_Backend.FLEX_ATTENTION] if is_torch_equal_or_newer("2.9.0.dev0") else []
620 + [AttentionBackendEnum.FLEX_ATTENTION]
621 + if is_torch_equal_or_newer("2.9.0.dev0")
622 + else []
623 )
624 SMALL_BLOCK_BACKENDS = [
625 x for x in SLIDING_WINDOW_BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
tests/v1/attention/test_mla_backends.py
+13 -16
@@ -18,12 +18,11 @@ from tests.v1.attention.utils import (
18 try_get_attention_backend,
19 )
20 from vllm import _custom_ops as ops
21 -from vllm.attention.backends.registry import _Backend, backend_to_class_str
21 +from vllm.attention.backends.registry import AttentionBackendEnum
22 from vllm.attention.ops.flashmla import is_flashmla_dense_supported
23 from vllm.attention.utils.fa_utils import flash_attn_supports_mla
24 from vllm.config.vllm import set_current_vllm_config
25 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
26 -from vllm.utils.import_utils import resolve_obj_by_qualname
26 from vllm.utils.math_utils import cdiv
27 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
28 from vllm.v1.attention.backends.mla.common import QueryLenSupport
@@ -31,25 +30,25 @@ from vllm.v1.attention.backends.utils import CommonAttentionMetadata
30 from vllm.v1.kv_cache_interface import FullAttentionSpec
31
32 BACKENDS_TO_TEST = [
34 - _Backend.CUTLASS_MLA,
35 - _Backend.FLASHMLA,
36 - _Backend.FLASH_ATTN_MLA,
37 - _Backend.FLASHINFER_MLA,
38 - _Backend.TRITON_MLA,
33 + AttentionBackendEnum.CUTLASS_MLA,
34 + AttentionBackendEnum.FLASHMLA,
35 + AttentionBackendEnum.FLASH_ATTN_MLA,
36 + AttentionBackendEnum.FLASHINFER_MLA,
37 + AttentionBackendEnum.TRITON_MLA,
38 ]
39
40 # Remove sm100 backends from the list if not using sm100
41 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
43 - BACKENDS_TO_TEST.remove(_Backend.CUTLASS_MLA)
44 - BACKENDS_TO_TEST.remove(_Backend.FLASHINFER_MLA)
42 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
43 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER_MLA)
44
45 # Remove FLASH_ATTN_MLA from the list if not supported
46 if not flash_attn_supports_mla():
48 - BACKENDS_TO_TEST.remove(_Backend.FLASH_ATTN_MLA)
47 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASH_ATTN_MLA)
48
49 # Remove FLASHMLA from the list if not supported
50 if not is_flashmla_dense_supported()[0]:
52 - BACKENDS_TO_TEST.remove(_Backend.FLASHMLA)
51 + BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)
52
53 SPEC_DECODE_BACKENDS = []
54 for backend in BACKENDS_TO_TEST:
@@ -62,9 +61,7 @@ for backend in BACKENDS_TO_TEST:
61
62 BACKEND_BLOCK_SIZES = {}
63 for backend in BACKENDS_TO_TEST:
65 - backend_class_str = backend_to_class_str(backend)
66 - backend_class = resolve_obj_by_qualname(backend_class_str)
67 - supported_sizes = backend_class.get_supported_kernel_block_size()
64 + supported_sizes = backend.get_class().supported_kernel_block_sizes
65 if supported_sizes:
66 default_size = supported_sizes[0]
67 block_size = (
@@ -291,7 +288,7 @@ class MockMLAAttentionLayer(AttentionLayerBase):
288
289
290 def run_attention_backend(
294 - backend: _Backend,
291 + backend: AttentionBackendEnum,
292 kv_cache_spec: FullAttentionSpec,
293 layer_names: list[str],
294 vllm_config,
@@ -813,7 +810,7 @@ def test_backend_correctness(
810 # Create a summary for the single-line failure message
811 backend_names = []
812 for f in failures:
816 - if "[_Backend." in f:
813 + if "[AttentionBackendEnum." in f:
814 backend_name = f.split("[")[1].split("]")[0]
815 backend_names.append(backend_name)
816
tests/v1/attention/utils.py
+4 -6
@@ -8,7 +8,7 @@ import pytest
8 import torch
9
10 from vllm.attention.backends.abstract import AttentionImpl
11 -from vllm.attention.backends.registry import _Backend, backend_to_class_str
11 +from vllm.attention.backends.registry import AttentionBackendEnum
12 from vllm.config import (
13 CacheConfig,
14 CompilationConfig,
@@ -20,7 +20,6 @@ from vllm.config import (
20 VllmConfig,
21 )
22 from vllm.config.model import ModelDType
23 -from vllm.utils.import_utils import resolve_obj_by_qualname
23 from vllm.v1.attention.backends.utils import (
24 AttentionMetadataBuilder,
25 CommonAttentionMetadata,
@@ -120,15 +119,14 @@ def create_common_attn_metadata(
119
120
121 def try_get_attention_backend(
123 - backend: _Backend,
122 + backend: AttentionBackendEnum,
123 ) -> tuple[type[AttentionMetadataBuilder], type[AttentionImpl]]:
124 """Try to get the attention backend class, skipping test if not found."""
126 - backend_class_str = backend_to_class_str(backend)
125 try:
128 - backend_class = resolve_obj_by_qualname(backend_class_str)
126 + backend_class = backend.get_class()
127 return backend_class.get_builder_cls(), backend_class.get_impl_cls()
128 except ImportError as e:
131 - pytest.skip(f"{backend_class_str} not available: {e}")
129 + pytest.skip(f"{backend.name} not available: {e}")
130 raise AssertionError("unreachable") from None
131
132
tests/v1/spec_decode/test_eagle.py
+13 -5
@@ -13,7 +13,7 @@ from tests.v1.attention.utils import (
13 create_standard_kv_cache_spec,
14 try_get_attention_backend,
15 )
16 -from vllm.attention.backends.registry import _Backend
16 +from vllm.attention.backends.registry import AttentionBackendEnum
17 from vllm.config import (
18 CacheConfig,
19 DeviceConfig,
@@ -534,11 +534,17 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
534 sampling_metadata = mock.MagicMock()
535
536 if attn_backend == "FLASH_ATTN":
537 - attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.FLASH_ATTN)
537 + attn_metadata_builder_cls, _ = try_get_attention_backend(
538 + AttentionBackendEnum.FLASH_ATTN
539 + )
540 elif attn_backend == "TRITON_ATTN":
539 - attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.TRITON_ATTN)
541 + attn_metadata_builder_cls, _ = try_get_attention_backend(
542 + AttentionBackendEnum.TRITON_ATTN
543 + )
544 elif attn_backend == "TREE_ATTN":
541 - attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.TREE_ATTN)
545 + attn_metadata_builder_cls, _ = try_get_attention_backend(
546 + AttentionBackendEnum.TREE_ATTN
547 + )
548 else:
549 raise ValueError(f"Unsupported attention backend: {attn_backend}")
550
@@ -673,7 +679,9 @@ def test_propose_tree(spec_token_tree):
679 proposer.attn_layer_names = ["layer.0"]
680
681 # Get the tree attention metadata builder.
676 - attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.TREE_ATTN)
682 + attn_metadata_builder_cls, _ = try_get_attention_backend(
683 + AttentionBackendEnum.TREE_ATTN
684 + )
685 attn_metadata_builder = attn_metadata_builder_cls(
686 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
687 layer_names=proposer.attn_layer_names,
tests/v1/spec_decode/test_mtp.py
+4 -2
@@ -12,7 +12,7 @@ from tests.v1.attention.utils import (
12 create_standard_kv_cache_spec,
13 try_get_attention_backend,
14 )
15 -from vllm.attention.backends.registry import _Backend
15 +from vllm.attention.backends.registry import AttentionBackendEnum
16 from vllm.config import (
17 CacheConfig,
18 DeviceConfig,
@@ -177,7 +177,9 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):
177 sampling_metadata = mock.MagicMock()
178
179 # Setup attention metadata
180 - attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.FLASH_ATTN)
180 + attn_metadata_builder_cls, _ = try_get_attention_backend(
181 + AttentionBackendEnum.FLASH_ATTN
182 + )
183
184 attn_metadata_builder = attn_metadata_builder_cls(
185 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
tests/v1/spec_decode/test_tree_attention.py
+4 -4
@@ -10,7 +10,7 @@ from tests.v1.attention.utils import (
10 create_vllm_config,
11 try_get_attention_backend,
12 )
13 -from vllm.attention.backends.registry import _Backend
13 +from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.config import ParallelConfig, SpeculativeConfig
15 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
16
@@ -35,7 +35,7 @@ def forward_attention(
35 block_table: torch.Tensor,
36 slot_mapping: torch.Tensor,
37 seqlen_k: int,
38 - backend: _Backend,
38 + backend: AttentionBackendEnum,
39 spec_token_tree: str | None = None,
40 num_spec_tokens: int = 0,
41 ) -> torch.Tensor:
@@ -241,7 +241,7 @@ def test_tree_attn_correctness() -> None:
241 block_table=block_table,
242 slot_mapping=tree_slot_mapping,
243 seqlen_k=seqlen_k,
244 - backend=_Backend.TREE_ATTN,
244 + backend=AttentionBackendEnum.TREE_ATTN,
245 spec_token_tree=spec_token_tree,
246 num_spec_tokens=tree_size_q - 1,
247 ).view(batch_size, -1, num_heads, dim_per_head)
@@ -278,7 +278,7 @@ def test_tree_attn_correctness() -> None:
278 block_table=block_table,
279 slot_mapping=branch_slot_mapping,
280 seqlen_k=sequence_position + q_len,
281 - backend=_Backend.FLASH_ATTN,
281 + backend=AttentionBackendEnum.FLASH_ATTN,
282 ).view(batch_size, -1, num_heads, dim_per_head)
283
284 # Compare the outputs.
tests/v1/worker/test_gpu_model_runner.py
+15 -10
@@ -185,9 +185,7 @@ def _make_mock_backend_for_kernel_block_size(
185 supported_sizes: list[int | MultipleOf],
186 ):
187 class _MockBackend:
188 - @staticmethod
189 - def get_supported_kernel_block_size():
190 - return supported_sizes
188 + supported_kernel_block_sizes = supported_sizes
189
190 return _MockBackend()
191
@@ -466,13 +464,20 @@ def test_kv_cache_stride_order(monkeypatch, model_runner):
464 # This test checks if GPUModelRunner initializes correctly when an attention
465 # backend enforces a non-default KV cache stride order.
466 n_heads = model_runner.model_config.get_num_kv_heads(model_runner.parallel_config)
469 - expected_kv_cache_shape = [
470 - 2,
471 - NUM_BLOCKS,
472 - BLOCK_SIZE,
473 - n_heads,
474 - model_runner.model_config.get_head_size(),
475 - ]
467 + head_size = model_runner.model_config.get_head_size()
468 +
469 + # Get the expected shape from the backend's get_kv_cache_shape method
470 + # to ensure compatibility with different backends (triton vs flexattention)
471 + attn_backend = None
472 + for attn_group in model_runner._attn_group_iterator():
473 + attn_backend = attn_group.backend
474 + break
475 +
476 + assert attn_backend is not None, "No attention backend found"
477 + expected_kv_cache_shape = list(
478 + attn_backend.get_kv_cache_shape(NUM_BLOCKS, BLOCK_SIZE, n_heads, head_size)
479 + )
480 +
481 # TODO mla test
482 default_stride = tuple(range(5))
483 # Permutation that gets you back to expected kv shape
vllm/attention/backends/abstract.py
+139 -10
@@ -2,13 +2,18 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 -from typing import Generic, Protocol, TypeVar
5 +from typing import TYPE_CHECKING, ClassVar, Generic, Protocol, TypeVar, get_args
6
7 import torch
8
9 from vllm.model_executor.layers.linear import ColumnParallelLinear
10 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
11
12 +if TYPE_CHECKING:
13 + from vllm.config.cache import CacheDType
14 + from vllm.platforms.interface import DeviceCapability
15 + from vllm.v1.attention.backends.utils import KVCacheLayoutType
16 +
17
18 class AttentionType:
19 """
@@ -40,6 +45,9 @@ class AttentionBackend(ABC):
45 # calling the custom op. When piecewise cudagraph is enabled, this
46 # makes sure the output tensor is allocated inside the cudagraph.
47 accept_output_buffer: bool = False
48 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
49 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(1)]
50 + supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = ["auto"]
51
52 @staticmethod
53 @abstractmethod
@@ -51,10 +59,6 @@ class AttentionBackend(ABC):
59 def get_impl_cls() -> type["AttentionImpl"]:
60 raise NotImplementedError
61
54 - @classmethod
55 - def get_supported_kernel_block_size(cls) -> list[int | MultipleOf]:
56 - return cls.get_impl_cls().get_supported_kernel_block_size()
57 -
62 @staticmethod
63 @abstractmethod
64 def get_builder_cls(): # -> Type["AttentionMetadataBuilder"]:
@@ -79,6 +83,136 @@ class AttentionBackend(ABC):
83 def full_cls_name(cls) -> tuple[str, str]:
84 return (cls.__module__, cls.__qualname__)
85
86 + @classmethod
87 + def get_supported_head_sizes(cls) -> list[int]:
88 + return []
89 +
90 + @classmethod
91 + def supports_head_size(cls, head_size: int) -> bool:
92 + supported_head_sizes = cls.get_supported_head_sizes()
93 + return (not supported_head_sizes) or head_size in supported_head_sizes
94 +
95 + @classmethod
96 + def supports_dtype(cls, dtype: torch.dtype) -> bool:
97 + return dtype in cls.supported_dtypes
98 +
99 + @classmethod
100 + def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
101 + if kv_cache_dtype is None:
102 + return True
103 + return (not cls.supported_kv_cache_dtypes) or (
104 + kv_cache_dtype in cls.supported_kv_cache_dtypes
105 + )
106 +
107 + @classmethod
108 + def supports_block_size(cls, block_size: int | None) -> bool:
109 + from vllm.config.cache import BlockSize
110 +
111 + if block_size is None:
112 + return True
113 +
114 + valid_sizes = get_args(BlockSize)
115 + if block_size not in valid_sizes:
116 + return False
117 +
118 + if not cls.supported_kernel_block_sizes:
119 + return True
120 +
121 + for supported_size in cls.supported_kernel_block_sizes:
122 + is_multiple_of = (
123 + isinstance(supported_size, MultipleOf)
124 + and block_size % supported_size.base == 0
125 + )
126 + is_int_equal = (
127 + isinstance(supported_size, int) and block_size == supported_size
128 + )
129 + if is_multiple_of or is_int_equal:
130 + return True
131 + return False
132 +
133 + @classmethod
134 + def is_mla(cls) -> bool:
135 + return False
136 +
137 + @classmethod
138 + def supports_sink(cls) -> bool:
139 + return False
140 +
141 + @classmethod
142 + def is_sparse(cls) -> bool:
143 + return False
144 +
145 + @classmethod
146 + def supports_compute_capability(cls, capability: "DeviceCapability") -> bool:
147 + return True
148 +
149 + @classmethod
150 + def supports_combination(
151 + cls,
152 + head_size: int,
153 + dtype: torch.dtype,
154 + kv_cache_dtype: "CacheDType | None",
155 + block_size: int | None,
156 + use_mla: bool,
157 + has_sink: bool,
158 + use_sparse: bool,
159 + device_capability: "DeviceCapability",
160 + ) -> str | None:
161 + return None
162 +
163 + @classmethod
164 + def validate_configuration(
165 + cls,
166 + head_size: int,
167 + dtype: torch.dtype,
168 + kv_cache_dtype: "CacheDType | None",
169 + block_size: int | None,
170 + use_mla: bool,
171 + has_sink: bool,
172 + use_sparse: bool,
173 + device_capability: "DeviceCapability",
174 + ) -> list[str]:
175 + invalid_reasons = []
176 + if not cls.supports_head_size(head_size):
177 + invalid_reasons.append("head_size not supported")
178 + if not cls.supports_dtype(dtype):
179 + invalid_reasons.append("dtype not supported")
180 + if not cls.supports_kv_cache_dtype(kv_cache_dtype):
181 + invalid_reasons.append("kv_cache_dtype not supported")
182 + if not cls.supports_block_size(block_size):
183 + invalid_reasons.append("block_size not supported")
184 + if use_mla != cls.is_mla():
185 + if use_mla:
186 + invalid_reasons.append("MLA not supported")
187 + else:
188 + invalid_reasons.append("non-MLA not supported")
189 + if has_sink and not cls.supports_sink():
190 + invalid_reasons.append("sink setting not supported")
191 + if use_sparse != cls.is_sparse():
192 + if use_sparse:
193 + invalid_reasons.append("sparse not supported")
194 + else:
195 + invalid_reasons.append("non-sparse not supported")
196 + if not cls.supports_compute_capability(device_capability):
197 + invalid_reasons.append("compute capability not supported")
198 + combination_reason = cls.supports_combination(
199 + head_size,
200 + dtype,
201 + kv_cache_dtype,
202 + block_size,
203 + use_mla,
204 + has_sink,
205 + use_sparse,
206 + device_capability,
207 + )
208 + if combination_reason is not None:
209 + invalid_reasons.append(combination_reason)
210 + return invalid_reasons
211 +
212 + @classmethod
213 + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
214 + return None
215 +
216
217 class AttentionMetadata:
218 pass
@@ -151,11 +285,6 @@ class AttentionImpl(ABC, Generic[T]):
285 ) -> None:
286 raise NotImplementedError
287
154 - @staticmethod
155 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
156 - # TODO: implement this function for all backends.
157 - return [MultipleOf(1)]
158 -
288 @abstractmethod
289 def forward(
290 self,
vllm/attention/backends/registry.py
+168 -84
@@ -3,108 +3,192 @@
3 """Attention backend registry"""
4
5 import enum
6 +from collections.abc import Callable
7 +from typing import TYPE_CHECKING, cast
8
9 +from vllm.logger import init_logger
10 from vllm.utils.import_utils import resolve_obj_by_qualname
11
12 +if TYPE_CHECKING:
13 + from vllm.attention.backends.abstract import AttentionBackend
14
10 -class _Backend(enum.Enum):
11 - FLASH_ATTN = enum.auto()
12 - TRITON_ATTN = enum.auto()
13 - XFORMERS = enum.auto()
14 - ROCM_ATTN = enum.auto()
15 - ROCM_AITER_MLA = enum.auto()
16 - ROCM_AITER_FA = enum.auto() # used for ViT attn backend
17 - TORCH_SDPA = enum.auto()
18 - FLASHINFER = enum.auto()
19 - FLASHINFER_MLA = enum.auto()
20 - TRITON_MLA = enum.auto()
21 - CUTLASS_MLA = enum.auto()
22 - FLASHMLA = enum.auto()
23 - FLASHMLA_SPARSE = enum.auto()
24 - FLASH_ATTN_MLA = enum.auto()
25 - PALLAS = enum.auto()
26 - IPEX = enum.auto()
27 - NO_ATTENTION = enum.auto()
28 - FLEX_ATTENTION = enum.auto()
29 - TREE_ATTN = enum.auto()
30 - ROCM_AITER_UNIFIED_ATTN = enum.auto()
31 -
32 -
33 -BACKEND_MAP = {
34 - _Backend.FLASH_ATTN: "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend", # noqa: E501
35 - _Backend.TRITON_ATTN: "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend", # noqa: E501
36 - _Backend.XFORMERS: "vllm.v1.attention.backends.xformers.XFormersAttentionBackend", # noqa: E501
37 - _Backend.ROCM_ATTN: "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend", # noqa: E501
38 - _Backend.ROCM_AITER_MLA: "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend", # noqa: E501
39 - _Backend.ROCM_AITER_FA: "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend", # noqa: E501
40 - _Backend.TORCH_SDPA: "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend", # noqa: E501
41 - _Backend.FLASHINFER: "vllm.v1.attention.backends.flashinfer.FlashInferBackend", # noqa: E501
42 - _Backend.FLASHINFER_MLA: "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend", # noqa: E501
43 - _Backend.TRITON_MLA: "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend", # noqa: E501
44 - _Backend.CUTLASS_MLA: "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend", # noqa: E501
45 - _Backend.FLASHMLA: "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend", # noqa: E501
46 - _Backend.FLASHMLA_SPARSE: "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend", # noqa: E501
47 - _Backend.FLASH_ATTN_MLA: "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend", # noqa: E501
48 - _Backend.PALLAS: "vllm.v1.attention.backends.pallas.PallasAttentionBackend", # noqa: E501
49 - _Backend.FLEX_ATTENTION: "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend", # noqa: E501
50 - _Backend.TREE_ATTN: "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend", # noqa: E501
51 - _Backend.ROCM_AITER_UNIFIED_ATTN: "vllm.v1.attention.backends.rocm_aiter_unified_attn.RocmAiterUnifiedAttentionBackend", # noqa: E501
52 -}
53 -
54 -
55 -def register_attn_backend(backend: _Backend, class_path: str | None = None):
56 - """
57 - Decorator: register a custom attention backend into BACKEND_MAPPING.
58 - - If class_path is provided, use it.
59 - - Otherwise, auto-generate from the class object.
60 - Validation: only checks if 'backend' is a valid _Backend enum member.
61 - Overwriting existing mappings is allowed. This enables other hardware
62 - platforms to plug in custom out-of-tree backends.
63 - """
64 - if not isinstance(backend, _Backend):
65 - raise ValueError(f"{backend} is not a valid _Backend enum value.")
15 +logger = init_logger(__name__)
16
67 - def decorator(cls):
68 - path = class_path or f"{cls.__module__}.{cls.__qualname__}"
69 - BACKEND_MAP[backend] = path
70 - return cls
17
72 - return decorator
18 +class _AttentionBackendEnumMeta(enum.EnumMeta):
19 + """Metaclass for AttentionBackendEnum to provide better error messages."""
20
21 + def __getitem__(cls, name: str):
22 + """Get backend by name with helpful error messages."""
23 + try:
24 + return super().__getitem__(name)
25 + except KeyError:
26 + members = cast("dict[str, AttentionBackendEnum]", cls.__members__).values()
27 + valid_backends = ", ".join(m.name for m in members)
28 + raise ValueError(
29 + f"Unknown attention backend: '{name}'. "
30 + f"Valid options are: {valid_backends}"
31 + ) from None
32
75 -def backend_to_class_str(backend: _Backend) -> str:
76 - """Get the backend class string
33
78 - Args:
79 - backend: The backend enum value
34 +class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
35 + """Enumeration of all supported attention backends.
36
81 - Returns:
82 - The backend class string
83 - """
84 - return BACKEND_MAP[backend]
37 + The enum value is the default class path, but this can be overridden
38 + at runtime using register_backend().
39
40 + To get the actual backend class (respecting overrides), use:
41 + backend.get_class()
42 + """
43
87 -def backend_to_class(backend: _Backend) -> type:
88 - """Get the backend class.
44 + FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
45 + TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
46 + XFORMERS = "vllm.v1.attention.backends.xformers.XFormersAttentionBackend"
47 + ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
48 + ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
49 + ROCM_AITER_FA = (
50 + "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
51 + )
52 + TORCH_SDPA = "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend"
53 + FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
54 + FLASHINFER_MLA = (
55 + "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
56 + )
57 + TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
58 + CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"
59 + FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend"
60 + FLASHMLA_SPARSE = (
61 + "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
62 + )
63 + FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
64 + PALLAS = "vllm.v1.attention.backends.pallas.PallasAttentionBackend"
65 + IPEX = "vllm.v1.attention.backends.ipex.IpexAttentionBackend"
66 + NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
67 + FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
68 + TREE_ATTN = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend"
69 + ROCM_AITER_UNIFIED_ATTN = (
70 + "vllm.v1.attention.backends.rocm_aiter_unified_attn."
71 + "RocmAiterUnifiedAttentionBackend"
72 + )
73 + # Placeholder for third-party/custom backends - must be registered before use
74 + CUSTOM = ""
75 +
76 + def get_path(self, include_classname: bool = True) -> str:
77 + """Get the class path for this backend (respects overrides).
78 +
79 + Returns:
80 + The fully qualified class path string
81 +
82 + Raises:
83 + ValueError: If Backend.CUSTOM is used without being registered
84 + """
85 + path = _OVERRIDES.get(self, self.value)
86 + if not path:
87 + raise ValueError(
88 + f"Backend {self.name} must be registered before use. "
89 + f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')"
90 + )
91 + if not include_classname:
92 + path = path.rsplit(".", 1)[0]
93 + return path
94 +
95 + def get_class(self) -> "type[AttentionBackend]":
96 + """Get the backend class (respects overrides).
97 +
98 + Returns:
99 + The backend class
100 +
101 + Raises:
102 + ImportError: If the backend class cannot be imported
103 + ValueError: If Backend.CUSTOM is used without being registered
104 + """
105 + return resolve_obj_by_qualname(self.get_path())
106 +
107 + def is_overridden(self) -> bool:
108 + """Check if this backend has been overridden.
109 +
110 + Returns:
111 + True if the backend has a registered override
112 + """
113 + return self in _OVERRIDES
114 +
115 + def clear_override(self) -> None:
116 + """Clear any override for this backend, reverting to the default."""
117 + _OVERRIDES.pop(self, None)
118 +
119 +
120 +_OVERRIDES: dict[AttentionBackendEnum, str] = {}
121 +
122 +
123 +def register_backend(
124 + backend: AttentionBackendEnum, class_path: str | None = None
125 +) -> Callable[[type], type]:
126 + """Register or override a backend implementation.
127
128 Args:
91 - backend: The backend enum value
129 + backend: The AttentionBackendEnum member to register
130 + class_path: Optional class path. If not provided and used as
131 + decorator, will be auto-generated from the class.
132
133 Returns:
94 - The backend class
134 + Decorator function if class_path is None, otherwise a no-op
135 +
136 + Examples:
137 + # Override an existing backend
138 + @register_backend(AttentionBackendEnum.FLASH_ATTN)
139 + class MyCustomFlashAttn:
140 + ...
141 +
142 + # Register a custom third-party backend
143 + @register_backend(AttentionBackendEnum.CUSTOM)
144 + class MyCustomBackend:
145 + ...
146 +
147 + # Direct registration
148 + register_backend(
149 + AttentionBackendEnum.CUSTOM,
150 + "my.module.MyCustomBackend"
151 + )
152 """
96 - backend_class_name = backend_to_class_str(backend)
97 - return resolve_obj_by_qualname(backend_class_name)
153
154 + def decorator(cls: type) -> type:
155 + _OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
156 + return cls
157
100 -def backend_name_to_enum(backend_name: str) -> _Backend | None:
101 - """
102 - Convert a string backend name to a _Backend enum value.
158 + if class_path is not None:
159 + _OVERRIDES[backend] = class_path
160 + return lambda x: x
161
104 - Returns:
105 - _Backend: enum value if backend_name is a valid in-tree type
106 - None: otherwise it's an invalid in-tree type or an out-of-tree platform
107 - is loaded.
162 + return decorator
163 +
164 +
165 +# Backwards compatibility alias for plugins
166 +class _BackendMeta(type):
167 + """Metaclass to provide deprecation warnings when accessing _Backend."""
168 +
169 + def __getattribute__(cls, name: str):
170 + if name not in ("__class__", "__mro__", "__name__"):
171 + logger.warning(
172 + "_Backend has been renamed to AttentionBackendEnum. "
173 + "Please update your code to use AttentionBackendEnum instead. "
174 + "_Backend will be removed in a future release."
175 + )
176 + return getattr(AttentionBackendEnum, name)
177 +
178 + def __getitem__(cls, name: str):
179 + logger.warning(
180 + "_Backend has been renamed to AttentionBackendEnum. "
181 + "Please update your code to use AttentionBackendEnum instead. "
182 + "_Backend will be removed in a future release."
183 + )
184 + return AttentionBackendEnum[name]
185 +
186 +
187 +class _Backend(metaclass=_BackendMeta):
188 + """Deprecated: Use AttentionBackendEnum instead.
189 +
190 + This class is provided for backwards compatibility with plugins
191 + and will be removed in a future release.
192 """
109 - assert backend_name is not None
110 - return _Backend[backend_name] if backend_name in _Backend.__members__ else None
193 +
194 + pass
vllm/attention/layer.py
+39 -29
@@ -12,7 +12,7 @@ import torch.nn.functional as F
12 import vllm.envs as envs
13 from vllm.attention import AttentionType
14 from vllm.attention.backends.abstract import AttentionBackend, MLAAttentionImpl
15 -from vllm.attention.backends.registry import _Backend, backend_name_to_enum
15 +from vllm.attention.backends.registry import AttentionBackendEnum
16 from vllm.attention.selector import get_attn_backend
17 from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
18 from vllm.config import CacheConfig, get_current_vllm_config
@@ -99,40 +99,44 @@ def check_upstream_fa_availability(dtype: torch.dtype):
99
100
101 def maybe_get_vit_flash_attn_backend(
102 - attn_backend: _Backend,
102 + attn_backend: AttentionBackendEnum,
103 use_upstream_fa: bool,
104 - attn_backend_override: _Backend | None = None,
105 -) -> tuple[_Backend, Callable | None]:
104 + attn_backend_override: AttentionBackendEnum | None = None,
105 +) -> tuple[AttentionBackendEnum, Callable | None]:
106 if current_platform.is_rocm():
107 if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
108 - attn_backend = _Backend.ROCM_AITER_FA
108 + attn_backend = AttentionBackendEnum.ROCM_AITER_FA
109
110 elif (
111 check_upstream_fa_availability(torch.get_default_dtype())
112 and on_gfx9()
113 and attn_backend_override is None
114 ):
115 - attn_backend = _Backend.FLASH_ATTN
115 + attn_backend = AttentionBackendEnum.FLASH_ATTN
116 use_upstream_fa = True
117 else:
118 - return _Backend.TORCH_SDPA, None
118 + return AttentionBackendEnum.TORCH_SDPA, None
119
120 elif current_platform.is_cuda():
121 - if attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
122 - torch.get_default_dtype()
121 + if (
122 + attn_backend != AttentionBackendEnum.FLASH_ATTN
123 + and check_upstream_fa_availability(torch.get_default_dtype())
124 ):
124 - attn_backend = _Backend.FLASH_ATTN
125 + attn_backend = AttentionBackendEnum.FLASH_ATTN
126 use_upstream_fa = True
127 elif current_platform.is_xpu():
127 - assert attn_backend == _Backend.FLASH_ATTN, (
128 + assert attn_backend == AttentionBackendEnum.FLASH_ATTN, (
129 "XPU platform only supports FLASH_ATTN as vision attention backend."
130 )
131 use_upstream_fa = False
132 else:
132 - return _Backend.TORCH_SDPA, None
133 + return AttentionBackendEnum.TORCH_SDPA, None
134
134 - if attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
135 - if attn_backend == _Backend.ROCM_AITER_FA:
135 + if attn_backend in {
136 + AttentionBackendEnum.FLASH_ATTN,
137 + AttentionBackendEnum.ROCM_AITER_FA,
138 + }:
139 + if attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
140 from aiter import flash_attn_varlen_func
141 else:
142 if use_upstream_fa:
@@ -309,7 +313,7 @@ class Attention(nn.Module, AttentionLayerBase):
313 kv_sharing_target_layer_name,
314 **extra_impl_args,
315 )
312 - self.backend = backend_name_to_enum(self.attn_backend.get_name())
316 + self.backend = AttentionBackendEnum[self.attn_backend.get_name()]
317 self.dtype = dtype
318
319 # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
@@ -530,13 +534,13 @@ class MultiHeadAttention(nn.Module):
534 backend
535 if backend
536 in {
533 - _Backend.TORCH_SDPA,
534 - _Backend.XFORMERS,
535 - _Backend.PALLAS,
536 - _Backend.ROCM_AITER_FA,
537 - _Backend.FLASH_ATTN,
537 + AttentionBackendEnum.TORCH_SDPA,
538 + AttentionBackendEnum.XFORMERS,
539 + AttentionBackendEnum.PALLAS,
540 + AttentionBackendEnum.ROCM_AITER_FA,
541 + AttentionBackendEnum.FLASH_ATTN,
542 }
539 - else _Backend.TORCH_SDPA
543 + else AttentionBackendEnum.TORCH_SDPA
544 )
545
546 self.attn_backend, self._flash_attn_varlen_func = (
@@ -547,17 +551,23 @@ class MultiHeadAttention(nn.Module):
551 )
552 )
553
550 - if self.attn_backend == _Backend.XFORMERS and not check_xformers_availability():
551 - self.attn_backend = _Backend.TORCH_SDPA
554 + if (
555 + self.attn_backend == AttentionBackendEnum.XFORMERS
556 + and not check_xformers_availability()
557 + ):
558 + self.attn_backend = AttentionBackendEnum.TORCH_SDPA
559
560 self.is_flash_attn_backend = self.attn_backend in {
554 - _Backend.FLASH_ATTN,
555 - _Backend.ROCM_AITER_FA,
561 + AttentionBackendEnum.FLASH_ATTN,
562 + AttentionBackendEnum.ROCM_AITER_FA,
563 }
564
565 # this condition is just to make sure that the
566 # use_upstream_fa in the log is correct
560 - if current_platform.is_rocm() and self.attn_backend == _Backend.FLASH_ATTN:
567 + if (
568 + current_platform.is_rocm()
569 + and self.attn_backend == AttentionBackendEnum.FLASH_ATTN
570 + ):
571 use_upstream_fa = True
572
573 logger.info_once(
@@ -606,17 +616,17 @@ class MultiHeadAttention(nn.Module):
616 max_seqlen_k=kv_len,
617 softmax_scale=self.scale,
618 )
609 - elif self.attn_backend == _Backend.XFORMERS:
619 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
620 from xformers import ops as xops
621
622 out = xops.memory_efficient_attention_forward(
623 query, key, value, scale=self.scale
624 )
615 - elif self.attn_backend == _Backend.TORCH_SDPA:
625 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
626 query, key, value = (x.transpose(1, 2) for x in (query, key, value))
627 out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)
628 out = out.transpose(1, 2)
619 - elif self.attn_backend == _Backend.PALLAS:
629 + elif self.attn_backend == AttentionBackendEnum.PALLAS:
630 query, key, value = (x.transpose(1, 2) for x in (query, key, value))
631 from torch_xla.experimental.custom_kernel import flash_attention
632
vllm/attention/selector.py
+46 -78
@@ -4,14 +4,15 @@
4 import os
5 from collections.abc import Generator
6 from contextlib import contextmanager
7 -from dataclasses import dataclass
7 from functools import cache
8 +from typing import cast, get_args
9
10 import torch
11
12 import vllm.envs as envs
13 from vllm.attention.backends.abstract import AttentionBackend
14 -from vllm.attention.backends.registry import _Backend, backend_name_to_enum
14 +from vllm.attention.backends.registry import AttentionBackendEnum
15 +from vllm.config.cache import CacheDType
16 from vllm.logger import init_logger
17 from vllm.utils import STR_BACKEND_ENV_VAR
18 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -19,18 +20,18 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
20 logger = init_logger(__name__)
21
22
22 -def get_env_variable_attn_backend() -> _Backend | None:
23 +def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
24 """
25 Get the backend override specified by the vLLM attention
26 backend environment variable, if one is specified.
27
28 Returns:
29
29 - * _Backend enum value if an override is specified
30 + * AttentionBackendEnum value if an override is specified
31 * None otherwise
32 """
33 backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
33 - return None if backend_name is None else backend_name_to_enum(backend_name)
34 + return None if backend_name is None else AttentionBackendEnum[backend_name]
35
36
37 # Global state allows a particular choice of backend
@@ -40,10 +41,10 @@ def get_env_variable_attn_backend() -> _Backend | None:
41 #
42 # THIS SELECTION TAKES PRECEDENCE OVER THE
43 # VLLM_ATTENTION_BACKEND ENVIRONMENT VARIABLE
43 -forced_attn_backend: _Backend | None = None
44 +forced_attn_backend: AttentionBackendEnum | None = None
45
46
46 -def global_force_attn_backend(attn_backend: _Backend | None) -> None:
47 +def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None:
48 """
49 Force all attention operations to use a specified backend.
50
@@ -58,7 +59,7 @@ def global_force_attn_backend(attn_backend: _Backend | None) -> None:
59 forced_attn_backend = attn_backend
60
61
61 -def get_global_forced_attn_backend() -> _Backend | None:
62 +def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
63 """
64 Get the currently-forced choice of attention backend,
65 or None if auto-selection is currently enabled.
@@ -66,78 +67,28 @@ def get_global_forced_attn_backend() -> _Backend | None:
67 return forced_attn_backend
68
69
69 -@dataclass(frozen=True)
70 -class _IsSupported:
71 - can_import: bool
72 - head_size: bool
73 - dtype: bool
74 -
75 - def __bool__(self) -> bool:
76 - return self.can_import and self.head_size and self.dtype
77 -
78 -
79 -def is_attn_backend_supported(
80 - attn_backend: str | type[AttentionBackend],
81 - head_size: int,
82 - dtype: torch.dtype,
83 - *,
84 - allow_import_error: bool = True,
85 -) -> _IsSupported:
86 - if isinstance(attn_backend, str):
87 - try:
88 - attn_backend = resolve_obj_by_qualname(attn_backend)
89 - except ImportError:
90 - if not allow_import_error:
91 - raise
92 -
93 - return _IsSupported(can_import=False, head_size=False, dtype=False)
94 -
95 - assert isinstance(attn_backend, type)
96 -
97 - # TODO: Update the interface once V0 is removed
98 - if get_supported_head_sizes := getattr(
99 - attn_backend, "get_supported_head_sizes", None
100 - ):
101 - is_head_size_supported = head_size in get_supported_head_sizes()
102 - elif validate_head_size := getattr(attn_backend, "validate_head_size", None):
103 - try:
104 - validate_head_size(head_size)
105 - is_head_size_supported = True
106 - except Exception:
107 - is_head_size_supported = False
108 - else:
109 - raise NotImplementedError(
110 - f"{attn_backend.__name__} does not support head size validation"
111 - )
112 -
113 - if get_supported_dtypes := getattr(attn_backend, "get_supported_dtypes", None):
114 - is_dtype_supported = dtype in get_supported_dtypes()
115 - else:
116 - raise NotImplementedError(
117 - f"{attn_backend.__name__} does not support dtype validation"
118 - )
119 -
120 - return _IsSupported(
121 - can_import=True,
122 - head_size=is_head_size_supported,
123 - dtype=is_dtype_supported,
124 - )
125 -
126 -
70 def get_attn_backend(
71 head_size: int,
72 dtype: torch.dtype,
73 kv_cache_dtype: str | None,
131 - block_size: int,
74 + block_size: int | None,
75 use_mla: bool = False,
76 has_sink: bool = False,
77 use_sparse: bool = False,
78 ) -> type[AttentionBackend]:
79 """Selects which attention backend to use and lazily imports it."""
80 +
81 + if kv_cache_dtype is not None:
82 + valid_cache_dtypes = get_args(CacheDType)
83 + assert kv_cache_dtype in valid_cache_dtypes, (
84 + f"Invalid kv_cache_dtype: {kv_cache_dtype}. "
85 + f"Valid values are: {valid_cache_dtypes}"
86 + )
87 +
88 return _cached_get_attn_backend(
89 head_size=head_size,
90 dtype=dtype,
140 - kv_cache_dtype=kv_cache_dtype,
91 + kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype),
92 block_size=block_size,
93 use_mla=use_mla,
94 has_sink=has_sink,
@@ -149,8 +100,8 @@ def get_attn_backend(
100 def _cached_get_attn_backend(
101 head_size: int,
102 dtype: torch.dtype,
152 - kv_cache_dtype: str | None,
153 - block_size: int,
103 + kv_cache_dtype: CacheDType | None,
104 + block_size: int | None,
105 use_mla: bool = False,
106 has_sink: bool = False,
107 use_sparse: bool = False,
@@ -161,7 +112,9 @@ def _cached_get_attn_backend(
112 # THIS SELECTION OVERRIDES THE VLLM_ATTENTION_BACKEND
113 # ENVIRONMENT VARIABLE.
114 selected_backend = None
164 - backend_by_global_setting: _Backend | None = get_global_forced_attn_backend()
115 + backend_by_global_setting: AttentionBackendEnum | None = (
116 + get_global_forced_attn_backend()
117 + )
118 if backend_by_global_setting is not None:
119 selected_backend = backend_by_global_setting
120 else:
@@ -177,12 +130,13 @@ def _cached_get_attn_backend(
130 STR_BACKEND_ENV_VAR,
131 )
132 backend_by_env_var = backend_by_env_var.removesuffix("_VLLM_V1")
180 - selected_backend = backend_name_to_enum(backend_by_env_var)
181 - if selected_backend is None:
133 + try:
134 + selected_backend = AttentionBackendEnum[backend_by_env_var]
135 + except KeyError as e:
136 raise ValueError(
183 - f"Invalid attention backend: '{backend_by_env_var}'. "
184 - f"Valid backends are: {list(_Backend.__members__.keys())}"
185 - )
137 + f"Invalid attention backend: '{backend_by_env_var}'. Valid "
138 + f"backends are: {list(AttentionBackendEnum.__members__.keys())}"
139 + ) from e
140
141 # get device-specific attn_backend
142 from vllm.platforms import current_platform
@@ -202,12 +156,26 @@ def _cached_get_attn_backend(
156 raise ValueError(
157 f"Invalid attention backend for {current_platform.device_name}"
158 )
205 - return resolve_obj_by_qualname(attention_cls)
159 + backend = resolve_obj_by_qualname(attention_cls)
160 +
161 + # Adjust kv cache layout if the selected backend requires a specific one
162 + required_layout = backend.get_required_kv_cache_layout()
163 + if required_layout is not None:
164 + from vllm.v1.attention.backends.utils import set_kv_cache_layout
165 +
166 + set_kv_cache_layout(required_layout)
167 + logger.info(
168 + "Using %s KV cache layout for %s backend.",
169 + required_layout,
170 + backend.get_name(),
171 + )
172 +
173 + return backend
174
175
176 @contextmanager
177 def global_force_attn_backend_context_manager(
210 - attn_backend: _Backend,
178 + attn_backend: AttentionBackendEnum,
179 ) -> Generator[None, None, None]:
180 """
181 Globally force a vLLM attention backend override within a
vllm/config/cache.py
+9 -1
@@ -21,7 +21,15 @@ else:
21 logger = init_logger(__name__)
22
23 BlockSize = Literal[1, 8, 16, 32, 64, 128, 256]
24 -CacheDType = Literal["auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2", "fp8_inc"]
24 +CacheDType = Literal[
25 + "auto",
26 + "bfloat16",
27 + "fp8",
28 + "fp8_e4m3",
29 + "fp8_e5m2",
30 + "fp8_inc",
31 + "fp8_ds_mla",
32 +]
33 MambaDType = Literal["auto", "float32"]
34 PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor"]
35 KVOffloadingBackend = Literal["native", "lmcache"]
vllm/config/model.py
+4 -4
@@ -45,7 +45,7 @@ if TYPE_CHECKING:
45
46 import vllm.model_executor.layers.quantization as me_quant
47 import vllm.model_executor.models as me_models
48 - from vllm.attention.backends.registry import _Backend
48 + from vllm.attention.backends.registry import AttentionBackendEnum
49 from vllm.config.load import LoadConfig
50 from vllm.config.parallel import ParallelConfig
51 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -53,7 +53,7 @@ if TYPE_CHECKING:
53 else:
54 PretrainedConfig = Any
55
56 - _Backend = Any
56 + AttentionBackendEnum = Any
57 me_quant = LazyLoader(
58 "model_executor", globals(), "vllm.model_executor.layers.quantization"
59 )
@@ -302,7 +302,7 @@ class ModelConfig:
302 mm_processor_cache_type: InitVar[MMCacheType | None] = None
303 mm_shm_cache_max_object_size_mb: InitVar[int | None] = None
304 mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None
305 - mm_encoder_attn_backend: InitVar[_Backend | str | None] = None
305 + mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = None
306 interleave_mm_strings: InitVar[bool | None] = None
307 skip_mm_profiling: InitVar[bool | None] = None
308 video_pruning_rate: InitVar[float | None] = None
@@ -420,7 +420,7 @@ class ModelConfig:
420 mm_processor_cache_type: MMCacheType | None,
421 mm_shm_cache_max_object_size_mb: int | None,
422 mm_encoder_tp_mode: MMEncoderTPMode | None,
423 - mm_encoder_attn_backend: _Backend | str | None,
423 + mm_encoder_attn_backend: AttentionBackendEnum | str | None,
424 interleave_mm_strings: bool | None,
425 skip_mm_profiling: bool | None,
426 video_pruning_rate: float | None,
vllm/config/multimodal.py
+11 -21
@@ -11,9 +11,9 @@ from pydantic.dataclasses import dataclass
11 from vllm.config.utils import config
12
13 if TYPE_CHECKING:
14 - from vllm.attention.backends.registry import _Backend
14 + from vllm.attention.backends.registry import AttentionBackendEnum
15 else:
16 - _Backend = Any
16 + AttentionBackendEnum = Any
17
18
19 @dataclass
@@ -125,10 +125,10 @@ class MultiModalConfig:
125 DP (which is controlled by `--data-parallel-size`).
126 This is only supported on a per-model basis and falls back to
127 `"weights"` if the encoder does not support DP."""
128 - mm_encoder_attn_backend: _Backend | None = None
128 + mm_encoder_attn_backend: AttentionBackendEnum | None = None
129 """Optional override for the multi-modal encoder attention backend when
130 using vision transformers. Accepts any value from
131 - `vllm.attention.backends.registry._Backend` (e.g. `FLASH_ATTN`)."""
131 + `vllm.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
132 interleave_mm_strings: bool = False
133 """Enable fully interleaved support for multimodal prompts, while using
134 --chat-template-content-format=string."""
@@ -167,26 +167,16 @@ class MultiModalConfig:
167
168 @field_validator("mm_encoder_attn_backend", mode="before")
169 @classmethod
170 - def _validate_mm_encoder_attn_backend(cls, value: object) -> _Backend | None:
171 - from vllm.attention.backends.registry import (
172 - _Backend as BackendEnum,
173 - )
174 - from vllm.attention.backends.registry import (
175 - backend_name_to_enum,
176 - )
177 -
178 - if value is None or isinstance(value, BackendEnum):
170 + def _validate_mm_encoder_attn_backend(
171 + cls, value: str | AttentionBackendEnum | None
172 + ) -> AttentionBackendEnum | None:
173 + if value is None or isinstance(value, AttentionBackendEnum):
174 return value
175
181 - if isinstance(value, str):
182 - candidate = backend_name_to_enum(value.upper())
183 - if candidate is not None:
184 - return candidate
185 -
186 - valid_backends = ", ".join(sorted(BackendEnum.__members__.keys()))
187 - raise ValueError(
188 - f"Invalid mm encoder attention backend. Expected one of: {valid_backends}."
176 + assert isinstance(value, str), (
177 + "mm_encoder_attn_backend must be a string or an AttentionBackendEnum."
178 )
179 + return AttentionBackendEnum[value.upper()]
180
181 @model_validator(mode="after")
182 def _validate_multimodal_config(self):
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+4 -4
@@ -21,7 +21,7 @@ import torch
21 import zmq
22
23 from vllm import envs
24 -from vllm.attention.backends.registry import _Backend, backend_name_to_enum
24 +from vllm.attention.backends.registry import AttentionBackendEnum
25 from vllm.attention.selector import get_attn_backend
26 from vllm.config import VllmConfig
27 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -876,9 +876,9 @@ class NixlConnectorWorker:
876 use_mla=self.use_mla,
877 )
878 self.backend_name = backend.get_name()
879 - attn_backend = backend_name_to_enum(self.backend_name)
880 - self._use_flashinfer = attn_backend == _Backend.FLASHINFER
881 - self._use_pallas = attn_backend == _Backend.PALLAS
879 + attn_backend = AttentionBackendEnum[self.backend_name]
880 + self._use_flashinfer = attn_backend == AttentionBackendEnum.FLASHINFER
881 + self._use_pallas = attn_backend == AttentionBackendEnum.PALLAS
882 self.kv_cache_layout = get_kv_cache_layout()
883 self.host_buffer_kv_cache_layout = self.kv_cache_layout
884 logger.debug("Detected attention backend %s", self.backend_name)
vllm/engine/arg_utils.py
+2 -2
@@ -32,7 +32,7 @@ from pydantic.fields import FieldInfo
32 from typing_extensions import TypeIs, deprecated
33
34 import vllm.envs as envs
35 -from vllm.attention.backends.registry import _Backend
35 +from vllm.attention.backends.registry import AttentionBackendEnum
36 from vllm.config import (
37 CacheConfig,
38 CompilationConfig,
@@ -462,7 +462,7 @@ class EngineArgs:
462 MultiModalConfig.mm_shm_cache_max_object_size_mb
463 )
464 mm_encoder_tp_mode: MMEncoderTPMode = MultiModalConfig.mm_encoder_tp_mode
465 - mm_encoder_attn_backend: _Backend | str | None = (
465 + mm_encoder_attn_backend: AttentionBackendEnum | str | None = (
466 MultiModalConfig.mm_encoder_attn_backend
467 )
468 io_processor_plugin: str | None = None
vllm/envs.py
+3 -3
@@ -626,14 +626,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
626 # - "FLASH_ATTN_MLA": use FlashAttention for MLA
627 # - "FLASHINFER_MLA": use FlashInfer for MLA
628 # - "CUTLASS_MLA": use CUTLASS for MLA
629 - # All possible options loaded dynamically from _Backend enum
629 + # All possible options loaded dynamically from AttentionBackendEnum
630 "VLLM_ATTENTION_BACKEND": env_with_choices(
631 "VLLM_ATTENTION_BACKEND",
632 None,
633 lambda: list(
634 __import__(
635 - "vllm.attention.backends.registry", fromlist=["_Backend"]
636 - )._Backend.__members__.keys()
635 + "vllm.attention.backends.registry", fromlist=["AttentionBackendEnum"]
636 + ).AttentionBackendEnum.__members__.keys()
637 ),
638 ),
639 # If set, vllm will use flashinfer sampler
vllm/model_executor/models/dots_ocr.py
+19 -18
@@ -9,7 +9,7 @@ import torch.nn.functional as F
9 from torch.nn import LayerNorm
10 from transformers.models.qwen2_vl import Qwen2VLProcessor
11
12 -from vllm.attention.backends.registry import _Backend
12 +from vllm.attention.backends.registry import AttentionBackendEnum
13 from vllm.attention.layer import (
14 check_upstream_fa_availability,
15 maybe_get_vit_flash_attn_backend,
@@ -256,7 +256,7 @@ class DotsVisionAttention(nn.Module):
256 quant_config: QuantizationConfig | None = None,
257 prefix: str = "",
258 use_data_parallel: bool = False,
259 - attn_backend_override: _Backend | None = None,
259 + attn_backend_override: AttentionBackendEnum | None = None,
260 ) -> None:
261 super().__init__()
262
@@ -303,17 +303,17 @@ class DotsVisionAttention(nn.Module):
303 )
304 )
305 if self.attn_backend not in {
306 - _Backend.FLASH_ATTN,
307 - _Backend.TORCH_SDPA,
308 - _Backend.XFORMERS,
309 - _Backend.ROCM_AITER_FA,
306 + AttentionBackendEnum.FLASH_ATTN,
307 + AttentionBackendEnum.TORCH_SDPA,
308 + AttentionBackendEnum.XFORMERS,
309 + AttentionBackendEnum.ROCM_AITER_FA,
310 }:
311 raise RuntimeError(
312 f"Unsupported vision attention backend: {self.attn_backend}"
313 )
314 self.is_flash_attn_backend = self.attn_backend in {
315 - _Backend.FLASH_ATTN,
316 - _Backend.ROCM_AITER_FA,
315 + AttentionBackendEnum.FLASH_ATTN,
316 + AttentionBackendEnum.ROCM_AITER_FA,
317 }
318
319 def forward(
@@ -361,7 +361,7 @@ class DotsVisionAttention(nn.Module):
361 self.num_attention_heads_per_partition,
362 self.hidden_size_per_attention_head,
363 )
364 - elif self.attn_backend == _Backend.TORCH_SDPA:
364 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
365 outputs = []
366 for i in range(1, len(cu_seqlens)):
367 s = int(cu_seqlens[i - 1])
@@ -373,7 +373,7 @@ class DotsVisionAttention(nn.Module):
373 out_i = out_i.permute(0, 2, 1, 3)
374 outputs.append(out_i)
375 context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
376 - elif self.attn_backend == _Backend.XFORMERS:
376 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
377 from xformers import ops as xops
378 from xformers.ops.fmha.attn_bias import BlockDiagonalMask
379
@@ -514,7 +514,7 @@ class DotsVisionBlock(nn.Module):
514 quant_config: QuantizationConfig | None = None,
515 prefix: str = "",
516 use_data_parallel: bool = False,
517 - attn_backend_override: _Backend | None = None,
517 + attn_backend_override: AttentionBackendEnum | None = None,
518 ):
519 super().__init__()
520
@@ -567,7 +567,7 @@ class DotsVisionTransformer(nn.Module):
567 require_post_norm: bool | None = None,
568 prefix: str = "",
569 use_data_parallel: bool = False,
570 - attn_backend_override: _Backend | None = None,
570 + attn_backend_override: AttentionBackendEnum | None = None,
571 ) -> None:
572 super().__init__()
573 self.config = config
@@ -582,10 +582,11 @@ class DotsVisionTransformer(nn.Module):
582 dtype=torch.get_default_dtype(),
583 attn_backend_override=attn_backend_override,
584 )
585 - if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
586 - torch.get_default_dtype()
585 + if (
586 + self.attn_backend != AttentionBackendEnum.FLASH_ATTN
587 + and check_upstream_fa_availability(torch.get_default_dtype())
588 ):
588 - self.attn_backend = _Backend.FLASH_ATTN
589 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
590 self.out_hidden_size = config.hidden_size
591 # Keep blocks for compatibility with other vision towers
592 num_layers = (
@@ -666,11 +667,11 @@ class DotsVisionTransformer(nn.Module):
667 ) -> tuple[int | None, list[int] | None]:
668 max_seqlen, seqlens = None, None
669 if (
669 - self.attn_backend == _Backend.FLASH_ATTN
670 - or self.attn_backend == _Backend.ROCM_AITER_FA
670 + self.attn_backend == AttentionBackendEnum.FLASH_ATTN
671 + or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
672 ):
673 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
673 - elif self.attn_backend == _Backend.XFORMERS:
674 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
675 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
676 return max_seqlen, seqlens
677
vllm/model_executor/models/ernie45_vl.py
+19 -18
@@ -36,7 +36,7 @@ import torch.nn.functional as F
36 from einops import rearrange, repeat
37 from transformers import BatchFeature, PretrainedConfig
38
39 -from vllm.attention.backends.registry import _Backend
39 +from vllm.attention.backends.registry import AttentionBackendEnum
40 from vllm.attention.layer import (
41 check_upstream_fa_availability,
42 maybe_get_vit_flash_attn_backend,
@@ -164,7 +164,7 @@ class Ernie4_5_VisionAttention(nn.Module):
164 projection_size: int,
165 quant_config: QuantizationConfig | None = None,
166 prefix: str = "",
167 - attn_backend_override: _Backend | None = None,
167 + attn_backend_override: AttentionBackendEnum | None = None,
168 ) -> None:
169 super().__init__()
170 # Per attention head and per partition values.
@@ -211,17 +211,17 @@ class Ernie4_5_VisionAttention(nn.Module):
211 )
212
213 if self.attn_backend not in {
214 - _Backend.FLASH_ATTN,
215 - _Backend.TORCH_SDPA,
216 - _Backend.XFORMERS,
217 - _Backend.ROCM_AITER_FA,
214 + AttentionBackendEnum.FLASH_ATTN,
215 + AttentionBackendEnum.TORCH_SDPA,
216 + AttentionBackendEnum.XFORMERS,
217 + AttentionBackendEnum.ROCM_AITER_FA,
218 }:
219 raise RuntimeError(
220 f"Ernie45-VL does not support {self.attn_backend} backend now."
221 )
222 self.is_flash_attn_backend = self.attn_backend in {
223 - _Backend.FLASH_ATTN,
224 - _Backend.ROCM_AITER_FA,
223 + AttentionBackendEnum.FLASH_ATTN,
224 + AttentionBackendEnum.ROCM_AITER_FA,
225 }
226
227 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
@@ -291,7 +291,7 @@ class Ernie4_5_VisionAttention(nn.Module):
291 context_layer = rearrange(
292 output, "(b s) h d -> s b (h d)", b=batch_size
293 ).contiguous()
294 - elif self.attn_backend == _Backend.TORCH_SDPA:
294 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
295 # Execute attention entry by entry for speed & less VRAM.
296 outputs = []
297 for i in range(1, len(cu_seqlens)):
@@ -310,7 +310,7 @@ class Ernie4_5_VisionAttention(nn.Module):
310 context_layer = rearrange(
311 context_layer, "b s h d -> s b (h d)"
312 ).contiguous()
313 - elif self.attn_backend == _Backend.XFORMERS:
313 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
314 from xformers import ops as xops
315 from xformers.ops.fmha.attn_bias import BlockDiagonalMask
316
@@ -370,7 +370,7 @@ class Ernie4_5_VisionBlock(nn.Module):
370 norm_layer: Callable[[int], nn.Module] | None = None,
371 quant_config: QuantizationConfig | None = None,
372 prefix: str = "",
373 - attn_backend_override: _Backend | None = None,
373 + attn_backend_override: AttentionBackendEnum | None = None,
374 ) -> None:
375 super().__init__()
376
@@ -463,7 +463,7 @@ class Ernie4_5_VisionTransformer(nn.Module):
463 norm_eps: float = 1e-6,
464 quant_config: QuantizationConfig | None = None,
465 prefix: str = "",
466 - attn_backend_override: _Backend | None = None,
466 + attn_backend_override: AttentionBackendEnum | None = None,
467 ) -> None:
468 super().__init__()
469 patch_size = vision_config.patch_size
@@ -515,10 +515,11 @@ class Ernie4_5_VisionTransformer(nn.Module):
515 dtype=torch.get_default_dtype(),
516 attn_backend_override=attn_backend_override,
517 )
518 - if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
519 - torch.get_default_dtype()
518 + if (
519 + self.attn_backend != AttentionBackendEnum.FLASH_ATTN
520 + and check_upstream_fa_availability(torch.get_default_dtype())
521 ):
521 - self.attn_backend = _Backend.FLASH_ATTN
522 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
523
524 @property
525 def dtype(self) -> torch.dtype:
@@ -565,11 +566,11 @@ class Ernie4_5_VisionTransformer(nn.Module):
566 ) -> tuple[int | None, list[int] | None]:
567 max_seqlen, seqlens = None, None
568 if (
568 - self.attn_backend == _Backend.FLASH_ATTN
569 - or self.attn_backend == _Backend.ROCM_AITER_FA
569 + self.attn_backend == AttentionBackendEnum.FLASH_ATTN
570 + or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
571 ):
572 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
572 - elif self.attn_backend == _Backend.XFORMERS:
573 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
574 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
575 return max_seqlen, seqlens
576
vllm/model_executor/models/glm4_1v.py
+18 -17
@@ -46,7 +46,7 @@ from transformers.models.glm4v.image_processing_glm4v import (
46 from transformers.models.glm4v.video_processing_glm4v import Glm4vVideoProcessor
47 from transformers.video_utils import VideoMetadata
48
49 -from vllm.attention.backends.registry import _Backend
49 +from vllm.attention.backends.registry import AttentionBackendEnum
50 from vllm.attention.layer import (
51 check_upstream_fa_availability,
52 maybe_get_vit_flash_attn_backend,
@@ -252,7 +252,7 @@ class Glm4vVisionAttention(nn.Module):
252 quant_config: QuantizationConfig | None = None,
253 prefix: str = "",
254 use_data_parallel: bool = False,
255 - attn_backend_override: _Backend | None = None,
255 + attn_backend_override: AttentionBackendEnum | None = None,
256 ) -> None:
257 super().__init__()
258 # Per attention head and per partition values.
@@ -306,18 +306,18 @@ class Glm4vVisionAttention(nn.Module):
306 )
307
308 if self.attn_backend not in {
309 - _Backend.FLASH_ATTN,
310 - _Backend.TORCH_SDPA,
311 - _Backend.XFORMERS,
312 - _Backend.ROCM_AITER_FA,
309 + AttentionBackendEnum.FLASH_ATTN,
310 + AttentionBackendEnum.TORCH_SDPA,
311 + AttentionBackendEnum.XFORMERS,
312 + AttentionBackendEnum.ROCM_AITER_FA,
313 }:
314 raise RuntimeError(
315 f"GLM-4V does not support {self.attn_backend} backend now."
316 )
317
318 self.is_flash_attn_backend = self.attn_backend in {
319 - _Backend.FLASH_ATTN,
320 - _Backend.ROCM_AITER_FA,
319 + AttentionBackendEnum.FLASH_ATTN,
320 + AttentionBackendEnum.ROCM_AITER_FA,
321 }
322
323 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
@@ -377,7 +377,7 @@ class Glm4vVisionAttention(nn.Module):
377 context_layer = rearrange(
378 output, "(b s) h d -> s b (h d)", b=batch_size
379 ).contiguous()
380 - elif self.attn_backend == _Backend.TORCH_SDPA:
380 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
381 # Execute attention entry by entry for speed & less VRAM.
382 outputs = []
383 for i in range(1, len(cu_seqlens)):
@@ -396,7 +396,7 @@ class Glm4vVisionAttention(nn.Module):
396 context_layer = rearrange(
397 context_layer, "b s h d -> s b (h d)"
398 ).contiguous()
399 - elif self.attn_backend == _Backend.XFORMERS:
399 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
400 from xformers import ops as xops
401 from xformers.ops.fmha.attn_bias import BlockDiagonalMask
402
@@ -425,7 +425,7 @@ class Glm4vVisionBlock(nn.Module):
425 quant_config: QuantizationConfig | None = None,
426 prefix: str = "",
427 use_data_parallel: bool = False,
428 - attn_backend_override: _Backend | None = None,
428 + attn_backend_override: AttentionBackendEnum | None = None,
429 ) -> None:
430 super().__init__()
431 if norm_layer is None:
@@ -703,7 +703,7 @@ class Glm4vVisionTransformer(nn.Module):
703 quant_config: QuantizationConfig | None = None,
704 prefix: str = "",
705 use_data_parallel: bool = False,
706 - attn_backend_override: _Backend | None = None,
706 + attn_backend_override: AttentionBackendEnum | None = None,
707 ) -> None:
708 super().__init__()
709
@@ -772,10 +772,11 @@ class Glm4vVisionTransformer(nn.Module):
772 dtype=torch.get_default_dtype(),
773 attn_backend_override=attn_backend_override,
774 )
775 - if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
776 - torch.get_default_dtype()
775 + if (
776 + self.attn_backend != AttentionBackendEnum.FLASH_ATTN
777 + and check_upstream_fa_availability(torch.get_default_dtype())
778 ):
778 - self.attn_backend = _Backend.FLASH_ATTN
779 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
780
781 @property
782 def dtype(self) -> torch.dtype:
@@ -824,8 +825,8 @@ class Glm4vVisionTransformer(nn.Module):
825 max_seqlen, seqlens = None, None
826 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
827 if (
827 - self.attn_backend == _Backend.FLASH_ATTN
828 - or self.attn_backend == _Backend.ROCM_AITER_FA
828 + self.attn_backend == AttentionBackendEnum.FLASH_ATTN
829 + or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
830 ):
831 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
832 return max_seqlen, seqlens
vllm/model_executor/models/keye.py
+12 -12
@@ -16,7 +16,7 @@ from transformers.feature_extraction_utils import BatchFeature
16 from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
17 from transformers.utils import torch_int
18
19 -from vllm.attention.backends.registry import _Backend
19 +from vllm.attention.backends.registry import AttentionBackendEnum
20 from vllm.attention.layer import (
21 maybe_get_vit_flash_attn_backend,
22 )
@@ -360,7 +360,7 @@ class KeyeSiglipAttention(nn.Module):
360 config: PretrainedConfig,
361 quant_config: QuantizationConfig | None = None,
362 prefix: str = "",
363 - attn_backend_override: _Backend | None = None,
363 + attn_backend_override: AttentionBackendEnum | None = None,
364 ):
365 super().__init__()
366 self.config = config
@@ -414,17 +414,17 @@ class KeyeSiglipAttention(nn.Module):
414 )
415
416 if self.attn_backend not in {
417 - _Backend.FLASH_ATTN,
418 - _Backend.XFORMERS,
419 - _Backend.ROCM_AITER_FA,
417 + AttentionBackendEnum.FLASH_ATTN,
418 + AttentionBackendEnum.XFORMERS,
419 + AttentionBackendEnum.ROCM_AITER_FA,
420 }:
421 raise RuntimeError(
422 f"Keye-VL does not support {self.attn_backend} backend now."
423 )
424
425 self.is_flash_attn_backend = self.attn_backend in {
426 - _Backend.FLASH_ATTN,
427 - _Backend.ROCM_AITER_FA,
426 + AttentionBackendEnum.FLASH_ATTN,
427 + AttentionBackendEnum.ROCM_AITER_FA,
428 }
429
430 def forward(
@@ -489,7 +489,7 @@ class KeyeSiglipAttention(nn.Module):
489 softmax_scale=self.scale,
490 )
491 context_layer = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
492 - elif self.attn_backend == _Backend.XFORMERS:
492 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
493 from xformers import ops as xops
494 from xformers.ops.fmha.attn_bias import BlockDiagonalMask
495
@@ -536,7 +536,7 @@ class KeyeSiglipEncoderLayer(nn.Module):
536 config: PretrainedConfig,
537 quant_config: QuantizationConfig | None = None,
538 prefix: str = "",
539 - attn_backend_override: _Backend | None = None,
539 + attn_backend_override: AttentionBackendEnum | None = None,
540 ):
541 super().__init__()
542 self.embed_dim = config.hidden_size
@@ -590,7 +590,7 @@ class KeyeSiglipEncoder(nn.Module):
590 config: PretrainedConfig,
591 quant_config: QuantizationConfig | None = None,
592 prefix: str = "",
593 - attn_backend_override: _Backend | None = None,
593 + attn_backend_override: AttentionBackendEnum | None = None,
594 ):
595 super().__init__()
596 self.config = config
@@ -685,7 +685,7 @@ class KeyeSiglipVisionTransformer(nn.Module):
685 config: PretrainedConfig,
686 quant_config: QuantizationConfig | None = None,
687 prefix: str = "",
688 - attn_backend_override: _Backend | None = None,
688 + attn_backend_override: AttentionBackendEnum | None = None,
689 ):
690 super().__init__()
691 self.config = config
@@ -768,7 +768,7 @@ class KeyeSiglipVisionModel(nn.Module):
768 config: PretrainedConfig,
769 quant_config: QuantizationConfig | None = None,
770 prefix: str = "",
771 - attn_backend_override: _Backend | None = None,
771 + attn_backend_override: AttentionBackendEnum | None = None,
772 ):
773 super().__init__()
774
vllm/model_executor/models/ovis2_5.py
+3 -3
@@ -10,7 +10,7 @@ import torch
10 import torch.nn as nn
11 from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig
12
13 -from vllm.attention.backends.registry import _Backend
13 +from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.config import VllmConfig
15 from vllm.config.multimodal import BaseDummyOptions
16 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -106,7 +106,7 @@ class VisualTokenizer(torch.nn.Module):
106 quant_config: QuantizationConfig | None = None,
107 prefix: str = "",
108 use_data_parallel: bool = False,
109 - attn_backend_override: _Backend | None = None,
109 + attn_backend_override: AttentionBackendEnum | None = None,
110 ):
111 super().__init__()
112 self.config = config
@@ -135,7 +135,7 @@ class VisualTokenizer(torch.nn.Module):
135 quant_config: QuantizationConfig | None = None,
136 prefix: str = "",
137 use_data_parallel: bool = False,
138 - attn_backend_override: _Backend | None = None,
138 + attn_backend_override: AttentionBackendEnum | None = None,
139 ):
140 model_type = config.model_type
141 if model_type == "siglip2_navit":
vllm/model_executor/models/paddleocr_vl.py
+25 -22
@@ -31,7 +31,7 @@ from transformers.modeling_outputs import (
31 )
32 from transformers.utils import torch_int
33
34 -from vllm.attention.backends.registry import _Backend
34 +from vllm.attention.backends.registry import AttentionBackendEnum
35 from vllm.attention.layer import (
36 check_upstream_fa_availability,
37 maybe_get_vit_flash_attn_backend,
@@ -580,8 +580,8 @@ class SiglipAttention(nn.Module):
580 projection_size: int,
581 quant_config: QuantizationConfig | None = None,
582 prefix: str = "",
583 - attn_backend: _Backend = _Backend.TORCH_SDPA,
584 - attn_backend_override: _Backend | None = None,
583 + attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
584 + attn_backend_override: AttentionBackendEnum | None = None,
585 use_upstream_fa: bool = False,
586 ) -> None:
587 super().__init__()
@@ -621,8 +621,8 @@ class SiglipAttention(nn.Module):
621 )
622 )
623 self.is_flash_attn_backend = self.attn_backend in {
624 - _Backend.FLASH_ATTN,
625 - _Backend.ROCM_AITER_FA,
624 + AttentionBackendEnum.FLASH_ATTN,
625 + AttentionBackendEnum.ROCM_AITER_FA,
626 }
627
628 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
@@ -680,10 +680,10 @@ class SiglipAttention(nn.Module):
680 cu_seqlens,
681 max_seqlen,
682 batch_size,
683 - self.attn_backend == _Backend.ROCM_AITER_FA,
683 + self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA,
684 self.use_upstream_fa,
685 )
686 - elif self.attn_backend == _Backend.TORCH_SDPA:
686 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
687 outputs = []
688 for i in range(1, len(cu_seqlens)):
689 start_idx = cu_seqlens[i - 1]
@@ -702,7 +702,7 @@ class SiglipAttention(nn.Module):
702 context_layer = rearrange(
703 context_layer, "b s h d -> s b (h d)"
704 ).contiguous()
705 - elif self.attn_backend == _Backend.XFORMERS:
705 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
706 if seqlens is None:
707 raise ValueError("xFormers attention backend requires seqlens tensor.")
708 context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
@@ -786,8 +786,8 @@ class SiglipEncoderLayer(nn.Module):
786 quant_config: QuantizationConfig | None = None,
787 prefix: str = "",
788 *,
789 - attn_backend: _Backend = _Backend.TORCH_SDPA,
790 - attn_backend_override: _Backend | None = None,
789 + attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
790 + attn_backend_override: AttentionBackendEnum | None = None,
791 use_upstream_fa: bool = False,
792 ):
793 super().__init__()
@@ -847,7 +847,7 @@ class SiglipEncoder(nn.Module):
847 config: PretrainedConfig,
848 quant_config: QuantizationConfig | None = None,
849 prefix: str = "",
850 - attn_backend_override: _Backend | None = None,
850 + attn_backend_override: AttentionBackendEnum | None = None,
851 ):
852 super().__init__()
853 self.config = config
@@ -861,16 +861,16 @@ class SiglipEncoder(nn.Module):
861 )
862 self.use_upstream_fa = False
863 if self.attn_backend not in {
864 - _Backend.FLASH_ATTN,
865 - _Backend.ROCM_AITER_FA,
864 + AttentionBackendEnum.FLASH_ATTN,
865 + AttentionBackendEnum.ROCM_AITER_FA,
866 } and check_upstream_fa_availability(torch.get_default_dtype()):
867 - self.attn_backend = _Backend.FLASH_ATTN
867 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
868 self.use_upstream_fa = True
869 if self.attn_backend not in {
870 - _Backend.FLASH_ATTN,
871 - _Backend.TORCH_SDPA,
872 - _Backend.XFORMERS,
873 - _Backend.ROCM_AITER_FA,
870 + AttentionBackendEnum.FLASH_ATTN,
871 + AttentionBackendEnum.TORCH_SDPA,
872 + AttentionBackendEnum.XFORMERS,
873 + AttentionBackendEnum.ROCM_AITER_FA,
874 }:
875 raise RuntimeError(
876 f"PaddleOCR-VL does not support {self.attn_backend} backend now."
@@ -943,9 +943,12 @@ class SiglipEncoder(nn.Module):
943
944 max_seqlen = None
945 seqlens = None
946 - if self.attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
946 + if self.attn_backend in {
947 + AttentionBackendEnum.FLASH_ATTN,
948 + AttentionBackendEnum.ROCM_AITER_FA,
949 + }:
950 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
948 - elif self.attn_backend == _Backend.XFORMERS:
951 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
952 seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
953
954 hidden_states = inputs_embeds
@@ -966,7 +969,7 @@ class SiglipVisionTransformer(nn.Module):
969 config: PretrainedConfig,
970 quant_config: QuantizationConfig | None = None,
971 prefix: str = "",
969 - attn_backend_override: _Backend | None = None,
972 + attn_backend_override: AttentionBackendEnum | None = None,
973 ):
974 super().__init__()
975 self.config = config
@@ -1016,7 +1019,7 @@ class SiglipVisionModel(nn.Module):
1019 config,
1020 quant_config: QuantizationConfig | None = None,
1021 prefix: str = "",
1019 - attn_backend_override: _Backend | None = None,
1022 + attn_backend_override: AttentionBackendEnum | None = None,
1023 ):
1024 super().__init__()
1025
vllm/model_executor/models/qwen2_5_vl.py
+24 -18
@@ -42,7 +42,7 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
42 Qwen2_5_VLVisionConfig,
43 )
44
45 -from vllm.attention.backends.registry import _Backend
45 +from vllm.attention.backends.registry import AttentionBackendEnum
46 from vllm.attention.layer import maybe_get_vit_flash_attn_backend
47 from vllm.attention.ops.vit_attn_wrappers import (
48 vit_flash_attn_wrapper,
@@ -315,9 +315,9 @@ class Qwen2_5_VisionAttention(nn.Module):
315 quant_config: QuantizationConfig | None = None,
316 prefix: str = "",
317 use_data_parallel: bool = False,
318 - attn_backend: _Backend = _Backend.TORCH_SDPA,
318 + attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
319 use_upstream_fa: bool = False,
320 - attn_backend_override: _Backend | None = None,
320 + attn_backend_override: AttentionBackendEnum | None = None,
321 ) -> None:
322 super().__init__()
323 # Per attention head and per partition values.
@@ -364,13 +364,16 @@ class Qwen2_5_VisionAttention(nn.Module):
364 # On ROCm with FLASH_ATTN backend, upstream flash_attn is used
365 from vllm.platforms import current_platform
366
367 - if current_platform.is_rocm() and self.attn_backend == _Backend.FLASH_ATTN:
367 + if (
368 + current_platform.is_rocm()
369 + and self.attn_backend == AttentionBackendEnum.FLASH_ATTN
370 + ):
371 self.use_upstream_fa = True
372 if current_platform.is_xpu():
373 self.use_upstream_fa = False
374 self.is_flash_attn_backend = self.attn_backend in {
372 - _Backend.FLASH_ATTN,
373 - _Backend.ROCM_AITER_FA,
375 + AttentionBackendEnum.FLASH_ATTN,
376 + AttentionBackendEnum.ROCM_AITER_FA,
377 }
378
379 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
@@ -431,10 +434,10 @@ class Qwen2_5_VisionAttention(nn.Module):
434 cu_seqlens,
435 max_seqlen,
436 batch_size,
434 - self.attn_backend == _Backend.ROCM_AITER_FA,
437 + self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA,
438 self.use_upstream_fa,
439 )
437 - elif self.attn_backend == _Backend.TORCH_SDPA:
440 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
441 # Execute attention entry by entry for speed & less VRAM.
442 from vllm.platforms import current_platform
443
@@ -450,7 +453,7 @@ class Qwen2_5_VisionAttention(nn.Module):
453 v,
454 cu_seqlens,
455 )
453 - elif self.attn_backend == _Backend.XFORMERS:
456 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
457 context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
458
459 output, _ = self.proj(context_layer)
@@ -478,9 +481,9 @@ class Qwen2_5_VisionBlock(nn.Module):
481 quant_config: QuantizationConfig | None = None,
482 prefix: str = "",
483 use_data_parallel: bool = False,
481 - attn_backend: _Backend = _Backend.TORCH_SDPA,
484 + attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
485 use_upstream_fa: bool = False,
483 - attn_backend_override: _Backend | None = None,
486 + attn_backend_override: AttentionBackendEnum | None = None,
487 ) -> None:
488 super().__init__()
489 if norm_layer is None:
@@ -656,7 +659,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
659 quant_config: QuantizationConfig | None = None,
660 prefix: str = "",
661 use_data_parallel: bool = False,
659 - attn_backend_override: _Backend | None = None,
662 + attn_backend_override: AttentionBackendEnum | None = None,
663 ) -> None:
664 super().__init__()
665
@@ -708,10 +711,10 @@ class Qwen2_5_VisionTransformer(nn.Module):
711 )
712
713 if self.attn_backend not in {
711 - _Backend.FLASH_ATTN,
712 - _Backend.TORCH_SDPA,
713 - _Backend.XFORMERS,
714 - _Backend.ROCM_AITER_FA,
714 + AttentionBackendEnum.FLASH_ATTN,
715 + AttentionBackendEnum.TORCH_SDPA,
716 + AttentionBackendEnum.XFORMERS,
717 + AttentionBackendEnum.ROCM_AITER_FA,
718 }:
719 raise RuntimeError(
720 f"Qwen2.5-VL does not support {self.attn_backend} backend now."
@@ -850,9 +853,12 @@ class Qwen2_5_VisionTransformer(nn.Module):
853 ) -> tuple[torch.Tensor, torch.Tensor]:
854 max_seqlen = torch.zeros([], device=cu_seqlens.device)
855 seqlens = torch.zeros(1, device=cu_seqlens.device)
853 - if self.attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
856 + if self.attn_backend in {
857 + AttentionBackendEnum.FLASH_ATTN,
858 + AttentionBackendEnum.ROCM_AITER_FA,
859 + }:
860 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
855 - elif self.attn_backend == _Backend.XFORMERS:
861 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
862 seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
863 return max_seqlen, seqlens
864
vllm/model_executor/models/qwen2_vl.py
+21 -17
@@ -43,7 +43,7 @@ from transformers.models.qwen2_vl.configuration_qwen2_vl import (
43 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
44 from transformers.models.qwen2_vl.video_processing_qwen2_vl import Qwen2VLVideoProcessor
45
46 -from vllm.attention.backends.registry import _Backend
46 +from vllm.attention.backends.registry import AttentionBackendEnum
47 from vllm.attention.layer import (
48 check_upstream_fa_availability,
49 maybe_get_vit_flash_attn_backend,
@@ -329,7 +329,7 @@ class Qwen2VisionAttention(nn.Module):
329 quant_config: QuantizationConfig | None = None,
330 prefix: str = "",
331 use_data_parallel: bool = False,
332 - attn_backend_override: _Backend | None = None,
332 + attn_backend_override: AttentionBackendEnum | None = None,
333 ) -> None:
334 super().__init__()
335 # Per attention head and per partition values.
@@ -378,18 +378,18 @@ class Qwen2VisionAttention(nn.Module):
378 )
379
380 if self.attn_backend not in {
381 - _Backend.FLASH_ATTN,
382 - _Backend.TORCH_SDPA,
383 - _Backend.XFORMERS,
384 - _Backend.ROCM_AITER_FA,
381 + AttentionBackendEnum.FLASH_ATTN,
382 + AttentionBackendEnum.TORCH_SDPA,
383 + AttentionBackendEnum.XFORMERS,
384 + AttentionBackendEnum.ROCM_AITER_FA,
385 }:
386 raise RuntimeError(
387 f"Qwen2-VL does not support {self.attn_backend} backend now."
388 )
389
390 self.is_flash_attn_backend = self.attn_backend in {
391 - _Backend.FLASH_ATTN,
392 - _Backend.ROCM_AITER_FA,
391 + AttentionBackendEnum.FLASH_ATTN,
392 + AttentionBackendEnum.ROCM_AITER_FA,
393 }
394
395 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
@@ -460,7 +460,7 @@ class Qwen2VisionAttention(nn.Module):
460 context_layer = rearrange(
461 output, "(b s) h d -> s b (h d)", b=batch_size
462 ).contiguous()
463 - elif self.attn_backend == _Backend.TORCH_SDPA:
463 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
464 # Execute attention entry by entry for speed & less VRAM.
465 from vllm.platforms import current_platform
466
@@ -485,7 +485,7 @@ class Qwen2VisionAttention(nn.Module):
485 context_layer = rearrange(
486 context_layer, "b s h d -> s b (h d)"
487 ).contiguous()
488 - elif self.attn_backend == _Backend.XFORMERS:
488 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
489 from xformers import ops as xops
490 from xformers.ops.fmha.attn_bias import BlockDiagonalMask
491
@@ -515,7 +515,7 @@ class Qwen2VisionBlock(nn.Module):
515 quant_config: QuantizationConfig | None = None,
516 prefix: str = "",
517 use_data_parallel: bool = False,
518 - attn_backend_override: _Backend | None = None,
518 + attn_backend_override: AttentionBackendEnum | None = None,
519 ) -> None:
520 super().__init__()
521 if norm_layer is None:
@@ -679,7 +679,7 @@ class Qwen2VisionTransformer(nn.Module):
679 quant_config: QuantizationConfig | None = None,
680 prefix: str = "",
681 use_data_parallel: bool = False,
682 - attn_backend_override: _Backend | None = None,
682 + attn_backend_override: AttentionBackendEnum | None = None,
683 ) -> None:
684 super().__init__()
685
@@ -739,10 +739,11 @@ class Qwen2VisionTransformer(nn.Module):
739 dtype=torch.get_default_dtype(),
740 attn_backend_override=attn_backend_override,
741 )
742 - if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
743 - torch.get_default_dtype()
742 + if (
743 + self.attn_backend != AttentionBackendEnum.FLASH_ATTN
744 + and check_upstream_fa_availability(torch.get_default_dtype())
745 ):
745 - self.attn_backend = _Backend.FLASH_ATTN
746 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
747
748 @property
749 def dtype(self) -> torch.dtype:
@@ -789,9 +790,12 @@ class Qwen2VisionTransformer(nn.Module):
790 self, cu_seqlens: torch.Tensor
791 ) -> tuple[int | None, list[int] | None]:
792 max_seqlen, seqlens = None, None
792 - if self.attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
793 + if self.attn_backend in {
794 + AttentionBackendEnum.FLASH_ATTN,
795 + AttentionBackendEnum.ROCM_AITER_FA,
796 + }:
797 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
794 - elif self.attn_backend == _Backend.XFORMERS:
798 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
799 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
800 return max_seqlen, seqlens
801
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+8 -7
@@ -47,7 +47,7 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
47 )
48 from transformers.models.whisper import WhisperFeatureExtractor
49
50 -from vllm.attention.backends.registry import _Backend
50 +from vllm.attention.backends.registry import AttentionBackendEnum
51 from vllm.attention.layer import check_upstream_fa_availability
52 from vllm.compilation.decorators import support_torch_compile
53 from vllm.config import VllmConfig
@@ -301,7 +301,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):
301 norm_eps: float = 1e-6,
302 quant_config: QuantizationConfig | None = None,
303 prefix: str = "",
304 - attn_backend_override: _Backend | None = None,
304 + attn_backend_override: AttentionBackendEnum | None = None,
305 ) -> None:
306 super().__init__()
307 self.hidden_size = vision_config.hidden_size
@@ -377,10 +377,11 @@ class Qwen3Omni_VisionTransformer(nn.Module):
377 dtype=torch.get_default_dtype(),
378 attn_backend_override=attn_backend_override,
379 )
380 - if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
381 - torch.get_default_dtype()
380 + if (
381 + self.attn_backend != AttentionBackendEnum.FLASH_ATTN
382 + and check_upstream_fa_availability(torch.get_default_dtype())
383 ):
383 - self.attn_backend = _Backend.FLASH_ATTN
384 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
385
386 @property
387 def dtype(self) -> torch.dtype:
@@ -490,9 +491,9 @@ class Qwen3Omni_VisionTransformer(nn.Module):
491 ) -> tuple[torch.Tensor, torch.Tensor]:
492 max_seqlen = torch.zeros([], device=cu_seqlens.device)
493 seqlens = torch.zeros(1, device=cu_seqlens.device)
493 - if self.attn_backend == _Backend.FLASH_ATTN:
494 + if self.attn_backend == AttentionBackendEnum.FLASH_ATTN:
495 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
495 - elif self.attn_backend == _Backend.XFORMERS:
496 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
497 seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
498 return max_seqlen, seqlens
499
vllm/model_executor/models/qwen3_vl.py
+13 -13
@@ -49,7 +49,7 @@ from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
49 )
50 from transformers.video_utils import VideoMetadata
51
52 -from vllm.attention.backends.registry import _Backend
52 +from vllm.attention.backends.registry import AttentionBackendEnum
53 from vllm.attention.layer import check_upstream_fa_availability
54 from vllm.compilation.decorators import support_torch_compile
55 from vllm.config import VllmConfig
@@ -198,7 +198,7 @@ class Qwen3_VisionBlock(nn.Module):
198 quant_config: QuantizationConfig | None = None,
199 prefix: str = "",
200 use_data_parallel: bool = False,
201 - attn_backend: _Backend = _Backend.TORCH_SDPA,
201 + attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
202 use_upstream_fa: bool = False,
203 ) -> None:
204 super().__init__()
@@ -306,7 +306,7 @@ class Qwen3_VisionTransformer(nn.Module):
306 quant_config: QuantizationConfig | None = None,
307 prefix: str = "",
308 use_data_parallel: bool = False,
309 - attn_backend_override: _Backend | None = None,
309 + attn_backend_override: AttentionBackendEnum | None = None,
310 ) -> None:
311 super().__init__()
312 self.hidden_size = vision_config.hidden_size
@@ -372,18 +372,18 @@ class Qwen3_VisionTransformer(nn.Module):
372 )
373 use_upstream_fa = False
374 if (
375 - self.attn_backend != _Backend.FLASH_ATTN
376 - and self.attn_backend != _Backend.ROCM_AITER_FA
375 + self.attn_backend != AttentionBackendEnum.FLASH_ATTN
376 + and self.attn_backend != AttentionBackendEnum.ROCM_AITER_FA
377 and check_upstream_fa_availability(torch.get_default_dtype())
378 ):
379 - self.attn_backend = _Backend.FLASH_ATTN
379 + self.attn_backend = AttentionBackendEnum.FLASH_ATTN
380 use_upstream_fa = True
381
382 if self.attn_backend not in {
383 - _Backend.FLASH_ATTN,
384 - _Backend.TORCH_SDPA,
385 - _Backend.XFORMERS,
386 - _Backend.ROCM_AITER_FA,
383 + AttentionBackendEnum.FLASH_ATTN,
384 + AttentionBackendEnum.TORCH_SDPA,
385 + AttentionBackendEnum.XFORMERS,
386 + AttentionBackendEnum.ROCM_AITER_FA,
387 }:
388 raise RuntimeError(
389 f"Qwen3-VL does not support {self.attn_backend} backend now."
@@ -510,11 +510,11 @@ class Qwen3_VisionTransformer(nn.Module):
510 max_seqlen = torch.zeros([], device=cu_seqlens.device)
511 seqlens = torch.zeros(1, device=cu_seqlens.device)
512 if (
513 - self.attn_backend == _Backend.FLASH_ATTN
514 - or self.attn_backend == _Backend.ROCM_AITER_FA
513 + self.attn_backend == AttentionBackendEnum.FLASH_ATTN
514 + or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
515 ):
516 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
517 - elif self.attn_backend == _Backend.XFORMERS:
517 + elif self.attn_backend == AttentionBackendEnum.XFORMERS:
518 seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
519 return max_seqlen, seqlens
520
vllm/model_executor/models/siglip2navit.py
+13 -13
@@ -12,7 +12,7 @@ from torch.nn import functional as F
12 from transformers import Siglip2VisionConfig
13 from transformers.configuration_utils import PretrainedConfig
14
15 -from vllm.attention.backends.registry import _Backend
15 +from vllm.attention.backends.registry import AttentionBackendEnum
16 from vllm.attention.layer import maybe_get_vit_flash_attn_backend
17 from vllm.distributed import divide, get_tensor_model_parallel_world_size
18 from vllm.model_executor.layers.activation import get_act_fn
@@ -208,7 +208,7 @@ class Siglip2Attention(nn.Module):
208 quant_config: QuantizationConfig | None = None,
209 prefix: str = "",
210 use_data_parallel: bool = False,
211 - attn_backend_override: _Backend | None = None,
211 + attn_backend_override: AttentionBackendEnum | None = None,
212 ):
213 super().__init__()
214 self.config = config
@@ -264,14 +264,14 @@ class Siglip2Attention(nn.Module):
264 )
265
266 if self.attn_backend not in {
267 - _Backend.FLASH_ATTN,
268 - _Backend.TORCH_SDPA,
269 - _Backend.ROCM_AITER_FA,
267 + AttentionBackendEnum.FLASH_ATTN,
268 + AttentionBackendEnum.TORCH_SDPA,
269 + AttentionBackendEnum.ROCM_AITER_FA,
270 }:
271 - self.attn_backend = _Backend.TORCH_SDPA
271 + self.attn_backend = AttentionBackendEnum.TORCH_SDPA
272 self.is_flash_attn_backend = self.attn_backend in {
273 - _Backend.FLASH_ATTN,
274 - _Backend.ROCM_AITER_FA,
273 + AttentionBackendEnum.FLASH_ATTN,
274 + AttentionBackendEnum.ROCM_AITER_FA,
275 }
276
277 def forward(
@@ -308,7 +308,7 @@ class Siglip2Attention(nn.Module):
308 attn_output = self.flash_attn_varlen_func(
309 queries, keys, values, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen
310 ).reshape(seq_length, -1)
311 - elif self.attn_backend == _Backend.TORCH_SDPA:
311 + elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
312 # Execute attention entry by entry for speed & less VRAM.
313 batch_size = cu_seqlens.shape[0] - 1
314 outputs = []
@@ -376,7 +376,7 @@ class Siglip2EncoderLayer(nn.Module):
376 quant_config: QuantizationConfig | None = None,
377 prefix: str = "",
378 use_data_parallel: bool = False,
379 - attn_backend_override: _Backend | None = None,
379 + attn_backend_override: AttentionBackendEnum | None = None,
380 ):
381 super().__init__()
382 self.embed_dim = config.hidden_size
@@ -440,7 +440,7 @@ class Siglip2Encoder(nn.Module):
440 quant_config: QuantizationConfig | None = None,
441 prefix: str = "",
442 use_data_parallel: bool = False,
443 - attn_backend_override: _Backend | None = None,
443 + attn_backend_override: AttentionBackendEnum | None = None,
444 ):
445 super().__init__()
446 self.config = config
@@ -626,7 +626,7 @@ class Siglip2VisionTransformer(nn.Module):
626 quant_config: QuantizationConfig | None = None,
627 prefix: str = "",
628 use_data_parallel: bool = False,
629 - attn_backend_override: _Backend | None = None,
629 + attn_backend_override: AttentionBackendEnum | None = None,
630 ):
631 super().__init__()
632 self.config = config
@@ -667,7 +667,7 @@ class Siglip2NavitModel(torch.nn.Module):
667 quant_config: QuantizationConfig | None = None,
668 prefix: str = "",
669 use_data_parallel: bool = False,
670 - attn_backend_override: _Backend | None = None,
670 + attn_backend_override: AttentionBackendEnum | None = None,
671 ):
672 super().__init__()
673
vllm/model_executor/models/vision.py
+4 -4
@@ -10,7 +10,7 @@ from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar
10 import torch
11 from transformers import PretrainedConfig
12
13 -from vllm.attention.backends.registry import _Backend
13 +from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.config import VllmConfig
15 from vllm.distributed import (
16 get_tensor_model_parallel_rank,
@@ -83,8 +83,8 @@ def get_vit_attn_backend(
83 head_size: int,
84 dtype: torch.dtype,
85 *,
86 - attn_backend_override: _Backend | None = None,
87 -) -> _Backend:
86 + attn_backend_override: AttentionBackendEnum | None = None,
87 +) -> AttentionBackendEnum:
88 """
89 Get the available attention backend for Vision Transformer.
90 """
@@ -94,7 +94,7 @@ def get_vit_attn_backend(
94 # Lazy import to avoid circular dependency
95 from vllm.attention.selector import get_env_variable_attn_backend
96
97 - selected_backend: _Backend | None = get_env_variable_attn_backend()
97 + selected_backend: AttentionBackendEnum | None = get_env_variable_attn_backend()
98 if selected_backend is not None:
99 return selected_backend
100
vllm/platforms/cpu.py
+6 -6
@@ -23,10 +23,10 @@ from .interface import CpuArchEnum, Platform, PlatformEnum
23 logger = init_logger(__name__)
24
25 if TYPE_CHECKING:
26 - from vllm.attention.backends.registry import _Backend
26 + from vllm.attention.backends.registry import AttentionBackendEnum
27 from vllm.config import VllmConfig
28 else:
29 - _Backend = None
29 + AttentionBackendEnum = None
30 VllmConfig = None
31
32
@@ -127,7 +127,7 @@ class CpuPlatform(Platform):
127 @classmethod
128 def get_attn_backend_cls(
129 cls,
130 - selected_backend: "_Backend",
130 + selected_backend: "AttentionBackendEnum",
131 head_size: int,
132 dtype: torch.dtype,
133 kv_cache_dtype: str | None,
@@ -137,9 +137,9 @@ class CpuPlatform(Platform):
137 has_sink: bool,
138 use_sparse: bool,
139 ) -> str:
140 - from vllm.attention.backends.registry import _Backend
140 + from vllm.attention.backends.registry import AttentionBackendEnum
141
142 - if selected_backend and selected_backend != _Backend.TORCH_SDPA:
142 + if selected_backend and selected_backend != AttentionBackendEnum.TORCH_SDPA:
143 logger.info("Cannot use %s backend on CPU.", selected_backend)
144 if use_mla:
145 raise NotImplementedError("MLA is not supported on CPU.")
@@ -148,7 +148,7 @@ class CpuPlatform(Platform):
148 logger.info("Using Torch SDPA backend.")
149 if not use_v1:
150 raise ValueError("CPU backend only supports V1.")
151 - return "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend"
151 + return AttentionBackendEnum.TORCH_SDPA.get_path()
152
153 @classmethod
154 def get_device_total_memory(cls, device_id: int = 0) -> int:
vllm/platforms/cuda.py
+181 -181
@@ -22,10 +22,13 @@ 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.backends.registry import _Backend
25 + from vllm.attention.backends.registry import AttentionBackendEnum
26 from vllm.config import VllmConfig
27 + from vllm.config.cache import CacheDType
28 else:
28 - _Backend = None
29 + AttentionBackendEnum = None
30 + VllmConfig = None
31 + CacheDType = None
32
33 logger = init_logger(__name__)
34
@@ -39,6 +42,49 @@ pynvml = import_pynvml()
42 torch.backends.cuda.enable_cudnn_sdp(False)
43
44
45 +@cache
46 +def _get_backend_priorities(
47 + use_mla: bool,
48 + device_capability: DeviceCapability,
49 +) -> list[AttentionBackendEnum]:
50 + """Get backend priorities with lazy import to avoid circular dependency."""
51 + from vllm.attention.backends.registry import AttentionBackendEnum
52 +
53 + if use_mla:
54 + if device_capability.major == 10:
55 + return [
56 + AttentionBackendEnum.CUTLASS_MLA,
57 + AttentionBackendEnum.FLASHINFER_MLA,
58 + AttentionBackendEnum.FLASHMLA,
59 + AttentionBackendEnum.FLASH_ATTN_MLA,
60 + AttentionBackendEnum.TRITON_MLA,
61 + AttentionBackendEnum.FLASHMLA_SPARSE,
62 + ]
63 + else:
64 + return [
65 + AttentionBackendEnum.FLASHMLA,
66 + AttentionBackendEnum.FLASH_ATTN_MLA,
67 + AttentionBackendEnum.FLASHINFER_MLA,
68 + AttentionBackendEnum.TRITON_MLA,
69 + AttentionBackendEnum.FLASHMLA_SPARSE,
70 + ]
71 + else:
72 + if device_capability.major == 10:
73 + return [
74 + AttentionBackendEnum.FLASHINFER,
75 + AttentionBackendEnum.FLASH_ATTN,
76 + AttentionBackendEnum.TRITON_ATTN,
77 + AttentionBackendEnum.FLEX_ATTENTION,
78 + ]
79 + else:
80 + return [
81 + AttentionBackendEnum.FLASH_ATTN,
82 + AttentionBackendEnum.FLASHINFER,
83 + AttentionBackendEnum.TRITON_ATTN,
84 + AttentionBackendEnum.FLEX_ATTENTION,
85 + ]
86 +
87 +
88 def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
89 @wraps(fn)
90 def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
@@ -216,217 +262,171 @@ class CudaPlatformBase(Platform):
262 return torch.cuda.max_memory_allocated(device)
263
264 @classmethod
219 - def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":
220 - from vllm.attention.backends.registry import _Backend
265 + def get_vit_attn_backend(
266 + cls, head_size: int, dtype: torch.dtype
267 + ) -> "AttentionBackendEnum":
268 + from vllm.attention.backends.registry import AttentionBackendEnum
269
270 # For Blackwell GPUs, force TORCH_SDPA for now.
271 # See https://github.com/facebookresearch/xformers/issues/1317#issuecomment-3199392579 # noqa: E501
272 if cls.has_device_capability(100):
225 - return _Backend.TORCH_SDPA
273 + return AttentionBackendEnum.TORCH_SDPA
274
275 if dtype not in (torch.float16, torch.bfloat16):
228 - return _Backend.XFORMERS
276 + return AttentionBackendEnum.XFORMERS
277
278 if cls.has_device_capability(80):
231 - FLASH_ATTN_V1 = (
232 - "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
233 - )
234 - from vllm.attention.selector import is_attn_backend_supported
235 -
236 - is_default_fa_supported = is_attn_backend_supported(
237 - FLASH_ATTN_V1, head_size, dtype, allow_import_error=False
238 - )
239 - if is_default_fa_supported:
240 - return _Backend.FLASH_ATTN
279 + backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()
280 + if backend_class.supports_head_size(
281 + head_size
282 + ) and backend_class.supports_dtype(dtype):
283 + return AttentionBackendEnum.FLASH_ATTN
284 else:
242 - # Fallback to XFORMERS
243 - return _Backend.XFORMERS
285 + return AttentionBackendEnum.XFORMERS
286 else:
287 # Fallback for Volta/Turing GPUs or FA not supported
246 - return _Backend.XFORMERS
288 + return AttentionBackendEnum.XFORMERS
289
290 @classmethod
249 - def get_attn_backend_cls(
291 + def get_valid_backends(
292 cls,
251 - selected_backend,
293 head_size,
294 dtype,
295 kv_cache_dtype,
296 block_size,
256 - use_v1,
297 use_mla,
298 has_sink,
299 use_sparse,
260 - ) -> str:
261 - from vllm.attention.backends.registry import _Backend
262 -
263 - if use_mla:
264 - # explicitly reject non-MLA backends when MLA is enabled to avoid
265 - # silently selecting an incompatible backend (e.g., FLASHINFER).
266 - if selected_backend in {
267 - _Backend.FLASHINFER,
268 - _Backend.FLASH_ATTN,
269 - _Backend.TRITON_ATTN,
270 - _Backend.TREE_ATTN,
271 - _Backend.XFORMERS,
272 - }:
273 - raise ValueError(
274 - f"Attention backend {selected_backend} incompatible with MLA. "
275 - "Please use one of the MLA backends: FLASHINFER_MLA, CUTLASS_MLA, "
276 - "FLASHMLA, FLASH_ATTN_MLA, or TRITON_MLA. Alternatively, set "
277 - "VLLM_MLA_DISABLE=1 to disable MLA for this model."
300 + device_capability,
301 + ) -> tuple[
302 + list[tuple["AttentionBackendEnum", int]],
303 + dict["AttentionBackendEnum", list[str]],
304 + ]:
305 + valid_backends_priorities = []
306 + invalid_reasons = {}
307 +
308 + backend_priorities = _get_backend_priorities(use_mla, device_capability)
309 + for priority, backend in enumerate(backend_priorities):
310 + try:
311 + backend_class = backend.get_class()
312 + invalid_reasons_i = backend_class.validate_configuration(
313 + head_size,
314 + dtype,
315 + kv_cache_dtype,
316 + block_size,
317 + use_mla,
318 + has_sink,
319 + use_sparse,
320 + device_capability,
321 )
322 + except ImportError:
323 + invalid_reasons_i = ["ImportError"]
324 + if invalid_reasons_i:
325 + invalid_reasons[backend] = invalid_reasons_i
326 + else:
327 + valid_backends_priorities.append((backend, priority))
328
280 - from vllm.attention.ops.flashmla import is_flashmla_dense_supported
281 - from vllm.attention.utils.fa_utils import flash_attn_supports_mla
329 + return valid_backends_priorities, invalid_reasons
330
283 - if use_sparse:
284 - logger.info_once("Using Sparse MLA backend.")
285 - return (
286 - "vllm.v1.attention.backends.mla.flashmla_sparse."
287 - "FlashMLASparseBackend"
288 - )
289 -
290 - use_cutlassmla = selected_backend == _Backend.CUTLASS_MLA or (
291 - selected_backend is None
292 - and cls.is_device_capability(100)
293 - and block_size % 128 == 0
294 - )
295 - use_flashinfermla = selected_backend == _Backend.FLASHINFER_MLA or (
296 - selected_backend is None
297 - and cls.is_device_capability(100)
298 - and (block_size == 32 or block_size % 64 == 0)
299 - )
300 - use_flashmla = selected_backend == _Backend.FLASHMLA or (
301 - selected_backend is None and is_flashmla_dense_supported()[0]
302 - )
303 - use_flashattn = selected_backend == _Backend.FLASH_ATTN_MLA or (
304 - selected_backend is None and flash_attn_supports_mla()
305 - )
306 - use_triton = selected_backend == _Backend.TRITON_MLA or (
307 - selected_backend is None
331 + @classmethod
332 + def get_attn_backend_cls(
333 + cls,
334 + selected_backend: "AttentionBackendEnum",
335 + head_size: int,
336 + dtype: torch.dtype,
337 + kv_cache_dtype: "CacheDType | None",
338 + block_size: int | None,
339 + use_v1: bool,
340 + use_mla: bool,
341 + has_sink: bool,
342 + use_sparse: bool,
343 + ) -> str:
344 + if not use_v1:
345 + raise RuntimeError(
346 + "V0 attention backends have been removed. Set VLLM_USE_V1=1 "
347 + "to select a supported backend."
348 )
349
310 - if use_cutlassmla:
311 - logger.info_once("Using Cutlass MLA backend.", scope="local")
312 - return "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"
313 - if use_flashinfermla:
314 - from vllm.v1.attention.backends.utils import set_kv_cache_layout
315 -
316 - set_kv_cache_layout("HND")
317 - logger.info_once("Using FlashInfer MLA backend.")
318 - return (
319 - "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
350 + device_capability = cls.get_device_capability()
351 + assert device_capability is not None
352 +
353 + # First try checking just the selected backend, if there is one.
354 + if selected_backend is not None:
355 + try:
356 + backend_class = selected_backend.get_class()
357 + invalid_reasons = backend_class.validate_configuration(
358 + head_size,
359 + dtype,
360 + kv_cache_dtype,
361 + None,
362 + use_mla,
363 + has_sink,
364 + use_sparse,
365 + device_capability,
366 )
321 - if use_flashmla:
322 - if block_size % 64 != 0:
323 - logger.warning(
324 - "FlashMLA backend is not supported for block size %d"
325 - " (currently only supports block size 64).",
326 - block_size,
327 - )
328 - else:
329 - logger.info_once("Using FlashMLA backend.")
330 - return "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend"
331 - if use_flashattn:
332 - logger.info_once("Using FlashAttention MLA backend.")
333 - return (
334 - "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
367 + except ImportError:
368 + invalid_reasons = ["ImportError"]
369 + if invalid_reasons:
370 + raise ValueError(
371 + f"Selected backend {selected_backend} is not valid for "
372 + f"this configuration. Reason: {invalid_reasons}"
373 )
336 - if use_triton:
337 - logger.info_once("Using Triton MLA backend.")
338 - return "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
339 -
340 - FLASHINFER_V1 = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" # noqa: E501
341 - FLEX_ATTENTION_V1 = (
342 - "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501
374 + else:
375 + logger.info("Using %s backend.", selected_backend)
376 + return selected_backend.get_path()
377 +
378 + # No selected backend or the selected backend is invalid,
379 + # so we try finding a valid backend.
380 + valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
381 + head_size,
382 + dtype,
383 + kv_cache_dtype,
384 + None,
385 + use_mla,
386 + has_sink,
387 + use_sparse,
388 + device_capability,
389 )
344 - TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" # noqa: E501
345 - FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
346 - TREE_ATTN_V1 = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend" # noqa: E501
347 - XFORMERS_V1 = "vllm.v1.attention.backends.xformers.XFormersAttentionBackend" # noqa: E501
348 -
349 - use_fp8_kv_cache = kv_cache_dtype is not None and kv_cache_dtype.startswith(
350 - "fp8"
390 + reasons_str = (
391 + "{"
392 + + ", ".join(
393 + f"{backend.name}: [{', '.join(reasons)}]"
394 + for backend, reasons in invalid_reasons.items()
395 + )
396 + + "}"
397 )
398 + config_str = (
399 + f"head_size: {head_size}, dtype: {dtype}, "
400 + f"kv_cache_dtype: {kv_cache_dtype}, block_size: {block_size}, "
401 + f"use_mla: {use_mla}, has_sink: {has_sink}, use_sparse: {use_sparse}"
402 + )
403 + logger.debug_once(
404 + f"Some attention backends are not valid for {cls.device_name} with "
405 + f"{config_str}. Reasons: {reasons_str}."
406 + )
407 + if len(valid_backends_priorities) == 0:
408 + raise ValueError(
409 + f"No valid attention backend found for {cls.device_name} "
410 + f"with {config_str}. Reasons: {reasons_str}."
411 + )
412
353 - if selected_backend == _Backend.FLASHINFER:
354 - logger.info_once("Using FlashInfer backend.")
355 - if cls.has_device_capability(100):
356 - from vllm.v1.attention.backends.utils import set_kv_cache_layout
357 -
358 - set_kv_cache_layout("HND")
359 - return FLASHINFER_V1
360 - elif selected_backend == _Backend.FLEX_ATTENTION:
361 - logger.info_once("Using FlexAttention backend.")
362 - return FLEX_ATTENTION_V1
363 - elif selected_backend == _Backend.TRITON_ATTN:
364 - logger.info_once("Using Triton backend.")
365 - return TRITON_ATTN
366 - elif selected_backend == _Backend.FLASH_ATTN:
367 - logger.info_once("Using Flash Attention backend.")
368 - return FLASH_ATTN_V1
369 - elif selected_backend == _Backend.TREE_ATTN:
370 - logger.info_once("Using Tree Attention backend.")
371 - return TREE_ATTN_V1
372 - elif selected_backend == _Backend.XFORMERS:
373 - logger.info_once("Using XFormers backend.")
374 - return XFORMERS_V1
375 -
376 - from vllm.attention.selector import is_attn_backend_supported
377 -
378 - # Default backends for V1 engine
379 - # Prefer FlashInfer for Blackwell GPUs if installed
380 - if cls.is_device_capability(100):
381 - if is_default_backend_supported := is_attn_backend_supported(
382 - FLASHINFER_V1, head_size, dtype
383 - ):
384 - from vllm.v1.attention.backends.utils import set_kv_cache_layout
385 -
386 - logger.info_once(
387 - "Using FlashInfer backend with HND KV cache layout on "
388 - "V1 engine by default for Blackwell (SM 10.0) GPUs."
389 - )
390 - set_kv_cache_layout("HND")
391 -
392 - return FLASHINFER_V1
393 -
394 - if not is_default_backend_supported.can_import:
395 - logger.warning_once(
396 - "FlashInfer failed to import on Blackwell (SM 10.0) GPUs; "
397 - "it is recommended to install FlashInfer for better "
398 - "performance."
399 - )
400 -
401 - # FlashAttention is the default for SM 8.0+ GPUs
402 - if cls.has_device_capability(80):
403 - if (has_sink or use_fp8_kv_cache) and not cls.is_device_capability(90):
404 - logger.info_once("Using Triton backend.")
405 - return TRITON_ATTN
406 - elif is_default_backend_supported := is_attn_backend_supported(
407 - FLASH_ATTN_V1, head_size, dtype, allow_import_error=False
408 - ):
409 - logger.info_once("Using Flash Attention backend.")
410 - return FLASH_ATTN_V1
411 -
412 - # FlexAttention is the default for older GPUs
413 - else:
414 - logger.info_once("Using FlexAttention backend.")
415 - return FLEX_ATTENTION_V1
416 -
417 - assert not is_default_backend_supported
418 -
419 - use_flex_attention_reason = {}
420 - if not is_default_backend_supported.head_size:
421 - use_flex_attention_reason["head_size"] = head_size
422 - if not is_default_backend_supported.dtype:
423 - use_flex_attention_reason["dtype"] = dtype
424 -
425 - logger.info_once(
426 - "Using FlexAttention backend for %s.",
427 - ", ".join(f"{k}={v}" for k, v in use_flex_attention_reason.items()),
413 + # We have found some valid backends. Select the one with the
414 + # highest priority.
415 + logger.info(
416 + "Valid backends: %s", [b[0].name for b in valid_backends_priorities]
417 )
429 - return FLEX_ATTENTION_V1
418 + sorted_indices = sorted(
419 + range(len(valid_backends_priorities)),
420 + key=lambda i: valid_backends_priorities[i][1],
421 + )
422 + selected_index = sorted_indices[0]
423 + selected_backend = valid_backends_priorities[selected_index][0]
424 + logger.info(
425 + "Using %s backend.",
426 + selected_backend.name,
427 + )
428 +
429 + return selected_backend.get_path()
430
431 @classmethod
432 def get_punica_wrapper(cls) -> str:
vllm/platforms/interface.py
+35 -7
@@ -17,8 +17,9 @@ from vllm.logger import init_logger
17 if TYPE_CHECKING:
18 from torch.distributed import PrefixStore, ProcessGroup
19
20 - from vllm.attention.backends.registry import _Backend
20 + from vllm.attention.backends.registry import AttentionBackendEnum
21 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
@@ -58,6 +59,31 @@ class DeviceCapability(NamedTuple):
59 major: int
60 minor: int
61
62 + def __lt__(self, other: Any) -> bool:
63 + if not isinstance(other, DeviceCapability):
64 + return NotImplemented
65 + return (self.major, self.minor) < (other.major, other.minor)
66 +
67 + def __le__(self, other: Any) -> bool:
68 + if not isinstance(other, DeviceCapability):
69 + return NotImplemented
70 + return (self.major, self.minor) <= (other.major, other.minor)
71 +
72 + def __eq__(self, other: Any) -> bool:
73 + if not isinstance(other, DeviceCapability):
74 + return NotImplemented
75 + return (self.major, self.minor) == (other.major, other.minor)
76 +
77 + def __ge__(self, other: Any) -> bool:
78 + if not isinstance(other, DeviceCapability):
79 + return NotImplemented
80 + return (self.major, self.minor) >= (other.major, other.minor)
81 +
82 + def __gt__(self, other: Any) -> bool:
83 + if not isinstance(other, DeviceCapability):
84 + return NotImplemented
85 + return (self.major, self.minor) > (other.major, other.minor)
86 +
87 def as_version_str(self) -> str:
88 return f"{self.major}.{self.minor}"
89
@@ -173,19 +199,21 @@ class Platform:
199 import vllm._moe_C # noqa: F401
200
201 @classmethod
176 - def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":
177 - # Import _Backend here to avoid circular import.
178 - from vllm.attention.backends.registry import _Backend
202 + def get_vit_attn_backend(
203 + cls, head_size: int, dtype: torch.dtype
204 + ) -> "AttentionBackendEnum":
205 + # Import AttentionBackendEnum here to avoid circular import.
206 + from vllm.attention.backends.registry import AttentionBackendEnum
207
180 - return _Backend.TORCH_SDPA
208 + return AttentionBackendEnum.TORCH_SDPA
209
210 @classmethod
211 def get_attn_backend_cls(
212 cls,
185 - selected_backend: "_Backend",
213 + selected_backend: "AttentionBackendEnum",
214 head_size: int,
215 dtype: torch.dtype,
188 - kv_cache_dtype: str | None,
216 + kv_cache_dtype: "CacheDType | None",
217 block_size: int,
218 use_v1: bool,
219 use_mla: bool,
vllm/platforms/rocm.py
+24 -25
@@ -14,10 +14,10 @@ from vllm.utils.torch_utils import cuda_device_count_stateless
14 from .interface import DeviceCapability, Platform, PlatformEnum
15
16 if TYPE_CHECKING:
17 - from vllm.attention.backends.registry import _Backend
17 + from vllm.attention.backends.registry import AttentionBackendEnum
18 from vllm.config import VllmConfig
19 else:
20 - _Backend = None
20 + AttentionBackendEnum = None
21
22 logger = init_logger(__name__)
23
@@ -204,21 +204,23 @@ class RocmPlatform(Platform):
204 ]
205
206 @classmethod
207 - def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> _Backend:
207 + def get_vit_attn_backend(
208 + cls, head_size: int, dtype: torch.dtype
209 + ) -> AttentionBackendEnum:
210 from importlib.util import find_spec
211
212 from vllm._aiter_ops import rocm_aiter_ops
211 - from vllm.attention.backends.registry import _Backend
213 + from vllm.attention.backends.registry import AttentionBackendEnum
214
215 if rocm_aiter_ops.is_mha_enabled():
216 # Note: AITER FA is only supported for Qwen-VL models.
217 # TODO: Add support for other VL models in their model class.
216 - return _Backend.ROCM_AITER_FA
218 + return AttentionBackendEnum.ROCM_AITER_FA
219
220 if on_gfx9() and find_spec("flash_attn") is not None:
219 - return _Backend.FLASH_ATTN
221 + return AttentionBackendEnum.FLASH_ATTN
222
221 - return _Backend.TORCH_SDPA
223 + return AttentionBackendEnum.TORCH_SDPA
224
225 @classmethod
226 def get_attn_backend_cls(
@@ -234,7 +236,7 @@ class RocmPlatform(Platform):
236 use_sparse,
237 ) -> str:
238 from vllm._aiter_ops import rocm_aiter_ops
237 - from vllm.attention.backends.registry import _Backend
239 + from vllm.attention.backends.registry import AttentionBackendEnum
240
241 if use_sparse:
242 raise NotImplementedError("Sparse Attention is not supported on ROCm.")
@@ -248,55 +250,52 @@ class RocmPlatform(Platform):
250 if use_mla:
251 if selected_backend is None:
252 selected_backend = (
251 - _Backend.ROCM_AITER_MLA
253 + AttentionBackendEnum.ROCM_AITER_MLA
254 if rocm_aiter_ops.is_mla_enabled() or block_size == 1
253 - else _Backend.TRITON_MLA
255 + else AttentionBackendEnum.TRITON_MLA
256 )
257
256 - if selected_backend == _Backend.TRITON_MLA:
258 + if selected_backend == AttentionBackendEnum.TRITON_MLA:
259 if block_size != 1:
260 logger.info_once("Using Triton MLA backend.")
259 - return "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
261 + return AttentionBackendEnum.TRITON_MLA.get_path()
262 raise ValueError(
263 f" The selected backend, {selected_backend.name},"
264 f"does not support block size {block_size}."
265 )
264 - if selected_backend == _Backend.ROCM_AITER_MLA:
266 + if selected_backend == AttentionBackendEnum.ROCM_AITER_MLA:
267 logger.info("Using AITER MLA backend.")
266 - return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501
268 + return AttentionBackendEnum.ROCM_AITER_MLA.get_path()
269
270 raise ValueError(
271 f" The selected backend, {selected_backend.name},"
272 f"is not MLA type while requested for MLA backend."
273 )
274
273 - if selected_backend == _Backend.FLEX_ATTENTION:
275 + if selected_backend == AttentionBackendEnum.FLEX_ATTENTION:
276 logger.info("Using FlexAttention backend.")
277 return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
278 if (
279 rocm_aiter_ops.is_mha_enabled()
278 - ) or selected_backend == _Backend.ROCM_AITER_FA:
280 + ) or selected_backend == AttentionBackendEnum.ROCM_AITER_FA:
281 logger.info("Using Aiter Flash Attention backend.")
280 - return "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
282 + return AttentionBackendEnum.ROCM_AITER_FA.get_path()
283 if (
284 rocm_aiter_ops.is_triton_unified_attn_enabled()
283 - ) or selected_backend == _Backend.ROCM_AITER_UNIFIED_ATTN:
285 + ) or selected_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
286 logger.info("Using Aiter Unified Attention backend.")
285 - return (
286 - "vllm.v1.attention.backends."
287 - "rocm_aiter_unified_attn.RocmAiterUnifiedAttentionBackend"
288 - )
287 + return AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
288 if (
289 envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
291 - or selected_backend == _Backend.ROCM_ATTN
290 + or selected_backend == AttentionBackendEnum.ROCM_ATTN
291 ):
292 # rocm specific backend, with aiter and/or
293 # triton prefix-prefill
294 logger.info("Using Rocm Attention backend.")
296 - return "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
295 + return AttentionBackendEnum.ROCM_ATTN.get_path()
296 # default case, using triton unified attention
297 logger.info("Using Triton Attention backend.")
299 - return "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
298 + return AttentionBackendEnum.TRITON_ATTN.get_path()
299
300 @classmethod
301 def set_device(cls, device: torch.device) -> None:
vllm/platforms/tpu.py
+7 -8
@@ -15,16 +15,15 @@ from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
15 from .interface import Platform, PlatformEnum
16
17 if TYPE_CHECKING:
18 - from vllm.attention.backends.registry import _Backend
19 - from vllm.config import ModelConfig, VllmConfig
18 + from vllm.attention.backends.registry import AttentionBackendEnum
19 + from vllm.config import VllmConfig
20 from vllm.config.cache import BlockSize
21 from vllm.pooling_params import PoolingParams
22 else:
23 BlockSize = None
24 - ModelConfig = None
24 VllmConfig = None
25 PoolingParams = None
27 - _Backend = None
26 + AttentionBackendEnum = None
27
28 logger = init_logger(__name__)
29
@@ -54,7 +53,7 @@ class TpuPlatform(Platform):
53 @classmethod
54 def get_attn_backend_cls(
55 cls,
57 - selected_backend: "_Backend",
56 + selected_backend: "AttentionBackendEnum",
57 head_size: int,
58 dtype: torch.dtype,
59 kv_cache_dtype: str | None,
@@ -64,17 +63,17 @@ class TpuPlatform(Platform):
63 has_sink,
64 use_sparse,
65 ) -> str:
67 - from vllm.attention.backends.registry import _Backend
66 + from vllm.attention.backends.registry import AttentionBackendEnum
67
68 if use_sparse:
69 raise NotImplementedError("Sparse Attention is not supported on TPU.")
71 - if selected_backend != _Backend.PALLAS:
70 + if selected_backend != AttentionBackendEnum.PALLAS:
71 logger.info("Cannot use %s backend on TPU.", selected_backend)
72
73 if not use_v1:
74 raise ValueError("TPU backend only supports V1.")
75 logger.info("Using Pallas V1 backend.")
77 - return "vllm.v1.attention.backends.pallas.PallasAttentionBackend"
76 + return AttentionBackendEnum.PALLAS.get_path()
77
78 @classmethod
79 def set_device(cls, device: torch.device) -> None:
vllm/platforms/xpu.py
+17 -17
@@ -14,12 +14,11 @@ from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
14 from .interface import DeviceCapability, Platform, PlatformEnum
15
16 if TYPE_CHECKING:
17 - from vllm.attention.backends.registry import _Backend
18 - from vllm.config import ModelConfig, VllmConfig
17 + from vllm.attention.backends.registry import AttentionBackendEnum
18 + from vllm.config import VllmConfig
19 else:
20 - ModelConfig = None
20 VllmConfig = None
22 - _Backend = None
21 + AttentionBackendEnum = None
22
23 logger = init_logger(__name__)
24
@@ -44,7 +43,7 @@ class XPUPlatform(Platform):
43 @classmethod
44 def get_attn_backend_cls(
45 cls,
47 - selected_backend: "_Backend",
46 + selected_backend: "AttentionBackendEnum",
47 head_size: int,
48 dtype: torch.dtype,
49 kv_cache_dtype: str | None,
@@ -62,18 +61,19 @@ class XPUPlatform(Platform):
61 "only NHD layout is supported by XPU attention kernels."
62 )
63
65 - from vllm.attention.backends.registry import _Backend
64 + from vllm.attention.backends.registry import AttentionBackendEnum
65
66 if use_sparse:
67 raise NotImplementedError("Sparse Attention is not supported on XPU.")
69 - TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" # noqa: E501
70 - FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
71 - if selected_backend == _Backend.TRITON_ATTN:
68 + use_v1 = envs.VLLM_USE_V1
69 + if not use_v1:
70 + raise ValueError("XPU backend only supports V1.")
71 + if selected_backend == AttentionBackendEnum.TRITON_ATTN:
72 logger.info_once("Using Triton backend.")
73 - return TRITON_ATTN
74 - elif selected_backend == _Backend.FLASH_ATTN:
73 + return AttentionBackendEnum.TRITON_ATTN.get_path()
74 + elif selected_backend == AttentionBackendEnum.FLASH_ATTN:
75 logger.info_once("Using Flash Attention backend.")
76 - return FLASH_ATTN
76 + return AttentionBackendEnum.FLASH_ATTN.get_path()
77 elif selected_backend:
78 raise ValueError(
79 f"Invalid attention backend for {cls.device_name}, "
@@ -81,7 +81,7 @@ class XPUPlatform(Platform):
81 )
82
83 logger.info("Using Flash Attention backend.")
84 - return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
84 + return AttentionBackendEnum.FLASH_ATTN.get_path()
85
86 @classmethod
87 def set_device(cls, device: torch.device) -> None:
@@ -113,10 +113,10 @@ class XPUPlatform(Platform):
113 return device_props.total_memory
114
115 @classmethod
116 - def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> _Backend:
117 - from vllm.attention.backends.registry import _Backend
118 -
119 - return _Backend.FLASH_ATTN
116 + def get_vit_attn_backend(
117 + cls, head_size: int, dtype: torch.dtype
118 + ) -> AttentionBackendEnum:
119 + return AttentionBackendEnum.FLASH_ATTN
120
121 @classmethod
122 def inference_mode(cls):
vllm/v1/attention/backends/cpu_attn.py
+12 -20
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 -from typing import Optional
4 +from typing import ClassVar, Optional
5
6 import numpy as np
7 import torch
@@ -40,23 +40,16 @@ logger = init_logger(__name__)
40
41 class TorchSDPABackend(AttentionBackend):
42 accept_output_buffer: bool = False
43 + supported_dtypes: ClassVar[list[torch.dtype]] = [
44 + torch.float16,
45 + torch.bfloat16,
46 + torch.float32,
47 + ]
48
49 @classmethod
45 - def get_supported_dtypes(cls) -> list[torch.dtype]:
46 - return [torch.float16, torch.bfloat16, torch.float32]
47 -
48 - @classmethod
49 - def validate_head_size(cls, head_size: int) -> None:
50 + def get_supported_head_sizes(cls) -> list[int]:
51 attn_impl = _get_paged_attn_impl()
51 - is_valid, supported_head_sizes = attn_impl.validate_head_size(head_size)
52 - if not is_valid:
53 - attn_type = cls.__name__.removesuffix("Backend")
54 - raise ValueError(
55 - f"Head size {head_size} is not supported by {attn_type}. "
56 - f"Supported head sizes are: {supported_head_sizes}. "
57 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
58 - "FlexAttention backend which supports all head sizes."
59 - )
52 + return attn_impl.get_supported_head_sizes()
53
54 @staticmethod
55 def get_name() -> str:
@@ -759,9 +752,8 @@ def _make_sliding_window_bias(
752
753 class _PagedAttention:
754 @staticmethod
762 - def validate_head_size(head_size: int) -> tuple[bool, list[int]]:
763 - SUPPORT_HS = [32, 64, 80, 96, 112, 128, 192, 256]
764 - return head_size in SUPPORT_HS, SUPPORT_HS
755 + def get_supported_head_sizes() -> list[int]:
756 + return [32, 64, 80, 96, 112, 128, 192, 256]
757
758 @staticmethod
759 def get_kv_cache_shape(
@@ -861,8 +853,8 @@ class _PagedAttention:
853
854 class _IPEXPagedAttention(_PagedAttention):
855 @staticmethod
864 - def validate_head_size(head_size: int) -> tuple[bool, list[int]]:
865 - return True, []
856 + def get_supported_head_sizes() -> list[int]:
857 + return []
858
859 @staticmethod
860 def split_kv_cache(
vllm/v1/attention/backends/flash_attn.py
+41 -30
@@ -3,6 +3,7 @@
3 """Attention layer with FlashAttention."""
4
5 from dataclasses import dataclass
6 +from typing import ClassVar
7
8 import numpy as np
9 import torch
@@ -32,11 +33,13 @@ if is_flash_attn_varlen_func_available():
33 reshape_and_cache_flash,
34 )
35 from vllm.config import VllmConfig, get_layers_from_vllm_config
36 +from vllm.config.cache import CacheDType
37 from vllm.distributed.parallel_state import get_dcp_group
38 from vllm.logger import init_logger
39 from vllm.model_executor.layers.batch_invariant import (
40 vllm_is_batch_invariant,
41 )
42 +from vllm.platforms.interface import DeviceCapability
43 from vllm.utils.math_utils import cdiv
44 from vllm.v1.attention.backends.utils import (
45 AttentionCGSupport,
@@ -52,34 +55,12 @@ logger = init_logger(__name__)
55
56 class FlashAttentionBackend(AttentionBackend):
57 accept_output_buffer: bool = True
55 -
56 - @classmethod
57 - def get_supported_dtypes(cls) -> list[torch.dtype]:
58 - return [torch.float16, torch.bfloat16]
59 -
60 - @classmethod
61 - def get_supported_head_sizes(cls) -> list[int]:
62 - return [32, 64, 96, 128, 160, 192, 224, 256]
63 -
64 - @staticmethod
65 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
66 - # NOTE(tdoublep): while in principle, FA supports
67 - # MultipleOf(16), these are the block sizes that do not
68 - # suffer from the NaN propagation problem described here:
69 - # https://github.com/Dao-AILab/flash-attention/issues/1974
70 - return [16, 32, 64]
71 -
72 - @classmethod
73 - def validate_head_size(cls, head_size: int) -> None:
74 - supported_head_sizes = cls.get_supported_head_sizes()
75 - if head_size not in supported_head_sizes:
76 - attn_type = cls.__name__.removesuffix("Backend")
77 - raise ValueError(
78 - f"Head size {head_size} is not supported by {attn_type}. "
79 - f"Supported head sizes are: {supported_head_sizes}. "
80 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
81 - "FlexAttention backend which supports all head sizes."
82 - )
58 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
59 + # NOTE(tdoublep): while in principle, FA supports
60 + # MultipleOf(16), these are the block sizes that do not
61 + # suffer from the NaN propagation problem described here:
62 + # https://github.com/Dao-AILab/flash-attention/issues/1974
63 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [16, 32, 64]
64
65 @staticmethod
66 def get_name() -> str:
@@ -125,6 +106,38 @@ class FlashAttentionBackend(AttentionBackend):
106 else:
107 raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
108
109 + @classmethod
110 + def get_supported_head_sizes(cls) -> list[int]:
111 + return [32, 64, 96, 128, 160, 192, 224, 256]
112 +
113 + @classmethod
114 + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
115 + if kv_cache_dtype is None:
116 + return True
117 + if kv_cache_dtype.startswith("fp8"):
118 + return flash_attn_supports_fp8()
119 + return kv_cache_dtype in ["auto"]
120 +
121 + @classmethod
122 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
123 + return capability >= DeviceCapability(8, 0)
124 +
125 + @classmethod
126 + def supports_combination(
127 + cls,
128 + head_size: int,
129 + dtype: torch.dtype,
130 + kv_cache_dtype: CacheDType | None,
131 + block_size: int,
132 + use_mla: bool,
133 + has_sink: bool,
134 + use_sparse: bool,
135 + device_capability: DeviceCapability,
136 + ) -> str | None:
137 + if has_sink and device_capability < DeviceCapability(9, 0):
138 + return "sink not supported on compute capability < 9.0"
139 + return None
140 +
141
142 @dataclass
143 class FlashAttentionMetadata:
@@ -481,8 +494,6 @@ class FlashAttentionImpl(AttentionImpl):
494
495 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
496
484 - FlashAttentionBackend.validate_head_size(head_size)
485 -
497 self.attn_type = attn_type
498 self.vllm_flash_attn_version = get_flash_attn_version()
499 # Cache the batch invariant result for use in forward passes
vllm/v1/attention/backends/flashinfer.py
+34 -29
@@ -23,6 +23,7 @@ from vllm.attention.backends.abstract import (
23 MultipleOf,
24 )
25 from vllm.config import CUDAGraphMode, VllmConfig
26 +from vllm.config.cache import CacheDType
27 from vllm.logger import init_logger
28 from vllm.model_executor.layers.batch_invariant import (
29 vllm_is_batch_invariant,
@@ -33,6 +34,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
34 kNvfp4Quant,
35 )
36 from vllm.platforms import current_platform
37 +from vllm.platforms.interface import DeviceCapability
38 from vllm.triton_utils import tl, triton
39 from vllm.utils.flashinfer import (
40 can_use_trtllm_attention,
@@ -45,6 +47,7 @@ from vllm.v1.attention.backends.utils import (
47 AttentionCGSupport,
48 AttentionMetadataBuilder,
49 CommonAttentionMetadata,
50 + KVCacheLayoutType,
51 get_kv_cache_layout,
52 get_per_layer_parameters,
53 infer_global_hyperparameters,
@@ -158,34 +161,17 @@ def trtllm_prefill_attn_kvfp8_dequant(
161
162 class FlashInferBackend(AttentionBackend):
163 accept_output_buffer: bool = True
161 -
162 - @classmethod
163 - def get_supported_dtypes(cls) -> list[torch.dtype]:
164 - return [torch.float16, torch.bfloat16]
165 -
166 - @classmethod
167 - def get_supported_head_sizes(cls) -> list[int]:
168 - # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
169 - return [64, 128, 256]
170 -
171 - @staticmethod
172 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
173 - # Note: Not sure for all platforms,
174 - # but on Blackwell, only support a page size of
175 - # 16, 32, 64
176 - return [16, 32, 64]
177 -
178 - @classmethod
179 - def validate_head_size(cls, head_size: int) -> None:
180 - supported_head_sizes = cls.get_supported_head_sizes()
181 - if head_size not in supported_head_sizes:
182 - attn_type = cls.__name__.removesuffix("Backend")
183 - raise ValueError(
184 - f"Head size {head_size} is not supported by {attn_type}. "
185 - f"Supported head sizes are: {supported_head_sizes}. "
186 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
187 - "FlexAttention backend which supports all head sizes."
188 - )
164 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
165 + # Note: Not sure for all platforms,
166 + # but on Blackwell, only support a page size of
167 + # 16, 32, 64
168 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [16, 32, 64]
169 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
170 + "auto",
171 + "fp8",
172 + "fp8_e4m3",
173 + "fp8_e5m2",
174 + ]
175
176 @staticmethod
177 def get_name() -> str:
@@ -231,6 +217,26 @@ class FlashInferBackend(AttentionBackend):
217 else:
218 raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
219
220 + @classmethod
221 + def get_supported_head_sizes(cls) -> list[int]:
222 + # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
223 + return [64, 128, 256]
224 +
225 + @classmethod
226 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
227 + return capability >= DeviceCapability(7, 5) and capability <= DeviceCapability(
228 + 12, 1
229 + )
230 +
231 + @classmethod
232 + def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
233 + from vllm.platforms import current_platform
234 +
235 + capability = current_platform.get_device_capability()
236 + if capability is not None and capability.major == 10:
237 + return "HND"
238 + return None
239 +
240
241 @dataclass
242 class FlashInferMetadata:
@@ -328,7 +334,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
334 )
335 self.num_kv_heads = self.kv_cache_spec.num_kv_heads
336 self.head_dim = self.kv_cache_spec.head_size
331 - FlashInferBackend.validate_head_size(self.head_dim)
337 self.page_size = self.kv_cache_spec.block_size
338
339 self.cache_dtype = self.cache_config.cache_dtype
vllm/v1/attention/backends/flex_attention.py
+12 -9
@@ -4,6 +4,7 @@
4
5 import math
6 from dataclasses import dataclass
7 +from typing import ClassVar
8
9 import torch
10 import torch._dynamo.decorators
@@ -24,6 +25,7 @@ from vllm.attention.backends.abstract import (
25 is_quantized_kv_cache,
26 )
27 from vllm.config import VllmConfig
28 +from vllm.config.cache import CacheDType
29 from vllm.logger import init_logger
30 from vllm.model_executor.layers.batch_invariant import (
31 vllm_is_batch_invariant,
@@ -71,14 +73,12 @@ def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int):
73
74 class FlexAttentionBackend(AttentionBackend):
75 accept_output_buffer: bool = True
74 -
75 - @classmethod
76 - def get_supported_dtypes(cls) -> list[torch.dtype]:
77 - return [torch.float16, torch.bfloat16, torch.float32]
78 -
79 - @classmethod
80 - def validate_head_size(cls, head_size: int) -> None:
81 - return # FlexAttention supports any head size
76 + supported_dtypes: ClassVar[list[torch.dtype]] = [
77 + torch.float16,
78 + torch.bfloat16,
79 + torch.float32,
80 + ]
81 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
82
83 @staticmethod
84 def get_name() -> str:
@@ -106,6 +106,10 @@ class FlexAttentionBackend(AttentionBackend):
106 def use_cascade_attention(*args, **kwargs) -> bool:
107 return False
108
109 + @classmethod
110 + def get_supported_head_sizes(cls) -> list[int]:
111 + return []
112 +
113
114 # @torch.compile(fullgraph=True, mode="reduce-overhead")
115 def physical_to_logical_mapping(
@@ -720,7 +724,6 @@ class FlexAttentionImpl(AttentionImpl):
724 if kv_sharing_target_layer_name is not None:
725 raise NotImplementedError("FlexAttention does not support kv sharing yet.")
726
723 - FlexAttentionBackend.validate_head_size(head_size)
727 if is_quantized_kv_cache(self.kv_cache_dtype):
728 raise NotImplementedError(
729 "FlexAttention does not support quantized kv-cache. Yet"
vllm/v1/attention/backends/mla/common.py
+6 -16
@@ -308,25 +308,13 @@ class MLACommonBackend(AttentionBackend):
308 ) -> tuple[int, ...]:
309 return (num_blocks, block_size, head_size)
310
311 - @classmethod
312 - def get_supported_dtypes(cls) -> list[torch.dtype]:
313 - return [torch.float16, torch.bfloat16]
314 -
311 @classmethod
312 def get_supported_head_sizes(cls) -> list[int]:
313 return [576]
314
315 @classmethod
320 - def validate_head_size(cls, head_size: int) -> None:
321 - supported_head_sizes = cls.get_supported_head_sizes()
322 - if head_size not in supported_head_sizes:
323 - attn_type = cls.__name__.removesuffix("Backend")
324 - raise ValueError(
325 - f"Head size {head_size} is not supported by {attn_type}. "
326 - f"Supported head sizes are: {supported_head_sizes}. "
327 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
328 - "FlexAttention backend which supports all head sizes."
329 - )
316 + def is_mla(cls) -> bool:
317 + return True
318
319
320 @dataclass
@@ -425,8 +413,10 @@ class MLACommonMetadata(Generic[D]):
413 ) = None
414
415 def __post_init__(self):
428 - if self.head_dim is not None:
429 - MLACommonBackend.validate_head_size(self.head_dim)
416 + if self.head_dim is not None and not MLACommonBackend.supports_head_size(
417 + self.head_dim
418 + ):
419 + raise ValueError(f"Head dimension {self.head_dim} is not supported by MLA.")
420
421
422 M = TypeVar("M", bound=MLACommonMetadata)
vllm/v1/attention/backends/mla/cutlass_mla.py
+13 -3
@@ -13,7 +13,9 @@ from vllm.attention.backends.abstract import (
13 MultipleOf,
14 is_quantized_kv_cache,
15 )
16 +from vllm.config.cache import CacheDType
17 from vllm.logger import init_logger
18 +from vllm.platforms.interface import DeviceCapability
19 from vllm.v1.attention.backends.mla.common import (
20 MLACommonBackend,
21 MLACommonImpl,
@@ -33,6 +35,14 @@ class CutlassMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
35
36
37 class CutlassMLABackend(MLACommonBackend):
38 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
39 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [128]
40 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
41 + "auto",
42 + "fp8",
43 + "fp8_e4m3",
44 + ]
45 +
46 @staticmethod
47 def get_name() -> str:
48 return "CUTLASS_MLA"
@@ -45,9 +55,9 @@ class CutlassMLABackend(MLACommonBackend):
55 def get_builder_cls() -> type["CutlassMLAMetadataBuilder"]:
56 return CutlassMLAMetadataBuilder
57
48 - @staticmethod
49 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
50 - return [128]
58 + @classmethod
59 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
60 + return capability.major == 10
61
62
63 class SM100Workspace:
vllm/v1/attention/backends/mla/flashattn_mla.py
+27
@@ -10,6 +10,7 @@ from vllm import envs
10 from vllm.attention.backends.abstract import (
11 AttentionLayer,
12 AttentionType,
13 + MultipleOf,
14 is_quantized_kv_cache,
15 )
16 from vllm.attention.utils.fa_utils import (
@@ -17,10 +18,12 @@ from vllm.attention.utils.fa_utils import (
18 get_flash_attn_version,
19 )
20 from vllm.config import VllmConfig
21 +from vllm.config.cache import CacheDType
22 from vllm.logger import init_logger
23 from vllm.model_executor.layers.batch_invariant import (
24 vllm_is_batch_invariant,
25 )
26 +from vllm.platforms.interface import DeviceCapability
27 from vllm.v1.attention.backends.mla.common import (
28 MLACommonBackend,
29 MLACommonDecodeMetadata,
@@ -37,6 +40,10 @@ logger = init_logger(__name__)
40
41
42 class FlashAttnMLABackend(MLACommonBackend):
43 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
44 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
45 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
46 +
47 @staticmethod
48 def get_name() -> str:
49 return "FLASH_ATTN_MLA"
@@ -49,6 +56,26 @@ class FlashAttnMLABackend(MLACommonBackend):
56 def get_impl_cls() -> type["FlashAttnMLAImpl"]:
57 return FlashAttnMLAImpl
58
59 + @classmethod
60 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
61 + return capability.major == 9
62 +
63 + @classmethod
64 + def supports_combination(
65 + cls,
66 + head_size: int,
67 + dtype: torch.dtype,
68 + kv_cache_dtype: CacheDType | None,
69 + block_size: int,
70 + use_mla: bool,
71 + has_sink: bool,
72 + use_sparse: bool,
73 + device_capability: DeviceCapability,
74 + ) -> str | None:
75 + if not flash_attn_supports_mla():
76 + return "FlashAttention MLA not supported on this device"
77 + return None
78 +
79
80 @dataclass
81 class FlashAttnMLADecodeMetadata(MLACommonDecodeMetadata):
vllm/v1/attention/backends/mla/flashinfer_mla.py
+22 -4
@@ -6,8 +6,14 @@ from typing import ClassVar
6 import torch
7 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
8
9 -from vllm.attention.backends.abstract import AttentionLayer, AttentionType, MultipleOf
9 +from vllm.attention.backends.abstract import (
10 + AttentionLayer,
11 + AttentionType,
12 + MultipleOf,
13 +)
14 +from vllm.config.cache import CacheDType
15 from vllm.logger import init_logger
16 +from vllm.platforms.interface import DeviceCapability
17 from vllm.v1.attention.backends.mla.common import (
18 MLACommonBackend,
19 MLACommonImpl,
@@ -15,7 +21,7 @@ from vllm.v1.attention.backends.mla.common import (
21 MLACommonMetadataBuilder,
22 QueryLenSupport,
23 )
18 -from vllm.v1.attention.backends.utils import AttentionCGSupport
24 +from vllm.v1.attention.backends.utils import AttentionCGSupport, KVCacheLayoutType
25
26 logger = init_logger(__name__)
27
@@ -28,6 +34,14 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
34
35
36 class FlashInferMLABackend(MLACommonBackend):
37 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
38 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [32, 64]
39 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
40 + "auto",
41 + "fp8",
42 + "fp8_e4m3",
43 + ]
44 +
45 @staticmethod
46 def get_name() -> str:
47 return "FLASHINFER_MLA"
@@ -41,8 +55,12 @@ class FlashInferMLABackend(MLACommonBackend):
55 return FlashInferMLAMetadataBuilder
56
57 @classmethod
44 - def get_supported_kernel_block_size(cls) -> list[int | MultipleOf]:
45 - return [32, 64]
58 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
59 + return capability.major == 10
60 +
61 + @classmethod
62 + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
63 + return "HND"
64
65
66 g_fi_workspace = torch.zeros(
vllm/v1/attention/backends/mla/flashmla.py
+34 -3
@@ -13,10 +13,12 @@ from vllm.attention.ops.flashmla import (
13 is_flashmla_dense_supported,
14 )
15 from vllm.config import VllmConfig
16 +from vllm.config.cache import CacheDType
17 from vllm.logger import init_logger
18 from vllm.model_executor.layers.batch_invariant import (
19 vllm_is_batch_invariant,
20 )
21 +from vllm.platforms.interface import DeviceCapability
22 from vllm.v1.attention.backends.mla.common import (
23 MLACommonBackend,
24 MLACommonDecodeMetadata,
@@ -36,6 +38,14 @@ logger = init_logger(__name__)
38
39
40 class FlashMLABackend(MLACommonBackend):
41 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
42 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
43 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
44 + "auto",
45 + "fp8",
46 + "fp8_e4m3",
47 + ]
48 +
49 @staticmethod
50 def get_name() -> str:
51 return "FLASHMLA"
@@ -48,9 +58,30 @@ class FlashMLABackend(MLACommonBackend):
58 def get_impl_cls() -> type["FlashMLAImpl"]:
59 return FlashMLAImpl
60
51 - @staticmethod
52 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
53 - return [64]
61 + @classmethod
62 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
63 + return capability.major in [9, 10]
64 +
65 + @classmethod
66 + def supports_combination(
67 + cls,
68 + head_size: int,
69 + dtype: torch.dtype,
70 + kv_cache_dtype: CacheDType | None,
71 + block_size: int,
72 + use_mla: bool,
73 + has_sink: bool,
74 + use_sparse: bool,
75 + device_capability: DeviceCapability,
76 + ) -> str | None:
77 + if use_sparse:
78 + from vllm.attention.ops.flashmla import is_flashmla_sparse_supported
79 +
80 + return is_flashmla_sparse_supported()[1]
81 + else:
82 + from vllm.attention.ops.flashmla import is_flashmla_dense_supported
83 +
84 + return is_flashmla_dense_supported()[1]
85
86
87 @dataclass
vllm/v1/attention/backends/mla/flashmla_sparse.py
+22 -8
@@ -10,6 +10,7 @@ from vllm import _custom_ops as ops
10 from vllm.attention.backends.abstract import (
11 AttentionBackend,
12 AttentionLayer,
13 + MultipleOf,
14 )
15 from vllm.attention.backends.utils import get_mla_dims
16 from vllm.attention.ops.flashmla import (
@@ -18,8 +19,10 @@ from vllm.attention.ops.flashmla import (
19 get_mla_metadata,
20 )
21 from vllm.config import VllmConfig
22 +from vllm.config.cache import CacheDType
23 from vllm.logger import init_logger
24 from vllm.platforms import current_platform
25 +from vllm.platforms.interface import DeviceCapability
26 from vllm.triton_utils import tl, triton
27 from vllm.utils.math_utils import cdiv
28 from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl
@@ -51,6 +54,9 @@ structured as:
54
55 class FlashMLASparseBackend(AttentionBackend):
56 accept_output_buffer: bool = True
57 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
58 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
59 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "fp8_ds_mla"]
60
61 @staticmethod
62 def get_name() -> str:
@@ -64,6 +70,22 @@ class FlashMLASparseBackend(AttentionBackend):
70 def get_impl_cls() -> type["FlashMLASparseImpl"]:
71 return FlashMLASparseImpl
72
73 + @classmethod
74 + def get_supported_head_sizes(cls) -> list[int]:
75 + return [576]
76 +
77 + @classmethod
78 + def is_mla(cls) -> bool:
79 + return True
80 +
81 + @classmethod
82 + def is_sparse(cls) -> bool:
83 + return True
84 +
85 + @classmethod
86 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
87 + return capability.major in [9, 10]
88 +
89 @staticmethod
90 def get_kv_cache_shape(
91 num_blocks: int,
@@ -79,14 +101,6 @@ class FlashMLASparseBackend(AttentionBackend):
101 else:
102 return (num_blocks, block_size, head_size)
103
82 - @classmethod
83 - def get_supported_dtypes(cls) -> list[torch.dtype]:
84 - return [torch.bfloat16]
85 -
86 - @classmethod
87 - def get_supported_head_sizes(cls) -> list[int]:
88 - return [576]
89 -
104
105 @dataclass
106 class FlashMLASparseMetadata:
vllm/v1/attention/backends/mla/indexer.py
+2 -4
@@ -23,6 +23,8 @@ logger = init_logger(__name__)
23
24
25 class DeepseekV32IndexerBackend(AttentionBackend):
26 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
27 +
28 @classmethod
29 def get_supported_head_sizes(cls) -> list[int]:
30 return [32, 64, 128]
@@ -46,10 +48,6 @@ class DeepseekV32IndexerBackend(AttentionBackend):
48 def get_kv_cache_stride_order() -> tuple[int, ...]:
49 return (0, 1, 2)
50
49 - @classmethod
50 - def get_supported_kernel_block_size(cls) -> list[int | MultipleOf]:
51 - return [64]
52 -
51
52 @dataclass
53 class DeepseekV32IndexerPrefillChunkMetadata:
vllm/v1/attention/backends/mla/triton_mla.py
+10
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from typing import ClassVar
5
6 import torch
7
@@ -12,11 +13,13 @@ from vllm.attention.backends.abstract import (
13 )
14 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
15 from vllm.attention.ops.triton_flash_attention import triton_attention
16 +from vllm.config.cache import CacheDType
17 from vllm.logger import init_logger
18 from vllm.model_executor.layers.batch_invariant import (
19 vllm_is_batch_invariant,
20 )
21 from vllm.platforms import current_platform
22 +from vllm.platforms.interface import DeviceCapability
23 from vllm.triton_utils import HAS_TRITON
24 from vllm.v1.attention.backends.mla.common import (
25 MLACommonBackend,
@@ -28,6 +31,9 @@ logger = init_logger(__name__)
31
32
33 class TritonMLABackend(MLACommonBackend):
34 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
35 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
36 +
37 @staticmethod
38 def get_name() -> str:
39 return "TRITON_MLA"
@@ -36,6 +42,10 @@ class TritonMLABackend(MLACommonBackend):
42 def get_impl_cls() -> type["TritonMLAImpl"]:
43 return TritonMLAImpl
44
45 + @classmethod
46 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
47 + return True
48 +
49
50 class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
51 can_return_lse_for_decode: bool = True
vllm/v1/attention/backends/rocm_aiter_fa.py
+3 -22
@@ -3,6 +3,7 @@
3 """Attention layer with AiterFlashAttention."""
4
5 from dataclasses import dataclass
6 +from typing import ClassVar
7
8 import torch
9
@@ -445,31 +446,13 @@ class AiterFlashAttentionMetadataBuilder(
446
447 class AiterFlashAttentionBackend(AttentionBackend):
448 accept_output_buffer: bool = True
448 -
449 - @classmethod
450 - def get_supported_dtypes(cls) -> list[torch.dtype]:
451 - return [torch.float16, torch.bfloat16]
449 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
450 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
451
452 @classmethod
453 def get_supported_head_sizes(cls) -> list[int]:
454 return [64, 128, 256]
455
457 - @staticmethod
458 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
459 - return [MultipleOf(16)]
460 -
461 - @classmethod
462 - def validate_head_size(cls, head_size: int) -> None:
463 - supported_head_sizes = cls.get_supported_head_sizes()
464 - if head_size not in supported_head_sizes:
465 - attn_type = cls.__name__.removesuffix("Backend")
466 - raise ValueError(
467 - f"Head size {head_size} is not supported by {attn_type}. "
468 - f"Supported head sizes are: {supported_head_sizes}. "
469 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
470 - "FlexAttention backend which supports all head sizes."
471 - )
472 -
456 @staticmethod
457 def get_name() -> str:
458 return "FLASH_ATTN"
@@ -531,8 +514,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
514 assert self.num_heads % self.num_kv_heads == 0
515 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
516
534 - AiterFlashAttentionBackend.validate_head_size(head_size)
535 -
517 if attn_type != AttentionType.DECODER:
518 raise NotImplementedError(
519 "Encoder self-attention and "
vllm/v1/attention/backends/rocm_attn.py
+3 -7
@@ -152,10 +152,7 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat
152
153 class RocmAttentionBackend(AttentionBackend):
154 accept_output_buffer: bool = True
155 -
156 - @classmethod
157 - def get_supported_dtypes(cls) -> list[torch.dtype]:
158 - return [torch.float16, torch.bfloat16]
155 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
156
157 @classmethod
158 def get_supported_head_sizes(cls) -> list[int]:
@@ -163,12 +160,11 @@ class RocmAttentionBackend(AttentionBackend):
160
161 @classmethod
162 def validate_head_size(cls, head_size: int) -> None:
166 - supported_head_sizes = cls.get_supported_head_sizes()
167 - if head_size not in supported_head_sizes:
163 + if not cls.supports_head_size(head_size):
164 attn_type = cls.__name__.removesuffix("Backend")
165 raise ValueError(
166 f"Head size {head_size} is not supported by {attn_type}. "
171 - f"Supported head sizes are: {supported_head_sizes}. "
167 + f"Supported head sizes are: {cls.get_supported_head_sizes()}. "
168 "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
169 "FlexAttention backend which supports all head sizes."
170 )
vllm/v1/attention/backends/tree_attn.py
+3 -23
@@ -4,7 +4,7 @@
4
5 import ast
6 from dataclasses import dataclass
7 -from typing import Optional
7 +from typing import ClassVar, Optional
8
9 import torch
10
@@ -30,31 +30,13 @@ logger = init_logger(__name__)
30
31 class TreeAttentionBackend(AttentionBackend):
32 accept_output_buffer: bool = True
33 -
34 - @classmethod
35 - def get_supported_dtypes(cls) -> list[torch.dtype]:
36 - return [torch.float16, torch.bfloat16]
33 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
34 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
35
36 @classmethod
37 def get_supported_head_sizes(cls) -> list[int]:
38 return [32, 64, 96, 128, 160, 192, 224, 256]
39
42 - @staticmethod
43 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
44 - return [MultipleOf(16)]
45 -
46 - @classmethod
47 - def validate_head_size(cls, head_size: int) -> None:
48 - supported_head_sizes = cls.get_supported_head_sizes()
49 - if head_size not in supported_head_sizes:
50 - attn_type = cls.__name__.removesuffix("Backend")
51 - raise ValueError(
52 - f"Head size {head_size} is not supported by {attn_type}. "
53 - f"Supported head sizes are: {supported_head_sizes}. "
54 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
55 - "FlexAttention backend which supports all head sizes."
56 - )
57 -
40 @staticmethod
41 def get_name() -> str:
42 return "TREE_ATTN"
@@ -331,8 +313,6 @@ class TreeAttentionImpl(AttentionImpl):
313 else:
314 self.sliding_window = (sliding_window - 1, 0)
315
334 - TreeAttentionBackend.validate_head_size(head_size)
335 -
316 if attn_type != AttentionType.DECODER:
317 raise NotImplementedError(
318 "Encoder self-attention and "
vllm/v1/attention/backends/triton_attn.py
+26 -21
@@ -18,12 +18,14 @@ from vllm.attention.ops.triton_reshape_and_cache_flash import (
18 )
19 from vllm.attention.ops.triton_unified_attention import unified_attention
20 from vllm.config import VllmConfig
21 +from vllm.config.cache import CacheDType
22 from vllm.logger import init_logger
23 from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 QuantKey,
25 kFp8StaticTensorSym,
26 )
27 from vllm.platforms import current_platform
28 +from vllm.platforms.interface import DeviceCapability
29 from vllm.v1.attention.backends.utils import (
30 AttentionCGSupport,
31 AttentionMetadataBuilder,
@@ -147,25 +149,18 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
149
150 class TritonAttentionBackend(AttentionBackend):
151 accept_output_buffer: bool = True
150 -
151 - @classmethod
152 - def get_supported_dtypes(cls) -> list[torch.dtype]:
153 - return [torch.float16, torch.bfloat16, torch.float32]
154 -
155 - @staticmethod
156 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
157 - return [MultipleOf(16)]
158 -
159 - @classmethod
160 - def validate_head_size(cls, head_size: int) -> None:
161 - # Triton Attention supports any head size above 32
162 - if head_size < 32:
163 - raise ValueError(
164 - f"Head size {head_size} is not supported by TritonAttention."
165 - f"Head sizes need to be larger or equal 32 for this backend. "
166 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
167 - "FlexAttention backend which supports all head sizes."
168 - )
152 + supported_dtypes: ClassVar[list[torch.dtype]] = [
153 + torch.float16,
154 + torch.bfloat16,
155 + torch.float32,
156 + ]
157 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
158 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
159 + "auto",
160 + "fp8",
161 + "fp8_e4m3",
162 + "fp8_e5m2",
163 + ]
164
165 @staticmethod
166 def get_name() -> str:
@@ -195,6 +190,18 @@ class TritonAttentionBackend(AttentionBackend):
190 def get_builder_cls() -> type["TritonAttentionMetadataBuilder"]:
191 return TritonAttentionMetadataBuilder
192
193 + @classmethod
194 + def supports_head_size(cls, head_size: int) -> bool:
195 + return head_size >= 32
196 +
197 + @classmethod
198 + def supports_sink(cls) -> bool:
199 + return True
200 +
201 + @classmethod
202 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
203 + return True
204 +
205
206 class TritonAttentionImpl(AttentionImpl):
207 def fused_output_quant_supported(self, quant_key: QuantKey):
@@ -237,8 +244,6 @@ class TritonAttentionImpl(AttentionImpl):
244
245 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
246
240 - TritonAttentionBackend.validate_head_size(head_size)
241 -
247 if attn_type != AttentionType.DECODER:
248 raise NotImplementedError(
249 "Encoder self-attention and "
vllm/v1/attention/backends/xformers.py
+3 -23
@@ -3,7 +3,7 @@
3 """Attention layer with XFormersAttention."""
4
5 from dataclasses import dataclass
6 -from typing import Optional
6 +from typing import ClassVar, Optional
7
8 import torch
9
@@ -41,10 +41,8 @@ logger = init_logger(__name__)
41
42 class XFormersAttentionBackend(AttentionBackend):
43 accept_output_buffer: bool = True
44 -
45 - @classmethod
46 - def get_supported_dtypes(cls) -> list[torch.dtype]:
47 - return [torch.float16, torch.bfloat16]
44 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
45 + supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]
46
47 @classmethod
48 def get_supported_head_sizes(cls) -> list[int]:
@@ -80,22 +78,6 @@ class XFormersAttentionBackend(AttentionBackend):
78 256,
79 ]
80
83 - @staticmethod
84 - def get_supported_kernel_block_size() -> list[int | MultipleOf]:
85 - return [MultipleOf(16)]
86 -
87 - @classmethod
88 - def validate_head_size(cls, head_size: int) -> None:
89 - supported_head_sizes = cls.get_supported_head_sizes()
90 - if head_size not in supported_head_sizes:
91 - attn_type = cls.__name__.removesuffix("Backend")
92 - raise ValueError(
93 - f"Head size {head_size} is not supported by {attn_type}. "
94 - f"Supported head sizes are: {supported_head_sizes}. "
95 - "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
96 - "FlexAttention backend which supports all head sizes."
97 - )
98 -
81 @staticmethod
82 def get_name() -> str:
83 return "XFORMERS"
@@ -305,8 +287,6 @@ class XFormersAttentionImpl(AttentionImpl):
287 logits_soft_cap = 0
288 self.logits_soft_cap = logits_soft_cap
289
308 - XFormersAttentionBackend.validate_head_size(head_size)
309 -
290 if attn_type != AttentionType.DECODER:
291 raise NotImplementedError(
292 "Encoder self-attention and "
vllm/v1/spec_decode/eagle.py
+6 -2
@@ -150,11 +150,15 @@ class EagleProposer:
150 )
151
152 # Determine allowed attention backends once during initialization.
153 + from vllm.attention.backends.registry import AttentionBackendEnum
154 +
155 self.allowed_attn_types: tuple | None = None
156 if current_platform.is_rocm():
157 rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]
156 - # vllm.v1.attention.backends.rocm_aiter_fa is an optional backend
157 - if find_spec("vllm.v1.attention.backends.rocm_aiter_fa"):
158 + # ROCM_AITER_FA is an optional backend
159 + if find_spec(
160 + AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)
161 + ):
162 from vllm.v1.attention.backends.rocm_aiter_fa import (
163 AiterFlashAttentionMetadata,
164 )
vllm/v1/worker/gpu_model_runner.py
+2 -2
@@ -4371,7 +4371,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4371 """
4372 for backend in backends:
4373 is_supported = False
4374 - for supported_size in backend.get_supported_kernel_block_size():
4374 + for supported_size in backend.supported_kernel_block_sizes:
4375 if isinstance(supported_size, int):
4376 if block_size == supported_size:
4377 is_supported = True
@@ -4402,7 +4402,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4402 all_int_supported_sizes = set(
4403 supported_size
4404 for backend in backends
4405 - for supported_size in backend.get_supported_kernel_block_size()
4405 + for supported_size in backend.supported_kernel_block_sizes
4406 if isinstance(supported_size, int)
4407 )
4408