| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Abstract base class for MLA prefill backends.""" |
| 4 | |
| 5 | from abc import ABC, abstractmethod |
| 6 | from dataclasses import dataclass |
| 7 | from typing import TYPE_CHECKING, ClassVar |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | if TYPE_CHECKING: |
| 12 | from vllm.config import VllmConfig |
| 13 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 14 | MLACommonPrefillMetadata, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 17 | from vllm.platforms.interface import DeviceCapability |
| 18 | from vllm.v1.attention.backends.mla.prefill.selector import ( |
| 19 | MLAPrefillSelectorConfig, |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | @dataclass(frozen=True, kw_only=True) |
| 24 | class MLADimensions: |
| 25 | qk_nope_head_dim: int |
| 26 | qk_rope_head_dim: int |
| 27 | v_head_dim: int |
| 28 | |
| 29 | def __str__(self) -> str: |
| 30 | return ( |
| 31 | f"(qk_nope_head_dim={self.qk_nope_head_dim}, " |
| 32 | f"qk_rope_head_dim={self.qk_rope_head_dim}, " |
| 33 | f"v_head_dim={self.v_head_dim})" |
| 34 | ) |
| 35 | |
| 36 | |
| 37 | class MLAPrefillBackend(ABC): |
| 38 | """Abstract base class for MLA prefill backends.""" |
| 39 | |
| 40 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 41 | torch.float16, |
| 42 | torch.bfloat16, |
| 43 | ] |
| 44 | supported_mla_dimensions: ClassVar[list[MLADimensions]] = [] |
| 45 | |
| 46 | @staticmethod |
| 47 | @abstractmethod |
| 48 | def get_name() -> str: |
| 49 | raise NotImplementedError |
| 50 | |
| 51 | @classmethod |
| 52 | def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: |
| 53 | return True |
| 54 | |
| 55 | @classmethod |
| 56 | def supports_dtype(cls, dtype: torch.dtype) -> bool: |
| 57 | return dtype in cls.supported_dtypes |
| 58 | |
| 59 | @classmethod |
| 60 | def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool: |
| 61 | return ( |
| 62 | not cls.supported_mla_dimensions |
| 63 | or mla_dimensions in cls.supported_mla_dimensions |
| 64 | ) |
| 65 | |
| 66 | @classmethod |
| 67 | def is_available(cls) -> bool: |
| 68 | return True |
| 69 | |
| 70 | def supports_quant_output(self, quant_key: "QuantKey") -> bool: |
| 71 | """Whether `run_prefill_new_tokens` can write quantized output |
| 72 | directly (fused) for the given quant key, skipping the post-quant |
| 73 | pass. Overridden by backends that support it.""" |
| 74 | return False |
| 75 | |
| 76 | def supports_out(self) -> bool: |
| 77 | """Whether `run_prefill_new_tokens` and `run_prefill_context_chunk` honor |
| 78 | a caller-provided `out` tensor of shape |
| 79 | `[num_tokens, num_heads, v_head_dim]`, writing the result into it in place |
| 80 | and returning it. |
| 81 | |
| 82 | When True, callers may pass `out` and skip the post-hoc |
| 83 | slice/flatten/copy -- and, for context chunks, size the accumulating |
| 84 | partial before running any chunk. False for backends that ignore `out` or |
| 85 | emit a padded (`qk_head_dim`) output. Overridden by backends that support |
| 86 | it. |
| 87 | """ |
| 88 | return False |
| 89 | |
| 90 | @classmethod |
| 91 | def validate_configuration( |
| 92 | cls, |
| 93 | device_capability: "DeviceCapability", |
| 94 | selector_config: "MLAPrefillSelectorConfig", |
| 95 | ) -> list[str]: |
| 96 | invalid_reasons: list[str] = [] |
| 97 | |
| 98 | if not cls.supports_compute_capability(device_capability): |
| 99 | invalid_reasons.append( |
| 100 | f"compute capability {device_capability.major}." |
| 101 | f"{device_capability.minor} not supported" |
| 102 | ) |
| 103 | |
| 104 | if not cls.supports_dtype(selector_config.dtype): |
| 105 | invalid_reasons.append(f"dtype {selector_config.dtype} not supported") |
| 106 | |
| 107 | if not cls.is_available(): |
| 108 | invalid_reasons.append("required dependencies not available") |
| 109 | |
| 110 | mla_dimensions = selector_config.mla_dimensions |
| 111 | if not cls.supports_mla_dimensions(mla_dimensions): |
| 112 | reason = ( |
| 113 | f"Model does not have supported MLA dimensions (got {mla_dimensions}" |
| 114 | ) |
| 115 | if ( |
| 116 | cls.supported_mla_dimensions |
| 117 | and mla_dimensions not in cls.supported_mla_dimensions |
| 118 | ): |
| 119 | supported = ", ".join( |
| 120 | str(dims) for dims in cls.supported_mla_dimensions |
| 121 | ) |
| 122 | reason += f"; supported: {supported}" |
| 123 | invalid_reasons.append(reason + ")") |
| 124 | |
| 125 | return invalid_reasons |
| 126 | |
| 127 | def __init__( |
| 128 | self, |
| 129 | num_heads: int, |
| 130 | scale: float, |
| 131 | kv_lora_rank: int, |
| 132 | qk_nope_head_dim: int, |
| 133 | qk_rope_head_dim: int, |
| 134 | v_head_dim: int, |
| 135 | vllm_config: "VllmConfig", |
| 136 | ) -> None: |
| 137 | self.num_heads = num_heads |
| 138 | self.scale = scale |
| 139 | self.kv_lora_rank = kv_lora_rank |
| 140 | self.qk_nope_head_dim = qk_nope_head_dim |
| 141 | self.qk_rope_head_dim = qk_rope_head_dim |
| 142 | self.v_head_dim = v_head_dim |
| 143 | self.vllm_config = vllm_config |
| 144 | |
| 145 | def clone(self) -> "MLAPrefillBackend": |
| 146 | return self.__class__( |
| 147 | num_heads=self.num_heads, |
| 148 | scale=self.scale, |
| 149 | kv_lora_rank=self.kv_lora_rank, |
| 150 | qk_nope_head_dim=self.qk_nope_head_dim, |
| 151 | qk_rope_head_dim=self.qk_rope_head_dim, |
| 152 | v_head_dim=self.v_head_dim, |
| 153 | vllm_config=self.vllm_config, |
| 154 | ) |
| 155 | |
| 156 | def prepare_metadata( # noqa: B027 |
| 157 | self, |
| 158 | prefill_metadata: "MLACommonPrefillMetadata", |
| 159 | ) -> None: |
| 160 | """Prepare backend-specific metadata before the forward pass. |
| 161 | |
| 162 | Called by the metadata builder after constructing the prefill metadata. |
| 163 | """ |
| 164 | self._prefill_metadata = prefill_metadata |
| 165 | |
| 166 | @abstractmethod |
| 167 | def run_prefill_new_tokens( |
| 168 | self, |
| 169 | q: torch.Tensor, |
| 170 | k: torch.Tensor, |
| 171 | v: torch.Tensor, |
| 172 | return_softmax_lse: bool, |
| 173 | out: torch.Tensor | None = None, |
| 174 | output_scale: torch.Tensor | None = None, |
| 175 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 176 | raise NotImplementedError |
| 177 | |
| 178 | @abstractmethod |
| 179 | def run_prefill_context_chunk( |
| 180 | self, |
| 181 | chunk: "MLACommonPrefillMetadata.ContextChunk", |
| 182 | q: torch.Tensor, |
| 183 | k: torch.Tensor, |
| 184 | v: torch.Tensor, |
| 185 | out: torch.Tensor | None = None, |
| 186 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 187 | raise NotImplementedError |