[Platform] Allow BackendEnum dynamic register (#41091)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/test_attention_backend_registry.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
 import torch
 
 from vllm.platforms.interface import DeviceCapability
@@ -71,6 +72,48 @@ class CustomMambaAttentionBackend(AttentionBackend):
         return None
 
 
+class StringDecoratorAttentionBackend(AttentionBackend):
+    """Mock backend defined at module level so get_class() can resolve
+    the auto-generated class path (local classes are not importable)."""
+
+    @staticmethod
+    def get_name():
+        return "DECORATED"
+
+    @staticmethod
+    def get_impl_cls():
+        return CustomAttentionImpl
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_required_kv_cache_layout():
+        return None
+
+
+class StringDecoratorMambaBackend(AttentionBackend):
+    """Mock mamba backend defined at module level so get_class() can
+    resolve the auto-generated class path."""
+
+    @staticmethod
+    def get_name():
+        return "DECORATED"
+
+    @staticmethod
+    def get_impl_cls():
+        return CustomMambaAttentionImpl
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_required_kv_cache_layout():
+        return None
+
+
 def test_custom_attention_backend_requires_dcp_opt_in():
     """A new implementation remains usable without implicitly opting into DCP."""
     config = AttentionSelectorConfig(
@@ -178,3 +221,101 @@ def test_register_custom_mamba_backend_with_class_path():
     backend_cls = MambaAttentionBackendEnum.CUSTOM.get_class()
     assert backend_cls.get_name() == "CUSTOM_MAMBA"
     assert backend_cls.get_impl_cls() == CustomMambaAttentionImpl
+
+
+@pytest.fixture(params=[AttentionBackendEnum, MambaAttentionBackendEnum])
+def _enum_cls(request):
+    return request.param
+
+
+def _remove_dynamic_member(enum_cls, member):
+    value = member.value
+    name = member.name
+    enum_cls._member_map_.pop(name, None)
+    enum_cls._member_names_.remove(name)
+    delattr(enum_cls, name)
+    if enum_cls._value2member_map_.get(value) is member:
+        del enum_cls._value2member_map_[value]
+
+
+def test_register_dynamic_enum_member(_enum_cls):
+    backend_path = {
+        AttentionBackendEnum: (
+            "tests.test_attention_backend_registry.CustomAttentionBackend"
+        ),
+        MambaAttentionBackendEnum: (
+            "tests.test_attention_backend_registry.CustomMambaAttentionBackend"
+        ),
+    }[_enum_cls]
+
+    member = _enum_cls.register("DYNAMIC_TEST", backend_path)
+    assert member.name == "DYNAMIC_TEST"
+    assert member.value == backend_path
+    assert member is _enum_cls.DYNAMIC_TEST
+    assert _enum_cls["DYNAMIC_TEST"] is member
+    _remove_dynamic_member(_enum_cls, member)
+
+
+def test_register_dynamic_enum_member_duplicate_raises(_enum_cls):
+    _enum_cls.register("DUP_TEST", "some.module.Class")
+    with pytest.raises(ValueError, match="already exists"):
+        _enum_cls.register("DUP_TEST", "other.module.OtherClass")
+    _remove_dynamic_member(_enum_cls, _enum_cls.DUP_TEST)
+
+
+def test_register_dynamic_enum_member_reserved_name_raises(_enum_cls):
+    with pytest.raises(ValueError, match="Invalid or reserved backend name"):
+        _enum_cls.register("get_path", "some.module.Class")
+
+
+@pytest.fixture(
+    params=[
+        pytest.param((AttentionBackendEnum, False), id="attention"),
+        pytest.param((MambaAttentionBackendEnum, True), id="mamba"),
+    ]
+)
+def _enum_is_mamba(request):
+    return request.param
+
+
+def _backend_path(enum_cls):
+    if enum_cls is AttentionBackendEnum:
+        return "tests.test_attention_backend_registry.CustomAttentionBackend"
+    return "tests.test_attention_backend_registry.CustomMambaAttentionBackend"
+
+
+def test_register_backend_with_string_name_direct(_enum_is_mamba):
+    enum_cls, is_mamba = _enum_is_mamba
+    path = _backend_path(enum_cls)
+    register_backend("STRING_DIRECT", path, is_mamba=is_mamba)
+
+    member = enum_cls.STRING_DIRECT
+    assert member.is_overridden()
+    assert member.get_path() == path
+    expected = "CUSTOM" if not is_mamba else "CUSTOM_MAMBA"
+    assert member.get_class().get_name() == expected
+    member.clear_override()
+    _remove_dynamic_member(enum_cls, member)
+
+
+def test_register_backend_with_string_name_decorator(_enum_is_mamba):
+    enum_cls, is_mamba = _enum_is_mamba
+    # The decorated class must be module-level: the decorator derives a
+    # class path from it, and get_class() re-imports that path. A class
+    # defined inside this test would have a non-importable <locals> path.
+    decorated_cls = (
+        StringDecoratorMambaBackend if is_mamba else StringDecoratorAttentionBackend
+    )
+
+    result = register_backend("STRING_DECORATOR", is_mamba=is_mamba)(decorated_cls)
+    assert result is decorated_cls
+
+    member = enum_cls.STRING_DECORATOR
+    assert member.is_overridden()
+    assert member.get_path() == (
+        f"{decorated_cls.__module__}.{decorated_cls.__qualname__}"
+    )
+    assert member.get_class() is decorated_cls
+    assert member.get_class().get_name() == "DECORATED"
+    member.clear_override()
+    _remove_dynamic_member(enum_cls, member)

tests/v1/attention/test_mla_prefill_registry.py

@@ -54,10 +54,20 @@ def test_prefill_backend_clone_has_isolated_metadata():
 def cleanup_overrides():
     """Clear any overrides after each test."""
     yield
-    for member in MLAPrefillBackendEnum:
+    for member in list(MLAPrefillBackendEnum):
         member.clear_override()
 
 
+def _remove_dynamic_member(enum_cls, member):
+    value = member.value
+    name = member.name
+    enum_cls._member_map_.pop(name, None)
+    enum_cls._member_names_.remove(name)
+    delattr(enum_cls, name)
+    if enum_cls._value2member_map_.get(value) is member:
+        del enum_cls._value2member_map_[value]
+
+
 def test_custom_is_not_alias_of_any_backend():
     all_backends = list(MLAPrefillBackendEnum)
 
@@ -83,21 +93,11 @@ def test_custom_unregistered_raises():
 
 def test_register_custom_backend_with_class_path():
     register_mla_prefill_backend(
-        backend=MLAPrefillBackendEnum.CUSTOM,
+        MLAPrefillBackendEnum.CUSTOM,
-        class_path=(
+        "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend",
-            "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
-        ),
     )
-
     assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
-
+    assert MLAPrefillBackendEnum.CUSTOM.get_class().get_name() == "CUSTOM"
-    class_path = MLAPrefillBackendEnum.CUSTOM.get_path()
-    assert class_path == (
-        "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
-    )
-
-    backend_cls = MLAPrefillBackendEnum.CUSTOM.get_class()
-    assert backend_cls.get_name() == "CUSTOM"
 
 
 def test_register_custom_backend_as_decorator():
@@ -120,33 +120,20 @@ def test_register_custom_backend_as_decorator():
 
 
 def test_override_existing_backend():
-    original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
-
     register_mla_prefill_backend(
-        backend=MLAPrefillBackendEnum.FLASH_ATTN,
+        MLAPrefillBackendEnum.FLASH_ATTN,
-        class_path=(
+        "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend",
-            "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
-        ),
     )
-
+    assert MLAPrefillBackendEnum.FLASH_ATTN.get_class().get_name() == "CUSTOM"
     assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
-    assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() != original_path
-
-    backend_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
-    assert backend_cls.get_name() == "CUSTOM"
 
 
 def test_clear_override():
     original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
-
     register_mla_prefill_backend(
-        backend=MLAPrefillBackendEnum.FLASH_ATTN,
+        MLAPrefillBackendEnum.FLASH_ATTN,
-        class_path=(
+        "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend",
-            "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
-        ),
     )
-    assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
-
     MLAPrefillBackendEnum.FLASH_ATTN.clear_override()
     assert not MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
     assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() == original_path
@@ -172,3 +159,72 @@ def test_rocm_aiter_fa_registered():
     # The AITER FA path is the fp16/bf16 generic-varlen prefill path.
     assert backend_cls.supports_dtype(torch.bfloat16)
     assert backend_cls.supports_dtype(torch.float16)
+
+
+def test_register_dynamic_member():
+    p = "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
+    member = MLAPrefillBackendEnum.register("DYNAMIC_TEST", p)
+    assert member.name == "DYNAMIC_TEST"
+    assert member is MLAPrefillBackendEnum["DYNAMIC_TEST"]
+
+    _remove_dynamic_member(MLAPrefillBackendEnum, member)
+
+
+def test_register_dynamic_member_duplicate_raises():
+    path = "some.module.Class"
+    MLAPrefillBackendEnum.register("DUP_TEST", path)
+    with pytest.raises(ValueError, match="already exists"):
+        MLAPrefillBackendEnum.register("DUP_TEST", "other.module.OtherClass")
+    _remove_dynamic_member(MLAPrefillBackendEnum, MLAPrefillBackendEnum.DUP_TEST)
+
+
+class StringDecoratorMLAPrefillBackend(MLAPrefillBackend):
+    """Mock backend defined at module level so get_class() can resolve
+    the auto-generated class path (local classes are not importable)."""
+
+    supported_dtypes = [torch.bfloat16]
+    requires_r1_mla_dimensions = False
+
+    @staticmethod
+    def get_name() -> str:
+        return "DECORATED"
+
+    def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
+        raise NotImplementedError
+
+    def run_prefill_context_chunk(self, chunk_idx, q, k, v):
+        raise NotImplementedError
+
+
+def test_register_mla_prefill_backend_with_string_name_direct():
+    register_mla_prefill_backend(
+        "STRING_DIRECT",
+        "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend",
+    )
+    member = MLAPrefillBackendEnum.STRING_DIRECT
+    assert member.is_overridden()
+    assert member.get_class().get_name() == "CUSTOM"
+
+    member.clear_override()
+    _remove_dynamic_member(MLAPrefillBackendEnum, member)
+
+
+def test_register_mla_prefill_backend_with_string_name_decorator():
+    # The decorated class must be module-level: the decorator derives a
+    # class path from it, and get_class() re-imports that path. A class
+    # defined inside this test would have a non-importable <locals> path.
+    decorated_cls = StringDecoratorMLAPrefillBackend
+
+    result = register_mla_prefill_backend("STRING_DECORATOR")(decorated_cls)
+    assert result is decorated_cls
+
+    member = MLAPrefillBackendEnum.STRING_DECORATOR
+    assert member.is_overridden()
+    assert member.get_path() == (
+        f"{decorated_cls.__module__}.{decorated_cls.__qualname__}"
+    )
+    assert member.get_class() is decorated_cls
+    assert member.get_class().get_name() == "DECORATED"
+
+    member.clear_override()
+    _remove_dynamic_member(MLAPrefillBackendEnum, member)

vllm/v1/attention/backends/mla/prefill/registry.py

@@ -10,11 +10,14 @@ from collections.abc import Callable
 from enum import Enum, EnumMeta
 from typing import TYPE_CHECKING
 
+from vllm.logger import init_logger
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
 if TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
 
+logger = init_logger(__name__)
+
 
 class _MLAPrefillBackendEnumMeta(EnumMeta):
     """Metaclass for MLAPrefillBackendEnum to provide better error messages."""
@@ -61,6 +64,37 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
     # set to None to avoid alias with other backend, whose value is an empty string
     CUSTOM = None
 
+    @classmethod
+    def register(cls, name: str, value: str) -> "MLAPrefillBackendEnum":
+        """Dynamically register a new MLA prefill backend enum member.
+
+        Args:
+            name: The name for the new enum member
+            value: The fully qualified class path string
+
+        Returns:
+            The newly created enum member
+
+        """
+        if name in cls._member_map_:
+            raise ValueError(
+                f"MLA prefill backend {name} already exists in {cls.__name__}. "
+                f"Use register_mla_prefill_backend("
+                f"{cls.__name__}.{name}, '{value}') to override."
+            )
+        if not name.isidentifier() or hasattr(cls, name):
+            raise ValueError(f"Invalid or reserved backend name: {name}")
+
+        member = object.__new__(cls)
+        member._name_ = name
+        member._value_ = value
+        setattr(cls, name, member)
+        cls._member_map_[name] = member
+        cls._value2member_map_[value] = member
+        cls._member_names_.append(name)
+        logger.info("Registered new MLA prefill backend: %s -> %s", name, value)
+        return member
+
     def get_path(self) -> str:
         """Get the class path for this backend (respects overrides).
 
@@ -107,13 +141,15 @@ _MLA_PREFILL_OVERRIDES: dict[MLAPrefillBackendEnum, str] = {}
 
 
 def register_mla_prefill_backend(
-    backend: MLAPrefillBackendEnum,
+    backend: MLAPrefillBackendEnum | str,
     class_path: str | None = None,
 ) -> Callable[[type], type]:
     """Register or override an MLA prefill backend implementation.
 
     Args:
-        backend: The MLAPrefillBackendEnum member to register.
+        backend: The MLAPrefillBackendEnum member to register,
+                 or a string name for a new custom backend
+                 (e.g., "CUSTOM_MLA_PREFILL").
         class_path: Optional class path. If not provided and used as
             decorator, will be auto-generated from the class.
 
@@ -131,13 +167,27 @@ def register_mla_prefill_backend(
         class MyCustomPrefillBackend(MLAPrefillBackend):
             ...
 
+        # Register a new custom MLA prefill backend with a dynamic enum name
+        @register_mla_prefill_backend("CUSTOM_MLA_PREFILL")
+        class CustomMLAPrefillBackend(MLAPrefillBackend):
+            ...
+
         # Direct registration
         register_mla_prefill_backend(
             MLAPrefillBackendEnum.CUSTOM,
             "my.module.MyCustomPrefillBackend"
         )
 
+        # Direct registration with string name
+        register_mla_prefill_backend(
+            "CUSTOM_MLA_PREFILL",
+            "custom.attention.mla_prefill.CustomMLAPrefillBackend"
+        )
+
     """
+    # Handle dynamic enum creation for string backend names
+    if isinstance(backend, str):
+        backend = MLAPrefillBackendEnum.register(backend, class_path or "")
 
     def decorator(cls: type) -> type:
         _MLA_PREFILL_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"

vllm/v1/attention/backends/registry.py

@@ -157,6 +157,38 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     # set to None to avoid alias with other backend, whose value is an empty string
     CUSTOM = None
 
+    @classmethod
+    def register(cls, name: str, value: str) -> "AttentionBackendEnum":
+        """Dynamically register a new backend enum member.
+
+        Args:
+            name: The name for the new enum member
+            value: The fully qualified class path string
+
+        Returns:
+            The newly created enum member
+
+        """
+        if name in cls._member_map_:
+            raise ValueError(
+                f"Backend {name} already exists in {cls.__name__}. "
+                f"Use register_backend({cls.__name__}.{name}, '{value}') "
+                f"to override."
+            )
+        if not name.isidentifier() or hasattr(cls, name):
+            raise ValueError(f"Invalid or reserved backend name: {name}")
+
+        # Create new enum member dynamically
+        member = object.__new__(cls)
+        member._name_ = name
+        member._value_ = value
+        setattr(cls, name, member)
+        cls._member_map_[name] = member
+        cls._value2member_map_[value] = member
+        cls._member_names_.append(name)
+        logger.info("Registered new attention backend: %s -> %s", name, value)
+        return member
+
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).
 
@@ -214,6 +246,37 @@ class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
         backend.get_class()
     """
 
+    @classmethod
+    def register(cls, name: str, value: str) -> "MambaAttentionBackendEnum":
+        """Dynamically register a new mamba backend enum member.
+
+        Args:
+            name: The name for the new enum member
+            value: The fully qualified class path string
+
+        Returns:
+            The newly created enum member
+
+        """
+        if name in cls._member_map_:
+            raise ValueError(
+                f"Backend {name} already exists in {cls.__name__}. "
+                f"Use register_backend({cls.__name__}.{name}, '{value}') "
+                f"to override."
+            )
+        if not name.isidentifier() or hasattr(cls, name):
+            raise ValueError(f"Invalid or reserved backend name: {name}")
+
+        member = object.__new__(cls)
+        member._name_ = name
+        member._value_ = value
+        setattr(cls, name, member)
+        cls._member_map_[name] = member
+        cls._value2member_map_[value] = member
+        cls._member_names_.append(name)
+        logger.info("Registered new attention backend: %s -> %s", name, value)
+        return member
+
     MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend"
     MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend"
     SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend"
@@ -275,14 +338,15 @@ _MAMBA_ATTN_OVERRIDES: dict[MambaAttentionBackendEnum, str] = {}
 
 
 def register_backend(
-    backend: AttentionBackendEnum | MambaAttentionBackendEnum,
+    backend: AttentionBackendEnum | MambaAttentionBackendEnum | str,
     class_path: str | None = None,
     is_mamba: bool = False,
 ) -> Callable[[type], type]:
     """Register or override a backend implementation.
 
     Args:
-        backend: The AttentionBackendEnum member to register
+        backend: The AttentionBackendEnum/MambaAttentionBackendEnum member to register,
+                 or a string name for a new custom backend (e.g., "CUSTOM_MLA").
         class_path: Optional class path. If not provided and used as
             decorator, will be auto-generated from the class.
         is_mamba: Whether the backend is a Mamba attention backend.
@@ -306,13 +370,30 @@ def register_backend(
         class MyCustomBackend:
             ...
 
+        # Register a new custom attention backend with a dynamic enum name
+        @register_backend("CUSTOM_MLA")
+        class CustomMLABackend:
+            ...
+
         # Direct registration
         register_backend(
             AttentionBackendEnum.CUSTOM,
             "my.module.MyCustomBackend"
         )
 
+        # Direct registration with string name
+        register_backend(
+            "CUSTOM_MLA",
+            "custom.attention.mla_v1.CustomMLAABackend"
+        )
+
     """
+    # Handle dynamic enum creation for string backend names
+    if isinstance(backend, str):
+        if is_mamba:
+            backend = MambaAttentionBackendEnum.register(backend, class_path or "")
+        else:
+            backend = AttentionBackendEnum.register(backend, class_path or "")
 
     def decorator(cls: type) -> type:
         if is_mamba: