diff --git a/tests/test_attention_backend_registry.py b/tests/test_attention_backend_registry.py index 5885411ed9..14fe06699d 100644 --- a/tests/test_attention_backend_registry.py +++ b/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 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) diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index fc5efe3a5b..0068d53b88 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/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, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.CUSTOM, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - assert MLAPrefillBackendEnum.CUSTOM.is_overridden() - - 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" + assert MLAPrefillBackendEnum.CUSTOM.get_class().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, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.FLASH_ATTN, + "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, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.FLASH_ATTN, + "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 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) diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index f7ec31e153..19d37fbc22 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/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__}" diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index ad184f0913..207c112a38 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/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: