attention backends register themselves in the mapping

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Author: MatthewBonanni

tests/kernels/attention/test_attention_selector.py

@@ -32,7 +32,7 @@ def clear_cache():
 
 DEVICE_REGULAR_ATTN_BACKENDS = {
     "cuda": ["XFORMERS", "FLASHINFER", "FLASH_ATTN"],
-    "hip": ["ROCM_FLASH"],
+    "hip": ["ROCM_AITER_FA"],
     "cpu": ["TORCH_SDPA"],
 }
 

tests/kernels/attention/test_rocm_attention_selector.py

@@ -19,16 +19,16 @@ def clear_cache():
 @pytest.mark.skip(reason="Skipped for now. Should be revisited.")
 def test_selector(monkeypatch: pytest.MonkeyPatch):
     with monkeypatch.context() as m:
-        m.setenv(STR_BACKEND_ENV_VAR, "ROCM_FLASH")
+        m.setenv(STR_BACKEND_ENV_VAR, "ROCM_AITER_FA")
 
         # Set the current platform to ROCm using monkeypatch
         monkeypatch.setattr("vllm.attention.selector.current_platform",
                             RocmPlatform())
 
         # Test standard ROCm attention
         backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
-        assert (backend.get_name() == "ROCM_FLASH"
-                or backend.get_name() == "TRITON_ATTN")
+        assert (backend.get_name() == "ROCM_AITER_FA"
+                or backend.get_name() == "ROCM_ATTN")
 
         # MLA test for deepseek related
 

tests/v1/attention/test_attention_backends.py

@@ -11,7 +11,7 @@
 from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
                                       create_standard_kv_cache_spec,
                                       create_vllm_config,
-                                      get_attention_backend)
+                                      try_get_attention_backend)
 from vllm.attention.backends.registry import _Backend
 from vllm.config import ModelConfig
 from vllm.platforms import current_platform
@@ -210,7 +210,7 @@ def run_attention_backend(
         actual_backend = _Backend.FLEX_ATTENTION
         use_direct_block_mask = False
 
-    builder_cls, impl_cls = get_attention_backend(actual_backend)
+    builder_cls, impl_cls = try_get_attention_backend(actual_backend)
 
     # Mock flashinfer's get_per_layer_parameters if needed
     if actual_backend == _Backend.FLASHINFER:

tests/v1/attention/test_mla_backends.py

@@ -9,7 +9,7 @@
 from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
                                       create_standard_kv_cache_spec,
                                       create_vllm_config,
-                                      get_attention_backend)
+                                      try_get_attention_backend)
 from vllm import _custom_ops as ops
 from vllm.attention.backends.registry import _Backend
 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
@@ -232,7 +232,7 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
                           mock_kv_b_proj) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
 
-    builder_cls, impl_cls = get_attention_backend(backend)
+    builder_cls, impl_cls = try_get_attention_backend(backend)
 
     # Build metadata
     builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
@@ -393,7 +393,7 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
         # Determine if this is decode or prefill
         is_decode = []
         for i, backend in enumerate(BACKENDS_TO_TEST):
-            builder_cls, _ = get_attention_backend(backend)
+            builder_cls, _ = try_get_attention_backend(backend)
             is_decode.append(q_len <= builder_cls.reorder_batch_threshold)
 
         # Split q into nope and rope components

tests/v1/attention/utils.py

@@ -8,11 +8,10 @@
 import pytest
 import torch
 
-from vllm.attention.backends.registry import _Backend
+from vllm.attention.backends.registry import _Backend, backend_to_class_str
 from vllm.config import (CacheConfig, CompilationConfig, DeviceConfig,
                          LoadConfig, ModelConfig, ModelDType, ParallelConfig,
                          SchedulerConfig, VllmConfig)
-from vllm.platforms import current_platform
 from vllm.utils import resolve_obj_by_qualname
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.kv_cache_interface import FullAttentionSpec
@@ -110,54 +109,15 @@ def create_common_attn_metadata(
     )
 
 
-def get_attention_backend(backend_name: _Backend):
-    """Set up attention backend classes for testing.
-
-    Args:
-        backend_name: Name of the backend ("flash_attn", "flashinfer", etc.)
-        vllm_config: VllmConfig instance
-
-    Returns:
-        Tuple of (backend_builder_class, backend_impl_class)
-    """
-    backend_map = {
-        _Backend.FLASH_ATTN:
-        ("vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
-         if current_platform.is_cuda() else
-         "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
-         ),
-        _Backend.FLASHINFER:
-        "vllm.v1.attention.backends.flashinfer.FlashInferBackend",
-        _Backend.FLEX_ATTENTION:
-        "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend",
-        _Backend.TRITON_ATTN:
-        "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend",
-        _Backend.TREE_ATTN:
-        "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend",
-        _Backend.XFORMERS:
-        "vllm.v1.attention.backends.xformers.XFormersAttentionBackend",
-        _Backend.CUTLASS_MLA:
-        "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
-        _Backend.FLASHMLA:
-        "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
-        _Backend.FLASH_ATTN_MLA:
-        "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
-        _Backend.FLASHINFER_MLA:
-        "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
-        _Backend.TRITON_MLA:
-        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
-    }
-
-    if backend_name not in backend_map:
-        raise ValueError(f"Unknown backend: {backend_name}")
-
-    backend_class_name = backend_map[backend_name]
-
+def try_get_attention_backend(backend: _Backend) -> tuple[type, type]:
+    """Try to get the attention backend class, skipping test if not found."""
+    backend_class_str = backend_to_class_str(backend)
     try:
-        backend_class = resolve_obj_by_qualname(backend_class_name)
+        backend_class = resolve_obj_by_qualname(backend_class_str)
         return backend_class.get_builder_cls(), backend_class.get_impl_cls()
     except ImportError as e:
-        pytest.skip(f"{backend_name} not available: {e}")
+        pytest.skip(f"{backend_class_str} not available: {e}")
+        assert False  # unreachable -- satisfies mypy
 
 
 def create_standard_kv_cache_spec(

tests/v1/spec_decode/test_eagle.py

@@ -10,7 +10,7 @@
 from tests.utils import get_attn_backend_list_based_on_platform
 from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
                                       create_standard_kv_cache_spec,
-                                      get_attention_backend)
+                                      try_get_attention_backend)
 from vllm.attention.backends.registry import _Backend
 from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
                          ParallelConfig, SchedulerConfig, SpeculativeConfig,
@@ -515,13 +515,13 @@ def create_deterministic_logits(token_ids):
     sampling_metadata = mock.MagicMock()
 
     if attn_backend == "FLASH_ATTN":
-        attn_metadata_builder_cls, _ = get_attention_backend(
+        attn_metadata_builder_cls, _ = try_get_attention_backend(
             _Backend.FLASH_ATTN)
     elif attn_backend == "TRITON_ATTN":
-        attn_metadata_builder_cls, _ = get_attention_backend(
+        attn_metadata_builder_cls, _ = try_get_attention_backend(
             _Backend.TRITON_ATTN)
     elif attn_backend == "TREE_ATTN":
-        attn_metadata_builder_cls, _ = get_attention_backend(
+        attn_metadata_builder_cls, _ = try_get_attention_backend(
             _Backend.TREE_ATTN)
     else:
         raise ValueError(f"Unsupported attention backend: {attn_backend}")
@@ -653,7 +653,8 @@ def create_deterministic_logits(token_ids, k: int):
     proposer.attn_layer_names = ["layer.0"]
 
     # Get the tree attention metadata builder.
-    attn_metadata_builder_cls, _ = get_attention_backend(_Backend.TREE_ATTN)
+    attn_metadata_builder_cls, _ = try_get_attention_backend(
+        _Backend.TREE_ATTN)
     attn_metadata_builder = attn_metadata_builder_cls(
         kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
         layer_names=proposer.attn_layer_names,

tests/v1/spec_decode/test_mtp.py

@@ -8,7 +8,7 @@
 
 from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
                                       create_standard_kv_cache_spec,
-                                      get_attention_backend)
+                                      try_get_attention_backend)
 from vllm.attention.backends.registry import _Backend
 from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
                          ParallelConfig, SchedulerConfig, SpeculativeConfig,
@@ -174,7 +174,8 @@ def create_deterministic_logits(batch_size, vocab_size, token_offset):
     sampling_metadata = mock.MagicMock()
 
     # Setup attention metadata
-    attn_metadata_builder_cls, _ = get_attention_backend(_Backend.FLASH_ATTN)
+    attn_metadata_builder_cls, _ = try_get_attention_backend(
+        _Backend.FLASH_ATTN)
 
     attn_metadata_builder = attn_metadata_builder_cls(
         kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

tests/v1/spec_decode/test_tree_attention.py

@@ -8,7 +8,7 @@
 
 from tests.v1.attention.utils import (create_standard_kv_cache_spec,
                                       create_vllm_config,
-                                      get_attention_backend)
+                                      try_get_attention_backend)
 from vllm.attention.backends.registry import _Backend
 from vllm.config import ParallelConfig, SpeculativeConfig
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
@@ -60,7 +60,7 @@ def forward_attention(
 
     # Build common metadata.
     model_name = "meta-llama/Meta-Llama-3-8B"
-    builder_cls, impl_cls = get_attention_backend(backend)
+    builder_cls, impl_cls = try_get_attention_backend(backend)
     vllm_config = create_vllm_config(model_name=model_name,
                                      max_model_len=max(seq_lens))
     if spec_token_tree is not None:

vllm/attention/backends/registry.py

@@ -3,13 +3,16 @@
 """Attention backend registry"""
 
 import enum
+from typing import Optional, Type
+
+from vllm.utils import resolve_obj_by_qualname
 
 
 class _Backend(enum.Enum):
     FLASH_ATTN = enum.auto()
     TRITON_ATTN = enum.auto()
     XFORMERS = enum.auto()
-    ROCM_FLASH = enum.auto()
+    ROCM_ATTN = enum.auto()
     ROCM_AITER_MLA = enum.auto()
     ROCM_AITER_FA = enum.auto()  # used for ViT attn backend
     TORCH_SDPA = enum.auto()
@@ -24,4 +27,64 @@ class _Backend(enum.Enum):
     NO_ATTENTION = enum.auto()
     FLEX_ATTENTION = enum.auto()
     TREE_ATTN = enum.auto()
-    ROCM_ATTN = enum.auto()
+
+
+BACKEND_MAPPING = {}
+
+
+def register_attn_backend(backend: _Backend, class_path: str | None = None):
+    """
+    Decorator: register a custom attention backend into BACKEND_MAPPING.
+    - If class_path is provided, use it.
+    - Otherwise, auto-generate from the class object.
+    Validation: only checks if 'backend' is a valid _Backend enum member.
+    Overwriting existing mappings is allowed.
+    """
+    if not isinstance(backend, _Backend):
+        raise ValueError(f"{backend} is not a valid _Backend enum value.")
+
+    def decorator(cls):
+        path = class_path or f"{cls.__module__}.{cls.__qualname__}"
+        BACKEND_MAPPING[backend] = path
+        return cls
+
+    return decorator
+
+
+def backend_to_class_str(backend: _Backend) -> str:
+    """Get the backend class string
+    
+    Args:
+        backend: The backend enum value
+        
+    Returns:
+        The backend class string
+    """
+    return BACKEND_MAPPING[backend]
+
+
+def backend_to_class(backend: _Backend) -> Type:
+    """Get the backend class.
+
+    Args:
+        backend: The backend enum value
+
+    Returns:
+        The backend class
+    """
+    backend_class_name = backend_to_class_str(backend)
+    return resolve_obj_by_qualname(backend_class_name)
+
+
+def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
+    """
+    Convert a string backend name to a _Backend enum value.
+
+    Returns:
+        _Backend: enum value if backend_name is a valid in-tree type
+        None: otherwise it's an invalid in-tree type or an out-of-tree platform
+              is loaded.
+    """
+    assert backend_name is not None
+    return _Backend[backend_name] if backend_name in _Backend.__members__ else \
+          None

vllm/attention/layer.py

@@ -10,8 +10,8 @@
 import vllm.envs as envs
 from vllm.attention import AttentionType
 from vllm.attention.backends.abstract import AttentionBackend
-from vllm.attention.backends.registry import _Backend
-from vllm.attention.selector import backend_name_to_enum, get_attn_backend
+from vllm.attention.backends.registry import _Backend, backend_name_to_enum
+from vllm.attention.selector import get_attn_backend
 from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
 from vllm.config import CacheConfig, get_current_vllm_config
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,