| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Registry for MLA prefill backends. |
| 4 | |
| 5 | This module provides an enumeration of all available MLA prefill backends |
| 6 | and utilities for loading and registering them. |
| 7 | """ |
| 8 | |
| 9 | from collections.abc import Callable |
| 10 | from enum import Enum, EnumMeta |
| 11 | from typing import TYPE_CHECKING |
| 12 | |
| 13 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 14 | |
| 15 | if TYPE_CHECKING: |
| 16 | from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend |
| 17 | |
| 18 | |
| 19 | class _MLAPrefillBackendEnumMeta(EnumMeta): |
| 20 | """Metaclass for MLAPrefillBackendEnum to provide better error messages.""" |
| 21 | |
| 22 | def __getitem__(cls, name: str): |
| 23 | try: |
| 24 | return super().__getitem__(name) |
| 25 | except KeyError: |
| 26 | members = cls.__members__.keys() |
| 27 | valid_backends = ", ".join(members) |
| 28 | raise ValueError( |
| 29 | f"Unknown MLA prefill backend: '{name}'. " |
| 30 | f"Valid options are: {valid_backends}" |
| 31 | ) from None |
| 32 | |
| 33 | |
| 34 | class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): |
| 35 | """Enumeration of all supported MLA prefill backends.""" |
| 36 | |
| 37 | FLASH_ATTN = ( |
| 38 | "vllm.v1.attention.backends.mla.prefill.flash_attn.FlashAttnPrefillBackend" |
| 39 | ) |
| 40 | FLASHINFER = ( |
| 41 | "vllm.v1.attention.backends.mla.prefill.flashinfer.FlashInferPrefillBackend" |
| 42 | ) |
| 43 | TRTLLM_RAGGED = ( |
| 44 | "vllm.v1.attention.backends.mla.prefill.trtllm_ragged." |
| 45 | "TrtllmRaggedPrefillBackend" |
| 46 | ) |
| 47 | TOKENSPEED_MLA = ( |
| 48 | "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla." |
| 49 | "TokenspeedMLAPrefillBackend" |
| 50 | ) |
| 51 | ROCM_AITER_FA = ( |
| 52 | "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn." |
| 53 | "AiterFlashAttnPrefillBackend" |
| 54 | ) |
| 55 | CPU = "vllm.v1.attention.backends.mla.prefill.cpu_sdpa.CPUSDPAMLAPrefillBackend" |
| 56 | ZEN_CPU = ( |
| 57 | "vllm.v1.attention.backends.mla.prefill.zen_cpu_sdpa." |
| 58 | "ZenCPUSDPAMLAPrefillBackend" |
| 59 | ) |
| 60 | # Placeholder for third-party/custom backends - must be registered before use |
| 61 | # set to None to avoid alias with other backend, whose value is an empty string |
| 62 | CUSTOM = None |
| 63 | |
| 64 | def get_path(self) -> str: |
| 65 | """Get the class path for this backend (respects overrides). |
| 66 | |
| 67 | Returns: |
| 68 | The fully qualified class path string |
| 69 | |
| 70 | Raises: |
| 71 | ValueError: If Backend.CUSTOM is used without being registered |
| 72 | |
| 73 | """ |
| 74 | path = _MLA_PREFILL_OVERRIDES.get(self, self.value) |
| 75 | if not path: |
| 76 | raise ValueError( |
| 77 | f"MLA prefill backend {self.name} must be registered before " |
| 78 | f"use. Use register_mla_prefill_backend(" |
| 79 | f"MLAPrefillBackendEnum.{self.name}, " |
| 80 | f"'your.module.YourClass')" |
| 81 | ) |
| 82 | return path |
| 83 | |
| 84 | def get_class(self) -> "type[MLAPrefillBackend]": |
| 85 | """Get the backend class (respects overrides). |
| 86 | |
| 87 | Returns: |
| 88 | The backend class |
| 89 | |
| 90 | Raises: |
| 91 | ImportError: If the backend class cannot be imported |
| 92 | ValueError: If CUSTOM is used without being registered |
| 93 | |
| 94 | """ |
| 95 | return resolve_obj_by_qualname(self.get_path()) |
| 96 | |
| 97 | def is_overridden(self) -> bool: |
| 98 | """Check if this backend has been overridden.""" |
| 99 | return self in _MLA_PREFILL_OVERRIDES |
| 100 | |
| 101 | def clear_override(self) -> None: |
| 102 | """Clear any override for this backend, reverting to the default.""" |
| 103 | _MLA_PREFILL_OVERRIDES.pop(self, None) |
| 104 | |
| 105 | |
| 106 | _MLA_PREFILL_OVERRIDES: dict[MLAPrefillBackendEnum, str] = {} |
| 107 | |
| 108 | |
| 109 | def register_mla_prefill_backend( |
| 110 | backend: MLAPrefillBackendEnum, |
| 111 | class_path: str | None = None, |
| 112 | ) -> Callable[[type], type]: |
| 113 | """Register or override an MLA prefill backend implementation. |
| 114 | |
| 115 | Args: |
| 116 | backend: The MLAPrefillBackendEnum member to register. |
| 117 | class_path: Optional class path. If not provided and used as |
| 118 | decorator, will be auto-generated from the class. |
| 119 | |
| 120 | Returns: |
| 121 | Decorator function if class_path is None, otherwise a no-op. |
| 122 | |
| 123 | Examples: |
| 124 | # Override an existing MLA prefill backend |
| 125 | @register_mla_prefill_backend(MLAPrefillBackendEnum.FLASH_ATTN) |
| 126 | class MyCustomFlashAttn(MLAPrefillBackend): |
| 127 | ... |
| 128 | |
| 129 | # Register a custom third-party MLA prefill backend |
| 130 | @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM) |
| 131 | class MyCustomPrefillBackend(MLAPrefillBackend): |
| 132 | ... |
| 133 | |
| 134 | # Direct registration |
| 135 | register_mla_prefill_backend( |
| 136 | MLAPrefillBackendEnum.CUSTOM, |
| 137 | "my.module.MyCustomPrefillBackend" |
| 138 | ) |
| 139 | |
| 140 | """ |
| 141 | |
| 142 | def decorator(cls: type) -> type: |
| 143 | _MLA_PREFILL_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" |
| 144 | return cls |
| 145 | |
| 146 | if class_path is not None: |
| 147 | _MLA_PREFILL_OVERRIDES[backend] = class_path |
| 148 | return lambda x: x |
| 149 | |
| 150 | return decorator |