| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU MLA prefill backend implemented with dense per-request attention. |
| 4 | |
| 5 | This backend is selected on CPU devices, where flash-attn is unavailable. It |
| 6 | computes MLA prefill attention by slicing the ragged (varlen) q/k/v tensors |
| 7 | per request and evaluating dense attention in PyTorch. |
| 8 | """ |
| 9 | |
| 10 | from typing import TYPE_CHECKING |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 18 | MLACommonPrefillMetadata, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | class CPUSDPAMLAPrefillBackend(MLAPrefillBackend): |
| 23 | """MLA prefill backend for CPU using dense PyTorch attention. |
| 24 | |
| 25 | flash-attn is not installable on CPU, so this backend provides an |
| 26 | equivalent prefill implementation using per-request dense attention. |
| 27 | """ |
| 28 | |
| 29 | @staticmethod |
| 30 | def get_name() -> str: |
| 31 | return "CPU_SDPA_MLA" |
| 32 | |
| 33 | @classmethod |
| 34 | def is_available(cls) -> bool: |
| 35 | # SDPA is always available on CPU. |
| 36 | return True |
| 37 | |
| 38 | def supports_out(self) -> bool: |
| 39 | return True |
| 40 | |
| 41 | def _ragged_sdpa( |
| 42 | self, |
| 43 | q: torch.Tensor, |
| 44 | k: torch.Tensor, |
| 45 | v: torch.Tensor, |
| 46 | q_cu_seq_lens: torch.Tensor, |
| 47 | kv_cu_seq_lens: torch.Tensor, |
| 48 | causal: bool, |
| 49 | return_softmax_lse: bool, |
| 50 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 51 | # q/k/v: [num_tokens, num_heads, head_dim] (ragged, concatenated). |
| 52 | # q_cu_seq_lens/kv_cu_seq_lens: [num_reqs + 1] row offsets for each |
| 53 | # request's query and key/value tokens. |
| 54 | num_reqs = q_cu_seq_lens.numel() - 1 |
| 55 | assert num_reqs == kv_cu_seq_lens.numel() - 1 |
| 56 | |
| 57 | out_chunks: list[torch.Tensor] = [] |
| 58 | lse_chunks: list[torch.Tensor] = [] |
| 59 | for r in range(num_reqs): |
| 60 | q_start = int(q_cu_seq_lens[r]) |
| 61 | q_end = int(q_cu_seq_lens[r + 1]) |
| 62 | kv_start = int(kv_cu_seq_lens[r]) |
| 63 | kv_end = int(kv_cu_seq_lens[r + 1]) |
| 64 | |
| 65 | q_r = q[q_start:q_end].transpose(0, 1).float() # [NH, Sq, Hq] |
| 66 | k_r = k[kv_start:kv_end].transpose(0, 1).float() # [NH, Sk, Hq] |
| 67 | v_r = v[kv_start:kv_end].transpose(0, 1).float() # [NH, Sk, Hv] |
| 68 | |
| 69 | attn_scores = torch.matmul(q_r, k_r.transpose(-2, -1)) * self.scale |
| 70 | if causal: |
| 71 | causal_mask = torch.ones( |
| 72 | (q_end - q_start, kv_end - kv_start), |
| 73 | dtype=torch.bool, |
| 74 | device=attn_scores.device, |
| 75 | ).triu(1) |
| 76 | attn_scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf")) |
| 77 | |
| 78 | attn_probs = torch.softmax(attn_scores, dim=-1, dtype=torch.float32) |
| 79 | out_r = torch.matmul(attn_probs, v_r).transpose(0, 1).to(v.dtype) |
| 80 | out_chunks.append(out_r) # [Sq, NH, Hv] |
| 81 | if return_softmax_lse: |
| 82 | lse_chunks.append(torch.logsumexp(attn_scores, dim=-1)) |
| 83 | |
| 84 | out = torch.cat(out_chunks, dim=0) |
| 85 | if not return_softmax_lse: |
| 86 | return out |
| 87 | lse = torch.cat(lse_chunks, dim=1) |
| 88 | return out, lse |
| 89 | |
| 90 | def run_prefill_new_tokens( |
| 91 | self, |
| 92 | q: torch.Tensor, |
| 93 | k: torch.Tensor, |
| 94 | v: torch.Tensor, |
| 95 | return_softmax_lse: bool, |
| 96 | out: torch.Tensor | None = None, |
| 97 | output_scale: torch.Tensor | None = None, |
| 98 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 99 | assert output_scale is None, ( |
| 100 | "CPU MLA prefill backend does not support fused FP8 output" |
| 101 | ) |
| 102 | out_t = self._ragged_sdpa( |
| 103 | q, |
| 104 | k, |
| 105 | v, |
| 106 | self._prefill_metadata.query_start_loc, |
| 107 | self._prefill_metadata.query_start_loc, |
| 108 | causal=True, |
| 109 | return_softmax_lse=return_softmax_lse, |
| 110 | ) |
| 111 | if out is not None and not return_softmax_lse: |
| 112 | assert isinstance(out_t, torch.Tensor) |
| 113 | out.copy_(out_t) |
| 114 | return out |
| 115 | return out_t |
| 116 | |
| 117 | def run_prefill_context_chunk( |
| 118 | self, |
| 119 | chunk: "MLACommonPrefillMetadata.ContextChunk", |
| 120 | q: torch.Tensor, |
| 121 | k: torch.Tensor, |
| 122 | v: torch.Tensor, |
| 123 | out: torch.Tensor | None = None, |
| 124 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 125 | assert out is None, "CPU MLA context prefill does not support in-place output" |
| 126 | result = self._ragged_sdpa( |
| 127 | q, |
| 128 | k, |
| 129 | v, |
| 130 | chunk.query_start_loc, |
| 131 | chunk.cu_seq_lens, |
| 132 | causal=False, |
| 133 | return_softmax_lse=True, |
| 134 | ) |
| 135 | assert isinstance(result, tuple) |
| 136 | out, lse = result |
| 137 | return out, lse |