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.

MLAPrefillBackendEnum

Bases: Enum

Enumeration of all supported MLA prefill backends.

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
34
35
36
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
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"
    )
    # 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

    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_override() -> None

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

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

get_class 

get_class() -> type[MLAPrefillBackend]

Get the backend class (respects overrides).

Returns:

Type Description
type[MLAPrefillBackend]

The backend class

Raises:

Type Description
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 
74
75
76
77
78
79
80
81
82
83
84
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_path() -> str

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

Returns:

Type Description
str

The fully qualified class path string

Raises:

Type Description
ValueError

If Backend.CUSTOM is used without being registered
Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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 

is_overridden() -> bool

Check if this backend has been overridden.

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

_MLAPrefillBackendEnumMeta

Bases: EnumMeta

Metaclass for MLAPrefillBackendEnum to provide better error messages.

Source code in vllm/v1/attention/backends/mla/prefill/registry.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
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 

register_mla_prefill_backend(
    backend: MLAPrefillBackendEnum,
    class_path: str | None = None,
) -> Callable[[type], type]

Register or override an MLA prefill backend implementation.

Parameters:

Name Type Description Default
backend MLAPrefillBackendEnum

The MLAPrefillBackendEnum member to register.

 required
class_path str | None

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

 None

Returns:

Type Description
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): ...

Direct registration

register_mla_prefill_backend( MLAPrefillBackendEnum.CUSTOM, "my.module.MyCustomPrefillBackend" )

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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def register_mla_prefill_backend(
    backend: MLAPrefillBackendEnum,
    class_path: str | None = None,
) -> Callable[[type], type]:
    """Register or override an MLA prefill backend implementation.

    Args:
        backend: The MLAPrefillBackendEnum member to register.
        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):
            ...

        # Direct registration
        register_mla_prefill_backend(
            MLAPrefillBackendEnum.CUSTOM,
            "my.module.MyCustomPrefillBackend"
        )
    """

    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