| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Selector for MLA prefill backends. |
| 4 | |
| 5 | This module provides functions for selecting the appropriate MLA prefill |
| 6 | backend based on device capabilities and configuration. |
| 7 | """ |
| 8 | |
| 9 | from functools import cache |
| 10 | from typing import TYPE_CHECKING, NamedTuple |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.platforms.interface import DeviceCapability |
| 16 | from vllm.v1.attention.backends.mla.prefill.base import MLADimensions |
| 17 | from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum |
| 18 | |
| 19 | if TYPE_CHECKING: |
| 20 | from vllm.config import VllmConfig |
| 21 | from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | class MLAPrefillSelectorConfig(NamedTuple): |
| 27 | """Hashable configuration for MLA prefill backend selection. |
| 28 | |
| 29 | This is analogous to AttentionSelectorConfig and contains model-specific |
| 30 | configuration needed to select an MLA prefill backend, extracted from |
| 31 | VllmConfig into a hashable form for caching. |
| 32 | """ |
| 33 | |
| 34 | dtype: torch.dtype |
| 35 | mla_dimensions: MLADimensions = MLADimensions( |
| 36 | qk_nope_head_dim=0, |
| 37 | qk_rope_head_dim=0, |
| 38 | v_head_dim=0, |
| 39 | ) |
| 40 | |
| 41 | def __repr__(self): |
| 42 | return ( |
| 43 | f"MLAPrefillSelectorConfig(dtype={self.dtype}, " |
| 44 | f"mla_dimensions={self.mla_dimensions})" |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def _get_mla_prefill_backend_priorities( |
| 49 | device_capability: DeviceCapability, |
| 50 | mla_dimensions: MLADimensions, |
| 51 | ) -> list[MLAPrefillBackendEnum]: |
| 52 | """Get MLA prefill backend priorities based on device capability. |
| 53 | |
| 54 | Args: |
| 55 | device_capability: The device's compute capability. |
| 56 | mla_dimensions: The model's MLA head dimensions. |
| 57 | |
| 58 | Returns: |
| 59 | List of backends in priority order (highest priority first). |
| 60 | |
| 61 | """ |
| 62 | from vllm.platforms import current_platform |
| 63 | |
| 64 | if current_platform.is_rocm(): |
| 65 | return [ |
| 66 | MLAPrefillBackendEnum.ROCM_AITER_FA, |
| 67 | MLAPrefillBackendEnum.FLASH_ATTN, |
| 68 | ] |
| 69 | |
| 70 | if device_capability.major == 10: # Blackwell |
| 71 | if mla_dimensions == MLADimensions( |
| 72 | qk_nope_head_dim=192, |
| 73 | qk_rope_head_dim=64, |
| 74 | v_head_dim=256, |
| 75 | ): |
| 76 | return [ |
| 77 | MLAPrefillBackendEnum.TRTLLM_RAGGED, |
| 78 | MLAPrefillBackendEnum.FLASH_ATTN, |
| 79 | MLAPrefillBackendEnum.FLASHINFER, |
| 80 | MLAPrefillBackendEnum.TOKENSPEED_MLA, |
| 81 | ] |
| 82 | return [ |
| 83 | MLAPrefillBackendEnum.FLASH_ATTN, |
| 84 | MLAPrefillBackendEnum.TRTLLM_RAGGED, |
| 85 | MLAPrefillBackendEnum.FLASHINFER, |
| 86 | MLAPrefillBackendEnum.TOKENSPEED_MLA, |
| 87 | ] |
| 88 | else: # Hopper (SM90) and older |
| 89 | return [ |
| 90 | MLAPrefillBackendEnum.FLASH_ATTN, |
| 91 | ] |
| 92 | |
| 93 | |
| 94 | def get_mla_prefill_backend( |
| 95 | vllm_config: "VllmConfig", |
| 96 | ) -> "type[MLAPrefillBackend]": |
| 97 | """Select the MLA prefill backend based on configuration and device. |
| 98 | |
| 99 | This function first checks for explicit user preferences via |
| 100 | mla_prefill_backend in AttentionConfig, then falls back to automatic |
| 101 | priority-based selection. |
| 102 | |
| 103 | Args: |
| 104 | vllm_config: The vLLM configuration. |
| 105 | |
| 106 | Returns: |
| 107 | The selected prefill backend class. |
| 108 | |
| 109 | """ |
| 110 | from vllm.platforms import current_platform |
| 111 | |
| 112 | if current_platform.is_cpu(): |
| 113 | # CPUs have no compute capability, so the capability-driven priority |
| 114 | # path below does not apply. Prefer an accelerator-specific CPU backend |
| 115 | # when its kernels are present, else the generic SDPA one. |
| 116 | for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): |
| 117 | try: |
| 118 | cpu_backend_cls = backend_enum.get_class() |
| 119 | except ImportError: |
| 120 | continue |
| 121 | if cpu_backend_cls.is_available(): |
| 122 | logger.info_once("Using %s MLA prefill backend.", backend_enum.name) |
| 123 | return cpu_backend_cls |
| 124 | raise ValueError("No valid CPU MLA prefill backend found.") |
| 125 | |
| 126 | device_capability = current_platform.get_device_capability() |
| 127 | if device_capability is None: |
| 128 | logger.info_once( |
| 129 | "Device capability not available, using FlashAttention MLA prefill backend." |
| 130 | ) |
| 131 | return MLAPrefillBackendEnum.FLASH_ATTN.get_class() |
| 132 | |
| 133 | attention_config = vllm_config.attention_config |
| 134 | |
| 135 | model_config = vllm_config.model_config |
| 136 | if model_config is None: |
| 137 | selector_config = MLAPrefillSelectorConfig(dtype=torch.get_default_dtype()) |
| 138 | else: |
| 139 | hf_text_config = model_config.hf_text_config |
| 140 | selector_config = MLAPrefillSelectorConfig( |
| 141 | dtype=model_config.dtype, |
| 142 | mla_dimensions=MLADimensions( |
| 143 | qk_nope_head_dim=getattr(hf_text_config, "qk_nope_head_dim", 0), |
| 144 | qk_rope_head_dim=getattr(hf_text_config, "qk_rope_head_dim", 0), |
| 145 | v_head_dim=getattr(hf_text_config, "v_head_dim", 0), |
| 146 | ), |
| 147 | ) |
| 148 | |
| 149 | if attention_config.mla_prefill_backend is not None: |
| 150 | selected_backend = attention_config.mla_prefill_backend |
| 151 | backend_cls: type[MLAPrefillBackend] | None = None |
| 152 | try: |
| 153 | backend_cls = selected_backend.get_class() |
| 154 | invalid_reasons = backend_cls.validate_configuration( |
| 155 | device_capability, selector_config |
| 156 | ) |
| 157 | except ImportError: |
| 158 | invalid_reasons = ["ImportError"] |
| 159 | if invalid_reasons: |
| 160 | raise ValueError( |
| 161 | f"Selected MLA prefill backend {selected_backend.name} " |
| 162 | f"is not valid for this configuration. " |
| 163 | f"Reason: {invalid_reasons}" |
| 164 | ) |
| 165 | assert backend_cls is not None |
| 166 | logger.info_once("Using %s MLA prefill backend.", selected_backend.name) |
| 167 | return backend_cls |
| 168 | |
| 169 | return _auto_select_mla_prefill_backend( |
| 170 | device_capability, |
| 171 | selector_config, |
| 172 | ) |
| 173 | |
| 174 | |
| 175 | @cache |
| 176 | def _auto_select_mla_prefill_backend( |
| 177 | device_capability: DeviceCapability, |
| 178 | selector_config: MLAPrefillSelectorConfig, |
| 179 | ) -> "type[MLAPrefillBackend]": |
| 180 | """Auto-select the best available MLA prefill backend. |
| 181 | |
| 182 | Args: |
| 183 | device_capability: The device's compute capability. |
| 184 | selector_config: Hashable configuration for backend selection. |
| 185 | |
| 186 | Returns: |
| 187 | The selected prefill backend class. |
| 188 | |
| 189 | """ |
| 190 | priorities = _get_mla_prefill_backend_priorities( |
| 191 | device_capability, |
| 192 | selector_config.mla_dimensions, |
| 193 | ) |
| 194 | all_invalid_reasons: dict[str, list[str]] = {} |
| 195 | |
| 196 | for backend_enum in priorities: |
| 197 | backend_cls: type[MLAPrefillBackend] | None = None |
| 198 | try: |
| 199 | backend_cls = backend_enum.get_class() |
| 200 | invalid_reasons = backend_cls.validate_configuration( |
| 201 | device_capability, selector_config |
| 202 | ) |
| 203 | except ImportError: |
| 204 | invalid_reasons = ["ImportError"] |
| 205 | if not invalid_reasons: |
| 206 | assert backend_cls is not None |
| 207 | logger.info_once("Using %s MLA prefill backend.", backend_enum.name) |
| 208 | return backend_cls |
| 209 | all_invalid_reasons[backend_enum.name] = invalid_reasons |
| 210 | |
| 211 | reasons_str = ( |
| 212 | "{" |
| 213 | + ", ".join( |
| 214 | f"{name}: [{', '.join(reasons)}]" |
| 215 | for name, reasons in all_invalid_reasons.items() |
| 216 | ) |
| 217 | + "}" |
| 218 | ) |
| 219 | config_str = repr(selector_config) |
| 220 | logger.debug_once( |
| 221 | "Some MLA prefill backends are not valid with %s. Reasons: %s.", |
| 222 | config_str, |
| 223 | reasons_str, |
| 224 | ) |
| 225 | |
| 226 | raise ValueError( |
| 227 | f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}." |
| 228 | ) |