| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU MLA backend. |
| 4 | |
| 5 | This is a *reference-quality* MLA backend for the CPU platform. It is |
| 6 | intended to make DeepSeek-V2/V3 style models runnable on CPU (for |
| 7 | functional verification, tiny-model smoke tests and CI), not to be |
| 8 | performant. The prefill path is a plain PyTorch SDPA and the decode |
| 9 | path forwards to the existing CPU decode kernel |
| 10 | `torch.ops._C.mla_decode_kvcache` (see csrc/cpu/mla_decode.cpp). |
| 11 | |
| 12 | Key design points: |
| 13 | |
| 14 | * We inherit the shared MLA scaffolding (`MLACommonBackend`, |
| 15 | `MLACommonImpl`, `MLACommonMetadata`) so that the same |
| 16 | ``forward_impl`` in ``MLAAttention`` orchestrates weight-absorbed |
| 17 | decode (MQA) and non-absorbed prefill (MHA) for us. |
| 18 | * The parent ``MLACommonImpl.__init__`` tries to pick a GPU prefill |
| 19 | kernel (flash_attn / flashinfer / cudnn / trtllm) and raises when |
| 20 | none of them are available. On CPU none of them apply, so we |
| 21 | bypass that logic and set the attributes ourselves. |
| 22 | * Chunked prefill and prefix caching are disabled by the CPU platform |
| 23 | when ``use_mla`` is set. |
| 24 | * The CPU decode kernel only supports ``head_dim=576``, |
| 25 | ``v_head_dim=512`` and ``block_size=16`` today; the platform layer |
| 26 | forces ``block_size=16`` for MLA models. |
| 27 | """ |
| 28 | |
| 29 | from typing import ClassVar |
| 30 | |
| 31 | import torch |
| 32 | |
| 33 | from vllm import _custom_ops as ops |
| 34 | from vllm.config.cache import CacheDType |
| 35 | from vllm.logger import init_logger |
| 36 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 37 | MLACommonBackend, |
| 38 | MLACommonImpl, |
| 39 | MLACommonMetadata, |
| 40 | ) |
| 41 | from vllm.v1.attention.backend import ( |
| 42 | AttentionLayer, |
| 43 | AttentionType, |
| 44 | MultipleOf, |
| 45 | ) |
| 46 | |
| 47 | logger = init_logger(__name__) |
| 48 | |
| 49 | |
| 50 | class CPUMLABackend(MLACommonBackend): |
| 51 | """Attention backend descriptor for CPU MLA.""" |
| 52 | |
| 53 | supported_dtypes: ClassVar[list[torch.dtype]] = [ |
| 54 | torch.float16, |
| 55 | torch.bfloat16, |
| 56 | torch.float32, |
| 57 | ] |
| 58 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"] |
| 59 | |
| 60 | @classmethod |
| 61 | def get_supported_head_sizes(cls) -> list[int]: |
| 62 | # DeepSeek-V2/V3 latent cache size = kv_lora_rank(512)+rope(64)=576. |
| 63 | # The current CPU decode kernel only compiles for this size. |
| 64 | return [576] |
| 65 | |
| 66 | @staticmethod |
| 67 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 68 | return [16] |
| 69 | |
| 70 | @classmethod |
| 71 | def supports_block_size(cls, block_size: int | None) -> bool: |
| 72 | return block_size is None or block_size == 16 |
| 73 | |
| 74 | @staticmethod |
| 75 | def get_name() -> str: |
| 76 | return "CPU_MLA" |
| 77 | |
| 78 | @staticmethod |
| 79 | def get_impl_cls() -> type["CPUMLAImpl"]: |
| 80 | return CPUMLAImpl |
| 81 | |
| 82 | |
| 83 | class CPUMLAImpl(MLACommonImpl[MLACommonMetadata]): |
| 84 | """CPU implementation of MLA attention. |
| 85 | |
| 86 | See module docstring for the overall design. This class only |
| 87 | overrides the parts of the parent that would otherwise pull in |
| 88 | CUDA-only dependencies. |
| 89 | """ |
| 90 | |
| 91 | # Decode always returns the softmax lse if the parent needs it, but |
| 92 | # CPU decode currently only produces the output tensor. Report False |
| 93 | # so DCP / cross-batch reductions are not requested from us. |
| 94 | can_return_lse_for_decode: bool = False |
| 95 | |
| 96 | def __init__( |
| 97 | self, |
| 98 | num_heads: int, |
| 99 | head_size: int, |
| 100 | scale: float, |
| 101 | num_kv_heads: int, |
| 102 | alibi_slopes: list[float] | None, |
| 103 | sliding_window: int | None, |
| 104 | kv_cache_dtype: str, |
| 105 | logits_soft_cap: float | None, |
| 106 | attn_type: str, |
| 107 | kv_sharing_target_layer_name: str | None, |
| 108 | # MLA-specific arguments (forwarded from `MLAAttention`). |
| 109 | **mla_args, |
| 110 | ) -> None: |
| 111 | # NOTE: We deliberately skip ``MLACommonImpl.__init__`` because it |
| 112 | # tries to pick a GPU prefill kernel and raises on CPU. Instead we |
| 113 | # replicate the attribute setup here and then wire our own hooks. |
| 114 | if kv_sharing_target_layer_name is not None: |
| 115 | raise NotImplementedError("KV sharing is not supported for MLA") |
| 116 | if attn_type != AttentionType.DECODER: |
| 117 | raise NotImplementedError("CPU MLA only supports decoder self-attention.") |
| 118 | unsupported = [alibi_slopes, sliding_window, logits_soft_cap] |
| 119 | if any(unsupported): |
| 120 | raise NotImplementedError( |
| 121 | "CPU MLA does not support alibi_slopes / sliding_window / " |
| 122 | "logits_soft_cap." |
| 123 | ) |
| 124 | |
| 125 | self.num_heads = num_heads |
| 126 | self.head_size = head_size |
| 127 | self.scale = float(scale) |
| 128 | self.num_kv_heads = num_kv_heads |
| 129 | self.kv_cache_dtype = kv_cache_dtype |
| 130 | |
| 131 | self.q_lora_rank = mla_args["q_lora_rank"] |
| 132 | self.kv_lora_rank = mla_args["kv_lora_rank"] |
| 133 | self.qk_nope_head_dim = mla_args["qk_nope_head_dim"] |
| 134 | self.qk_rope_head_dim = mla_args["qk_rope_head_dim"] |
| 135 | self.qk_head_dim = mla_args["qk_head_dim"] |
| 136 | self.v_head_dim = mla_args["v_head_dim"] |
| 137 | self.kv_b_proj = mla_args["kv_b_proj"] |
| 138 | self.indexer = mla_args.get("indexer") |
| 139 | self.q_pad_num_heads = mla_args.get("q_pad_num_heads") |
| 140 | self.supports_quant_query_input = False |
| 141 | |
| 142 | # Never take the flashinfer specialisation of ``_concat_k_nope_k_pe``. |
| 143 | self._use_flashinfer_concat_mla_k = False |
| 144 | |
| 145 | # We do not need V padding since our SDPA prefill takes k / v with |
| 146 | # matching head dims. |
| 147 | self._pad_v = False |
| 148 | |
| 149 | # Distributed / context-parallel bookkeeping expected by the parent. |
| 150 | self.dcp_world_size = 1 |
| 151 | self.cp_kv_cache_interleave_size = 1 |
| 152 | |
| 153 | # ------------------------------------------------------------------ |
| 154 | # Decode (MQA) |
| 155 | # ------------------------------------------------------------------ |
| 156 | def forward_mqa( |
| 157 | self, |
| 158 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 159 | kv_c_and_k_pe_cache: torch.Tensor, |
| 160 | attn_metadata: MLACommonMetadata, |
| 161 | layer: AttentionLayer, |
| 162 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 163 | assert attn_metadata.decode is not None |
| 164 | # The parent hands us either a tuple (q_nope, q_pe) or a single |
| 165 | # tensor already concatenated. CPU decode kernel wants a single |
| 166 | # (bs, num_heads, head_dim=kv_lora_rank+rope) contiguous tensor. |
| 167 | if isinstance(q, tuple): |
| 168 | q = torch.cat(q, dim=-1) |
| 169 | q = q.contiguous() |
| 170 | |
| 171 | bs, num_heads, _ = q.shape |
| 172 | out = torch.zeros( |
| 173 | bs, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device |
| 174 | ) |
| 175 | |
| 176 | # The CPU kernel is templated on int32 block_tables / seq_lens. |
| 177 | block_table = attn_metadata.decode.block_table.to(torch.int32) |
| 178 | seq_lens = attn_metadata.decode.seq_lens.to(torch.int32) |
| 179 | |
| 180 | ops.mla_decode_kvcache_cpu( |
| 181 | out, |
| 182 | q, |
| 183 | kv_c_and_k_pe_cache, |
| 184 | self.scale, |
| 185 | block_table, |
| 186 | seq_lens, |
| 187 | ) |
| 188 | # The parent only consumes the LSE when DCP is active, which we |
| 189 | # disable, so returning None here is safe. |
| 190 | return out, None |
| 191 | |
| 192 | # ------------------------------------------------------------------ |
| 193 | # KV cache write |
| 194 | # ------------------------------------------------------------------ |
| 195 | def do_kv_cache_update( |
| 196 | self, |
| 197 | kv_c_normed: torch.Tensor, |
| 198 | k_pe: torch.Tensor, |
| 199 | kv_cache: torch.Tensor, |
| 200 | slot_mapping: torch.Tensor, |
| 201 | kv_cache_dtype: str, |
| 202 | k_scale: torch.Tensor, |
| 203 | ) -> None: |
| 204 | if kv_cache.numel() == 0: |
| 205 | return |
| 206 | k_pe_sq = k_pe.squeeze(1) |
| 207 | kv_lora_rank = kv_c_normed.shape[-1] |
| 208 | pe_dim = k_pe_sq.shape[-1] |
| 209 | flat = kv_cache.view(-1, kv_lora_rank + pe_dim) |
| 210 | slots = slot_mapping.flatten().to(torch.long) |
| 211 | # Only write valid slots. Padding slots are set to -1 by the |
| 212 | # scheduler and must be skipped. |
| 213 | valid_mask = slots >= 0 |
| 214 | if not valid_mask.all().item(): |
| 215 | slots = slots[valid_mask] |
| 216 | kv_c_normed = kv_c_normed[valid_mask] |
| 217 | k_pe_sq = k_pe_sq[valid_mask] |
| 218 | target_dtype = flat.dtype |
| 219 | flat[slots, :kv_lora_rank] = kv_c_normed.to(target_dtype) |
| 220 | flat[slots, kv_lora_rank:] = k_pe_sq.view(-1, pe_dim).to(target_dtype) |