| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU MLA prefill backend using zentorch SDPA on AMD Zen CPUs. |
| 4 | |
| 5 | Subclasses the generic CPU SDPA backend and replaces the per-request attention |
| 6 | with ``zentorch_sdpa``, which fuses the scale, causal mask, softmax and both |
| 7 | matmuls of the base implementation into a single call. |
| 8 | |
| 9 | ``zentorch_sdpa`` does not return a usable log-sum-exp, so requests that need |
| 10 | one (context chunks, prefix/suffix merge) are delegated to the base backend. |
| 11 | |
| 12 | Availability is reported through :meth:`is_available`, so the prefill selector |
| 13 | gates this backend once at selection time rather than per attention call. |
| 14 | """ |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op |
| 19 | from vllm.v1.attention.backends.mla.prefill.cpu_sdpa import CPUSDPAMLAPrefillBackend |
| 20 | |
| 21 | |
| 22 | class ZenCPUSDPAMLAPrefillBackend(CPUSDPAMLAPrefillBackend): |
| 23 | """MLA prefill backend for AMD Zen CPUs backed by ``zentorch_sdpa``.""" |
| 24 | |
| 25 | @staticmethod |
| 26 | def get_name() -> str: |
| 27 | return "ZEN_CPU_SDPA_MLA" |
| 28 | |
| 29 | @classmethod |
| 30 | def is_available(cls) -> bool: |
| 31 | return has_zentorch_op(["zentorch_sdpa"]) |
| 32 | |
| 33 | @staticmethod |
| 34 | def _sdpa_layout( |
| 35 | q_r: torch.Tensor, |
| 36 | k_r: torch.Tensor, |
| 37 | v_r: torch.Tensor, |
| 38 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: |
| 39 | """Convert one request's [S, NH, D] tensors to the [1, NH, S, D] layout. |
| 40 | |
| 41 | MLA has ``v_head_dim < qk_head_dim``, which the fused kernel does not |
| 42 | accept, so V is padded up to the query head dim. The original V head |
| 43 | dim is returned so the caller can slice the padding off afterwards. |
| 44 | """ |
| 45 | h_q = q_r.shape[-1] |
| 46 | h_v = v_r.shape[-1] |
| 47 | pad = h_q - h_v |
| 48 | |
| 49 | qs = q_r.transpose(0, 1).unsqueeze(0).contiguous() |
| 50 | ks = k_r.transpose(0, 1).unsqueeze(0).contiguous() |
| 51 | vs = v_r.transpose(0, 1).unsqueeze(0).contiguous() |
| 52 | if pad > 0: |
| 53 | vs = torch.nn.functional.pad(vs, (0, pad)) |
| 54 | return qs, ks, vs, h_v |
| 55 | |
| 56 | def _ragged_sdpa( |
| 57 | self, |
| 58 | q: torch.Tensor, |
| 59 | k: torch.Tensor, |
| 60 | v: torch.Tensor, |
| 61 | q_cu_seq_lens: torch.Tensor, |
| 62 | kv_cu_seq_lens: torch.Tensor, |
| 63 | causal: bool, |
| 64 | return_softmax_lse: bool, |
| 65 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 66 | if return_softmax_lse: |
| 67 | # zentorch_sdpa's second return is not a usable LSE, so fall back |
| 68 | # to the base fp32 implementation whenever one is required. |
| 69 | return super()._ragged_sdpa( |
| 70 | q, |
| 71 | k, |
| 72 | v, |
| 73 | q_cu_seq_lens, |
| 74 | kv_cu_seq_lens, |
| 75 | causal, |
| 76 | return_softmax_lse, |
| 77 | ) |
| 78 | |
| 79 | num_reqs = q_cu_seq_lens.numel() - 1 |
| 80 | assert num_reqs == kv_cu_seq_lens.numel() - 1 |
| 81 | |
| 82 | out_chunks: list[torch.Tensor] = [] |
| 83 | for r in range(num_reqs): |
| 84 | q_start = int(q_cu_seq_lens[r]) |
| 85 | q_end = int(q_cu_seq_lens[r + 1]) |
| 86 | kv_start = int(kv_cu_seq_lens[r]) |
| 87 | kv_end = int(kv_cu_seq_lens[r + 1]) |
| 88 | |
| 89 | qs, ks, vs, h_v = self._sdpa_layout( |
| 90 | q[q_start:q_end], |
| 91 | k[kv_start:kv_end], |
| 92 | v[kv_start:kv_end], |
| 93 | ) |
| 94 | |
| 95 | attn, _ = torch.ops.zentorch.zentorch_sdpa( |
| 96 | qs, |
| 97 | ks, |
| 98 | vs, |
| 99 | 0.0, |
| 100 | causal, |
| 101 | attn_mask=None, |
| 102 | scale=self.scale, |
| 103 | ) |
| 104 | out_chunks.append(attn[0].transpose(0, 1)[..., :h_v].to(v.dtype)) |
| 105 | |
| 106 | return torch.cat(out_chunks, dim=0) |