A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Attention] Implement universal BACKEND_MAP (#25900)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 8, 2025 at 15:00 UTC
76879cc160425960b9101eaade27a0565a461afc
12 files changed
+119
-75
tests/kernels/attention/test_attention_selector.py
+2
-2
@@ -34,7 +34,7 @@ DEVICE_MLA_BACKENDS = {
34
35
DEVICE_REGULAR_ATTN_BACKENDS = {
36
"cuda": ["XFORMERS", "FLASHINFER", "FLASH_ATTN"],
37
- "hip": ["ROCM_FLASH"],
37
+ "hip": ["ROCM_ATTN"],
38
"cpu": ["TORCH_SDPA"],
39
}
40
@@ -122,7 +122,7 @@ def test_env(
122
backend = get_attn_backend(
123
16, torch.float16, None, block_size, use_mla=use_mla
124
)
125
- expected = "TRITON_ATTN"
125
+ expected = "ROCM_ATTN"
126
assert backend.get_name() == expected
127
128
elif device == "cuda":
tests/kernels/attention/test_rocm_attention_selector.py
+1
-1
@@ -18,7 +18,7 @@ def clear_cache():
18
@pytest.mark.skip(reason="Skipped for now. Should be revisited.")
19
def test_selector(monkeypatch: pytest.MonkeyPatch):
20
with monkeypatch.context() as m:
21
- m.setenv(STR_BACKEND_ENV_VAR, "ROCM_FLASH")
21
+ m.setenv(STR_BACKEND_ENV_VAR, "ROCM_ATTN")
22
23
# Set the current platform to ROCm using monkeypatch
24
monkeypatch.setattr("vllm.attention.selector.current_platform", RocmPlatform())
tests/v1/attention/test_attention_backends.py
+2
-2
@@ -14,7 +14,7 @@ from tests.v1.attention.utils import (
14
create_common_attn_metadata,
15
create_standard_kv_cache_spec,
16
create_vllm_config,
17
- get_attention_backend,
17
+ try_get_attention_backend,
18
)
19
from vllm.attention.backends.registry import _Backend
20
from vllm.config import ModelConfig
@@ -214,7 +214,7 @@ def run_attention_backend(
214
actual_backend = _Backend.FLEX_ATTENTION
215
use_direct_block_mask = False
216
217
- builder_cls, impl_cls = get_attention_backend(actual_backend)
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:
tests/v1/attention/test_mla_backends.py
+3
-3
@@ -12,7 +12,7 @@ from tests.v1.attention.utils import (
12
create_common_attn_metadata,
13
create_standard_kv_cache_spec,
14
create_vllm_config,
15
- get_attention_backend,
15
+ try_get_attention_backend,
16
)
17
from vllm import _custom_ops as ops
18
from vllm.attention.backends.registry import _Backend
@@ -239,7 +239,7 @@ def run_attention_backend(
239
) -> torch.Tensor:
240
"""Run attention computation using the specified backend's AttentionImpl."""
241
242
- builder_cls, impl_cls = get_attention_backend(backend)
242
+ builder_cls, impl_cls = try_get_attention_backend(backend)
243
244
# Build metadata
245
builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
@@ -400,7 +400,7 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
400
# Determine if this is decode or prefill
401
is_decode = []
402
for i, backend in enumerate(BACKENDS_TO_TEST):
403
- builder_cls, _ = get_attention_backend(backend)
403
+ builder_cls, _ = try_get_attention_backend(backend)
404
is_decode.append(q_len <= builder_cls.reorder_batch_threshold)
405
406
# Split q into nope and rope components
tests/v1/attention/utils.py
+14
-38
@@ -8,7 +8,8 @@ from typing import Optional, Union
8
import pytest
9
import torch
10
11
-from vllm.attention.backends.registry import _Backend
11
+from vllm.attention.backends.abstract import AttentionImpl
12
+from vllm.attention.backends.registry import _Backend, backend_to_class_str
13
from vllm.config import (
14
CacheConfig,
15
CompilationConfig,
@@ -20,9 +21,11 @@ from vllm.config import (
21
VllmConfig,
22
)
23
from vllm.config.model import ModelDType
23
-from vllm.platforms import current_platform
24
from vllm.utils import resolve_obj_by_qualname
25
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
25
+from vllm.v1.attention.backends.utils import (
26
+ AttentionMetadataBuilder,
27
+ CommonAttentionMetadata,
28
+)
29
from vllm.v1.kv_cache_interface import FullAttentionSpec
30
31
@@ -117,44 +120,17 @@ def create_common_attn_metadata(
120
)
121
122
120
-def get_attention_backend(backend_name: _Backend):
121
- """Set up attention backend classes for testing.
122
-
123
- Args:
124
- backend_name: Name of the backend ("flash_attn", "flashinfer", etc.)
125
- vllm_config: VllmConfig instance
126
-
127
- Returns:
128
- Tuple of (backend_builder_class, backend_impl_class)
129
- """
130
- backend_map = {
131
- _Backend.FLASH_ATTN: (
132
- "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
133
- if current_platform.is_cuda()
134
- else "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
135
- ),
136
- _Backend.FLASHINFER: "vllm.v1.attention.backends.flashinfer.FlashInferBackend",
137
- _Backend.FLEX_ATTENTION: "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend", # noqa: E501
138
- _Backend.TRITON_ATTN: "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend", # noqa: E501
139
- _Backend.TREE_ATTN: "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend",
140
- _Backend.XFORMERS: "vllm.v1.attention.backends.xformers.XFormersAttentionBackend", # noqa: E501
141
- _Backend.CUTLASS_MLA: "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend", # noqa: E501
142
- _Backend.FLASHMLA: "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
143
- _Backend.FLASH_ATTN_MLA: "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend", # noqa: E501
144
- _Backend.FLASHINFER_MLA: "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend", # noqa: E501
145
- _Backend.TRITON_MLA: "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend", # noqa: E501
146
- }
147
-
148
- if backend_name not in backend_map:
149
- raise ValueError(f"Unknown backend: {backend_name}")
150
-
151
- backend_class_name = backend_map[backend_name]
152
-
123
+def try_get_attention_backend(
124
+ backend: _Backend,
125
+) -> tuple[type[AttentionMetadataBuilder], type[AttentionImpl]]:
126
+ """Try to get the attention backend class, skipping test if not found."""
127
+ backend_class_str = backend_to_class_str(backend)
128
try:
154
- backend_class = resolve_obj_by_qualname(backend_class_name)
129
+ backend_class = resolve_obj_by_qualname(backend_class_str)
130
return backend_class.get_builder_cls(), backend_class.get_impl_cls()
131
except ImportError as e:
157
- pytest.skip(f"{backend_name} not available: {e}")
132
+ pytest.skip(f"{backend_class_str} not available: {e}")
133
+ raise AssertionError("unreachable") from None
134
135
136
def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
tests/v1/spec_decode/test_eagle.py
+5
-5
@@ -12,7 +12,7 @@ from tests.v1.attention.utils import (
12
BatchSpec,
13
create_common_attn_metadata,
14
create_standard_kv_cache_spec,
15
- get_attention_backend,
15
+ try_get_attention_backend,
16
)
17
from vllm.attention.backends.registry import _Backend
18
from vllm.config import (
@@ -535,11 +535,11 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
535
sampling_metadata = mock.MagicMock()
536
537
if attn_backend == "FLASH_ATTN":
538
- attn_metadata_builder_cls, _ = get_attention_backend(_Backend.FLASH_ATTN)
538
+ attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.FLASH_ATTN)
539
elif attn_backend == "TRITON_ATTN":
540
- attn_metadata_builder_cls, _ = get_attention_backend(_Backend.TRITON_ATTN)
540
+ attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.TRITON_ATTN)
541
elif attn_backend == "TREE_ATTN":
542
- attn_metadata_builder_cls, _ = get_attention_backend(_Backend.TREE_ATTN)
542
+ attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.TREE_ATTN)
543
else:
544
raise ValueError(f"Unsupported attention backend: {attn_backend}")
545
@@ -674,7 +674,7 @@ def test_propose_tree(spec_token_tree):
674
proposer.attn_layer_names = ["layer.0"]
675
676
# Get the tree attention metadata builder.
677
- attn_metadata_builder_cls, _ = get_attention_backend(_Backend.TREE_ATTN)
677
+ attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.TREE_ATTN)
678
attn_metadata_builder = attn_metadata_builder_cls(
679
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
680
layer_names=proposer.attn_layer_names,
tests/v1/spec_decode/test_mtp.py
+2
-2
@@ -10,7 +10,7 @@ from tests.v1.attention.utils import (
10
BatchSpec,
11
create_common_attn_metadata,
12
create_standard_kv_cache_spec,
13
- get_attention_backend,
13
+ try_get_attention_backend,
14
)
15
from vllm.attention.backends.registry import _Backend
16
from vllm.config import (
@@ -177,7 +177,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):
177
sampling_metadata = mock.MagicMock()
178
179
# Setup attention metadata
180
- attn_metadata_builder_cls, _ = get_attention_backend(_Backend.FLASH_ATTN)
180
+ attn_metadata_builder_cls, _ = try_get_attention_backend(_Backend.FLASH_ATTN)
181
182
attn_metadata_builder = attn_metadata_builder_cls(
183
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
tests/v1/spec_decode/test_tree_attention.py
+2
-2
@@ -9,7 +9,7 @@ import torch
9
from tests.v1.attention.utils import (
10
create_standard_kv_cache_spec,
11
create_vllm_config,
12
- get_attention_backend,
12
+ try_get_attention_backend,
13
)
14
from vllm.attention.backends.registry import _Backend
15
from vllm.config import ParallelConfig, SpeculativeConfig
@@ -63,7 +63,7 @@ def forward_attention(
63
64
# Build common metadata.
65
model_name = "meta-llama/Meta-Llama-3-8B"
66
- builder_cls, impl_cls = get_attention_backend(backend)
66
+ builder_cls, impl_cls = try_get_attention_backend(backend)
67
vllm_config = create_vllm_config(model_name=model_name, max_model_len=max(seq_lens))
68
if spec_token_tree is not None:
69
# Create speculative config if token tree is specified.
vllm/attention/backends/registry.py
+83
-2
@@ -3,13 +3,16 @@
3
"""Attention backend registry"""
4
5
import enum
6
+from typing import Optional
7
+
8
+from vllm.utils import resolve_obj_by_qualname
9
10
11
class _Backend(enum.Enum):
12
FLASH_ATTN = enum.auto()
13
TRITON_ATTN = enum.auto()
14
XFORMERS = enum.auto()
12
- ROCM_FLASH = enum.auto()
15
+ ROCM_ATTN = enum.auto()
16
ROCM_AITER_MLA = enum.auto()
17
ROCM_AITER_FA = enum.auto() # used for ViT attn backend
18
TORCH_SDPA = enum.auto()
@@ -24,5 +27,83 @@ class _Backend(enum.Enum):
27
NO_ATTENTION = enum.auto()
28
FLEX_ATTENTION = enum.auto()
29
TREE_ATTN = enum.auto()
27
- ROCM_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.FLASH_ATTN_MLA: "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend", # noqa: E501
47
+ _Backend.PALLAS: "vllm.v1.attention.backends.pallas.PallasAttentionBackend", # noqa: E501
48
+ _Backend.FLEX_ATTENTION: "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend", # noqa: E501
49
+ _Backend.TREE_ATTN: "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend", # noqa: E501
50
+ _Backend.ROCM_AITER_UNIFIED_ATTN: "vllm.v1.attention.backends.rocm_aiter_unified_attn.RocmAiterUnifiedAttentionBackend", # noqa: E501
51
+}
52
+
53
+
54
+def register_attn_backend(backend: _Backend, class_path: Optional[str] = None):
55
+ """
56
+ Decorator: register a custom attention backend into BACKEND_MAPPING.
57
+ - If class_path is provided, use it.
58
+ - Otherwise, auto-generate from the class object.
59
+ Validation: only checks if 'backend' is a valid _Backend enum member.
60
+ Overwriting existing mappings is allowed. This enables other hardware
61
+ platforms to plug in custom out-of-tree backends.
62
+ """
63
+ if not isinstance(backend, _Backend):
64
+ raise ValueError(f"{backend} is not a valid _Backend enum value.")
65
+
66
+ def decorator(cls):
67
+ path = class_path or f"{cls.__module__}.{cls.__qualname__}"
68
+ BACKEND_MAP[backend] = path
69
+ return cls
70
+
71
+ return decorator
72
+
73
+
74
+def backend_to_class_str(backend: _Backend) -> str:
75
+ """Get the backend class string
76
+
77
+ Args:
78
+ backend: The backend enum value
79
+
80
+ Returns:
81
+ The backend class string
82
+ """
83
+ return BACKEND_MAP[backend]
84
+
85
+
86
+def backend_to_class(backend: _Backend) -> type:
87
+ """Get the backend class.
88
+
89
+ Args:
90
+ backend: The backend enum value
91
+
92
+ Returns:
93
+ The backend class
94
+ """
95
+ backend_class_name = backend_to_class_str(backend)
96
+ return resolve_obj_by_qualname(backend_class_name)
97
+
98
+
99
+def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
100
+ """
101
+ Convert a string backend name to a _Backend enum value.
102
+
103
+ Returns:
104
+ _Backend: enum value if backend_name is a valid in-tree type
105
+ None: otherwise it's an invalid in-tree type or an out-of-tree platform
106
+ is loaded.
107
+ """
108
+ assert backend_name is not None
109
+ return _Backend[backend_name] if backend_name in _Backend.__members__ else None
vllm/attention/layer.py
+2
-2
@@ -11,8 +11,8 @@ import torch.nn.functional as F
11
import vllm.envs as envs
12
from vllm.attention import AttentionType
13
from vllm.attention.backends.abstract import AttentionBackend
14
-from vllm.attention.backends.registry import _Backend
15
-from vllm.attention.selector import backend_name_to_enum, get_attn_backend
14
+from vllm.attention.backends.registry import _Backend, backend_name_to_enum
15
+from vllm.attention.selector import get_attn_backend
16
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
17
from vllm.config import CacheConfig, get_current_vllm_config
18
from vllm.distributed.kv_transfer import (
vllm/attention/selector.py
+1
-14
@@ -12,7 +12,7 @@ import torch
12
13
import vllm.envs as envs
14
from vllm.attention.backends.abstract import AttentionBackend
15
-from vllm.attention.backends.registry import _Backend
15
+from vllm.attention.backends.registry import _Backend, backend_name_to_enum
16
from vllm.logger import init_logger
17
from vllm.platforms import current_platform
18
from vllm.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
@@ -20,19 +20,6 @@ from vllm.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
20
logger = init_logger(__name__)
21
22
23
-def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
24
- """
25
- Convert a string backend name to a _Backend enum value.
26
-
27
- Returns:
28
- * _Backend: enum value if backend_name is a valid in-tree type
29
- * None: otherwise it's an invalid in-tree type or an out-of-tree platform is
30
- loaded.
31
- """
32
- assert backend_name is not None
33
- return _Backend[backend_name] if backend_name in _Backend.__members__ else None
34
-
35
-
23
def get_env_variable_attn_backend() -> Optional[_Backend]:
24
"""
25
Get the backend override specified by the vLLM attention
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+2
-2
@@ -21,8 +21,8 @@ import torch
21
import zmq
22
23
from vllm import envs
24
-from vllm.attention.backends.registry import _Backend
25
-from vllm.attention.selector import backend_name_to_enum, get_attn_backend
24
+from vllm.attention.backends.registry import _Backend, backend_name_to_enum
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 (
28
CopyBlocksOp,