Skip to content

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

Registry for MLA prefill backends.

This module provides an enumeration of all available MLA prefill backends and utilities for loading and registering them.

Classes:

Functions:

MLAPrefillBackendEnum

Bases: Enum

Enumeration of all supported MLA prefill backends.

Methods:

  • clear_override –

    Clear any override for this backend, reverting to the default.

  • get_class –

    Get the backend class (respects overrides).

  • get_path –

    Get the class path for this backend (respects overrides).

  • is_overridden –

    Check if this backend has been overridden.

  • register –

    Dynamically register a new MLA prefill backend enum member.

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
    """Enumeration of all supported MLA prefill backends."""

    FLASH_ATTN = (
        "vllm.v1.attention.backends.mla.prefill.flash_attn.FlashAttnPrefillBackend"
    )
    FLASHINFER = (
        "vllm.v1.attention.backends.mla.prefill.flashinfer.FlashInferPrefillBackend"
    )
    TRTLLM_RAGGED = (
        "vllm.v1.attention.backends.mla.prefill.trtllm_ragged."
        "TrtllmRaggedPrefillBackend"
    )
    TOKENSPEED_MLA = (
        "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla."
        "TokenspeedMLAPrefillBackend"
    )
    ROCM_AITER_FA = (
        "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
        "AiterFlashAttnPrefillBackend"
    )
    CPU = "vllm.v1.attention.backends.mla.prefill.cpu_sdpa.CPUSDPAMLAPrefillBackend"
    ZEN_CPU = (
        "vllm.v1.attention.backends.mla.prefill.zen_cpu_sdpa."
        "ZenCPUSDPAMLAPrefillBackend"
    )
    # Placeholder for third-party/custom backends - must be registered before use
    # 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).

        Returns:
            The fully qualified class path string

        Raises:
            ValueError: If Backend.CUSTOM is used without being registered

        """
        path = _MLA_PREFILL_OVERRIDES.get(self, self.value)
        if not path:
            raise ValueError(
                f"MLA prefill backend {self.name} must be registered before "
                f"use. Use register_mla_prefill_backend("
                f"MLAPrefillBackendEnum.{self.name}, "
                f"'your.module.YourClass')"
            )
        return path

    def get_class(self) -> "type[MLAPrefillBackend]":
        """Get the backend class (respects overrides).

        Returns:
            The backend class

        Raises:
            ImportError: If the backend class cannot be imported
            ValueError: If CUSTOM is used without being registered

        """
        return resolve_obj_by_qualname(self.get_path())

    def is_overridden(self) -> bool:
        """Check if this backend has been overridden."""
        return self in _MLA_PREFILL_OVERRIDES

    def clear_override(self) -> None:
        """Clear any override for this backend, reverting to the default."""
        _MLA_PREFILL_OVERRIDES.pop(self, None)

clear_override()

Clear any override for this backend, reverting to the default.

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
135
136
137
def clear_override(self) -> None:
    """Clear any override for this backend, reverting to the default."""
    _MLA_PREFILL_OVERRIDES.pop(self, None)

get_class()

Get the backend class (respects overrides).

Returns:

Raises:

  • ImportError –

    If the backend class cannot be imported

  • ValueError –

    If CUSTOM is used without being registered

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
118
119
120
121
122
123
124
125
126
127
128
129
def get_class(self) -> "type[MLAPrefillBackend]":
    """Get the backend class (respects overrides).

    Returns:
        The backend class

    Raises:
        ImportError: If the backend class cannot be imported
        ValueError: If CUSTOM is used without being registered

    """
    return resolve_obj_by_qualname(self.get_path())

get_path()

Get the class path for this backend (respects overrides).

Returns:

  • str –

    The fully qualified class path string

Raises:

  • ValueError –

    If Backend.CUSTOM is used without being registered

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def get_path(self) -> str:
    """Get the class path for this backend (respects overrides).

    Returns:
        The fully qualified class path string

    Raises:
        ValueError: If Backend.CUSTOM is used without being registered

    """
    path = _MLA_PREFILL_OVERRIDES.get(self, self.value)
    if not path:
        raise ValueError(
            f"MLA prefill backend {self.name} must be registered before "
            f"use. Use register_mla_prefill_backend("
            f"MLAPrefillBackendEnum.{self.name}, "
            f"'your.module.YourClass')"
        )
    return path

is_overridden()

Check if this backend has been overridden.

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
131
132
133
def is_overridden(self) -> bool:
    """Check if this backend has been overridden."""
    return self in _MLA_PREFILL_OVERRIDES

register(name, value) classmethod

Dynamically register a new MLA prefill backend enum member.

Parameters:

  • name

    (str) –

    The name for the new enum member

  • value

    (str) –

    The fully qualified class path string

Returns:

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@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

_MLAPrefillBackendEnumMeta

Bases: EnumMeta

Metaclass for MLAPrefillBackendEnum to provide better error messages.

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
class _MLAPrefillBackendEnumMeta(EnumMeta):
    """Metaclass for MLAPrefillBackendEnum to provide better error messages."""

    def __getitem__(cls, name: str):
        try:
            return super().__getitem__(name)
        except KeyError:
            members = cls.__members__.keys()
            valid_backends = ", ".join(members)
            raise ValueError(
                f"Unknown MLA prefill backend: '{name}'. "
                f"Valid options are: {valid_backends}"
            ) from None

register_mla_prefill_backend(backend, class_path=None)

Register or override an MLA prefill backend implementation.

Parameters:

  • backend

    (MLAPrefillBackendEnum | str) –

    The MLAPrefillBackendEnum member to register, or a string name for a new custom backend (e.g., "CUSTOM_MLA_PREFILL").

  • class_path

    (str | None, default: None ) –

    Optional class path. If not provided and used as decorator, will be auto-generated from the class.

Returns:

  • Callable[[type], type] –

    Decorator function if class_path is None, otherwise a no-op.

Examples:

Override an existing MLA prefill backend

@register_mla_prefill_backend(MLAPrefillBackendEnum.FLASH_ATTN) class MyCustomFlashAttn(MLAPrefillBackend): ...

Register a custom third-party MLA prefill backend

@register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM) 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" )

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def register_mla_prefill_backend(
    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,
                 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.

    Returns:
        Decorator function if class_path is None, otherwise a no-op.

    Examples:
        # Override an existing MLA prefill backend
        @register_mla_prefill_backend(MLAPrefillBackendEnum.FLASH_ATTN)
        class MyCustomFlashAttn(MLAPrefillBackend):
            ...

        # Register a custom third-party MLA prefill backend
        @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM)
        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__}"
        return cls

    if class_path is not None:
        _MLA_PREFILL_OVERRIDES[backend] = class_path
        return lambda x: x

    return decorator