| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention backend registry.""" |
| 4 | |
| 5 | from collections.abc import Callable |
| 6 | from enum import Enum, EnumMeta |
| 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.v1.attention.backend import AttentionBackend |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | class _AttentionBackendEnumMeta(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, Enum]", cls.__members__).keys() |
| 27 | valid_backends = ", ".join(members) |
| 28 | raise ValueError( |
| 29 | f"Unknown attention backend: '{name}'. " |
| 30 | f"Valid options are: {valid_backends}" |
| 31 | ) from None |
| 32 | |
| 33 | |
| 34 | class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): |
| 35 | """Enumeration of all supported attention backends. |
| 36 | |
| 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 | |
| 44 | FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" |
| 45 | FLASH_ATTN_DIFFKV = ( |
| 46 | "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend" |
| 47 | ) |
| 48 | TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" |
| 49 | TRITON_ATTN_DIFFKV = ( |
| 50 | "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend" |
| 51 | ) |
| 52 | TRITON_FLASH_ATTN = ( |
| 53 | "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend" |
| 54 | ) |
| 55 | ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" |
| 56 | ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" |
| 57 | ROCM_AITER_TRITON_MLA = ( |
| 58 | "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend" |
| 59 | ) |
| 60 | ROCM_AITER_FA = ( |
| 61 | "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend" |
| 62 | ) |
| 63 | ROCM_AITER_MLA_SPARSE = ( |
| 64 | "vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse.ROCMAiterMLASparseBackend" |
| 65 | ) |
| 66 | XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend" |
| 67 | TORCH_SDPA = "" # this tag is only used for ViT |
| 68 | FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" |
| 69 | TRITON_FLASHINFER = ( |
| 70 | "vllm.v1.attention.backends.triton_flashinfer.TritonFlashInferBackend" |
| 71 | ) |
| 72 | FLASHINFER_MLA = ( |
| 73 | "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend" |
| 74 | ) |
| 75 | TOKENSPEED_MLA = ( |
| 76 | "vllm.v1.attention.backends.mla.tokenspeed_mla.TokenspeedMLABackend" |
| 77 | ) |
| 78 | FLASHINFER_MLA_SPARSE = ( |
| 79 | "vllm.v1.attention.backends.mla.flashinfer_mla_sparse." |
| 80 | "FlashInferMLASparseTRTLLMBackend" |
| 81 | ) |
| 82 | FLASHINFER_MLA_SPARSE_SM120 = ( |
| 83 | "vllm.v1.attention.backends.mla.flashinfer_mla_sparse." |
| 84 | "FlashInferMLASparseSM120Backend" |
| 85 | ) |
| 86 | FLASHINFER_MLA_SPARSE_SM90 = ( |
| 87 | "vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90." |
| 88 | "FlashInferMLASparseSM90Backend" |
| 89 | ) |
| 90 | TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend" |
| 91 | CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend" |
| 92 | FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend" |
| 93 | FLASHMLA_SPARSE = ( |
| 94 | "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend" |
| 95 | ) |
| 96 | # DeepSeek V4 sparse MLA backends (model-driven; selected via the V4 layer). |
| 97 | FLASHMLA_SPARSE_DSV4 = ( |
| 98 | "vllm.models.deepseek_v4.sparse_mla.DeepseekV4FlashMLABackend" |
| 99 | ) |
| 100 | FLASHINFER_MLA_SPARSE_DSV4 = ( |
| 101 | "vllm.models.deepseek_v4.nvidia.flashinfer_sparse." |
| 102 | "DeepseekV4FlashInferMLASparseBackend" |
| 103 | ) |
| 104 | ROCM_FLASHMLA_SPARSE_DSV4 = ( |
| 105 | "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend" |
| 106 | ) |
| 107 | # DeepSeek V4.1 sparse MLA backends (model-driven; selected via the V4.1 |
| 108 | # layer). Separate names from DSV4 so a V4.1 model never resolves the |
| 109 | # V4.0 backend classes through this enum. |
| 110 | FLASHMLA_SPARSE_DSV41 = ( |
| 111 | "vllm.models.deepseek_v41.sparse_mla.DeepseekV4FlashMLABackend" |
| 112 | ) |
| 113 | FLASHINFER_MLA_SPARSE_DSV41 = ( |
| 114 | "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." |
| 115 | "DeepseekV4FlashInferMLASparseBackend" |
| 116 | ) |
| 117 | FLASHMLA_MEGA_ATTN_DSV41 = ( |
| 118 | "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend" |
| 119 | ) |
| 120 | B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend" |
| 121 | FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" |
| 122 | FLASH_ATTN_MLA_SPARSE = ( |
| 123 | "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" |
| 124 | ) |
| 125 | MINIMAX_M3_SPARSE = ( |
| 126 | "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend" |
| 127 | ) |
| 128 | CUTLASS_MSA = ( |
| 129 | "vllm.models.minimax_m3.nvidia.sparse_attention_msa." |
| 130 | "MiniMaxM3SparseCutlassBackend" |
| 131 | ) |
| 132 | TRITON_MSA = ( |
| 133 | "vllm.models.minimax_m3.nvidia.sparse_attention_msa." |
| 134 | "MiniMaxM3SparseTritonBackend" |
| 135 | ) |
| 136 | NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend" |
| 137 | FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" |
| 138 | # HPC Attention Backend: |
| 139 | # powered by operators from https://github.com/Tencent/hpc-ops. |
| 140 | # Only supported on NVIDIA Hopper GPUs (e.g. H20, H200), |
| 141 | # currently limited to the Hy3 model, |
| 142 | # and requires a block size of 64. |
| 143 | HPC_ATTN = "vllm.v1.attention.backends.hpc_attn.HpcAttentionBackend" |
| 144 | ROCM_AITER_UNIFIED_ATTN = ( |
| 145 | "vllm.v1.attention.backends.rocm_aiter_unified_attn." |
| 146 | "RocmAiterUnifiedAttentionBackend" |
| 147 | ) |
| 148 | CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend" |
| 149 | CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend" |
| 150 | AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend" |
| 151 | TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" |
| 152 | # Placeholder for third-party/custom backends - must be registered before use |
| 153 | # set to None to avoid alias with other backend, whose value is an empty string |
| 154 | CUSTOM = None |
| 155 | |
| 156 | def get_path(self, include_classname: bool = True) -> str: |
| 157 | """Get the class path for this backend (respects overrides). |
| 158 | |
| 159 | Returns: |
| 160 | The fully qualified class path string |
| 161 | |
| 162 | Raises: |
| 163 | ValueError: If Backend.CUSTOM is used without being registered |
| 164 | |
| 165 | """ |
| 166 | path = _ATTN_OVERRIDES.get(self, self.value) |
| 167 | if not path: |
| 168 | raise ValueError( |
| 169 | f"Backend {self.name} must be registered before use. " |
| 170 | f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')" |
| 171 | ) |
| 172 | if not include_classname: |
| 173 | path = path.rsplit(".", 1)[0] |
| 174 | return path |
| 175 | |
| 176 | def get_class(self) -> "type[AttentionBackend]": |
| 177 | """Get the backend class (respects overrides). |
| 178 | |
| 179 | Returns: |
| 180 | The backend class |
| 181 | |
| 182 | Raises: |
| 183 | ImportError: If the backend class cannot be imported |
| 184 | ValueError: If Backend.CUSTOM is used without being registered |
| 185 | |
| 186 | """ |
| 187 | return resolve_obj_by_qualname(self.get_path()) |
| 188 | |
| 189 | def is_overridden(self) -> bool: |
| 190 | """Check if this backend has been overridden. |
| 191 | |
| 192 | Returns: |
| 193 | True if the backend has a registered override |
| 194 | |
| 195 | """ |
| 196 | return self in _ATTN_OVERRIDES |
| 197 | |
| 198 | def clear_override(self) -> None: |
| 199 | """Clear any override for this backend, reverting to the default.""" |
| 200 | _ATTN_OVERRIDES.pop(self, None) |
| 201 | |
| 202 | |
| 203 | class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): |
| 204 | """Enumeration of all supported mamba attention backends. |
| 205 | |
| 206 | The enum value is the default class path, but this can be overridden |
| 207 | at runtime using register_backend(). |
| 208 | |
| 209 | To get the actual backend class (respecting overrides), use: |
| 210 | backend.get_class() |
| 211 | """ |
| 212 | |
| 213 | MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend" |
| 214 | MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend" |
| 215 | SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend" |
| 216 | LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend" |
| 217 | GDN_ATTN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend" |
| 218 | # Placeholder for third-party/custom backends - must be registered before use |
| 219 | # set to None to avoid alias with other backend, whose value is an empty string |
| 220 | CUSTOM = None |
| 221 | |
| 222 | def get_path(self, include_classname: bool = True) -> str: |
| 223 | """Get the class path for this backend (respects overrides). |
| 224 | |
| 225 | Returns: |
| 226 | The fully qualified class path string |
| 227 | |
| 228 | Raises: |
| 229 | ValueError: If Backend.CUSTOM is used without being registered |
| 230 | |
| 231 | """ |
| 232 | path = _MAMBA_ATTN_OVERRIDES.get(self, self.value) |
| 233 | if not path: |
| 234 | raise ValueError( |
| 235 | f"Backend {self.name} must be registered before use. " |
| 236 | f"Use register_backend(Backend.{self.name}, 'your.module.YourClass')" |
| 237 | ) |
| 238 | if not include_classname: |
| 239 | path = path.rsplit(".", 1)[0] |
| 240 | return path |
| 241 | |
| 242 | def get_class(self) -> "type[AttentionBackend]": |
| 243 | """Get the backend class (respects overrides). |
| 244 | |
| 245 | Returns: |
| 246 | The backend class |
| 247 | |
| 248 | Raises: |
| 249 | ImportError: If the backend class cannot be imported |
| 250 | ValueError: If Backend.CUSTOM is used without being registered |
| 251 | |
| 252 | """ |
| 253 | return resolve_obj_by_qualname(self.get_path()) |
| 254 | |
| 255 | def is_overridden(self) -> bool: |
| 256 | """Check if this backend has been overridden. |
| 257 | |
| 258 | Returns: |
| 259 | True if the backend has a registered override |
| 260 | |
| 261 | """ |
| 262 | return self in _MAMBA_ATTN_OVERRIDES |
| 263 | |
| 264 | def clear_override(self) -> None: |
| 265 | """Clear any override for this backend, reverting to the default.""" |
| 266 | _MAMBA_ATTN_OVERRIDES.pop(self, None) |
| 267 | |
| 268 | |
| 269 | _ATTN_OVERRIDES: dict[AttentionBackendEnum, str] = {} |
| 270 | _MAMBA_ATTN_OVERRIDES: dict[MambaAttentionBackendEnum, str] = {} |
| 271 | |
| 272 | |
| 273 | def register_backend( |
| 274 | backend: AttentionBackendEnum | MambaAttentionBackendEnum, |
| 275 | class_path: str | None = None, |
| 276 | is_mamba: bool = False, |
| 277 | ) -> Callable[[type], type]: |
| 278 | """Register or override a backend implementation. |
| 279 | |
| 280 | Args: |
| 281 | backend: The AttentionBackendEnum member to register |
| 282 | class_path: Optional class path. If not provided and used as |
| 283 | decorator, will be auto-generated from the class. |
| 284 | is_mamba: Whether the backend is a Mamba attention backend. |
| 285 | |
| 286 | Returns: |
| 287 | Decorator function if class_path is None, otherwise a no-op |
| 288 | |
| 289 | Examples: |
| 290 | # Override an existing attention backend |
| 291 | @register_backend(AttentionBackendEnum.FLASH_ATTN) |
| 292 | class MyCustomFlashAttn: |
| 293 | ... |
| 294 | |
| 295 | # Override an existing mamba attention backend |
| 296 | @register_backend(MambaAttentionBackendEnum.LINEAR, is_mamba=True) |
| 297 | class MyCustomMambaAttn: |
| 298 | ... |
| 299 | |
| 300 | # Register a custom third-party attention backend |
| 301 | @register_backend(AttentionBackendEnum.CUSTOM) |
| 302 | class MyCustomBackend: |
| 303 | ... |
| 304 | |
| 305 | # Direct registration |
| 306 | register_backend( |
| 307 | AttentionBackendEnum.CUSTOM, |
| 308 | "my.module.MyCustomBackend" |
| 309 | ) |
| 310 | |
| 311 | """ |
| 312 | |
| 313 | def decorator(cls: type) -> type: |
| 314 | if is_mamba: |
| 315 | _MAMBA_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index] |
| 316 | else: |
| 317 | _ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" # type: ignore[index] |
| 318 | return cls |
| 319 | |
| 320 | if class_path is not None: |
| 321 | if is_mamba: |
| 322 | _MAMBA_ATTN_OVERRIDES[backend] = class_path # type: ignore[index] |
| 323 | else: |
| 324 | _ATTN_OVERRIDES[backend] = class_path # type: ignore[index] |
| 325 | return lambda x: x |
| 326 | |
| 327 | return decorator |