| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Encoder attention through the zentorch SDPA kernel on Zen CPUs.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op |
| 12 | from vllm.v1.attention.backend import AttentionType |
| 13 | |
| 14 | if TYPE_CHECKING: |
| 15 | from vllm.v1.attention.backends.cpu_attn import CPUAttentionMetadata |
| 16 | |
| 17 | __all__ = ["should_use_zentorch_sdpa", "zentorch_sdpa_attn"] |
| 18 | |
| 19 | |
| 20 | def should_use_zentorch_sdpa( |
| 21 | attn_type: str, |
| 22 | dtype: torch.dtype, |
| 23 | ) -> bool: |
| 24 | """True when encoder attention should dispatch to zentorch_sdpa. |
| 25 | |
| 26 | Full, sliding-window, and ALiBi encoder layers all dispatch. A window is |
| 27 | a single-head [1, 1, S, S] additive mask; ALiBi is a per-query-head |
| 28 | [1, H, S, S] bias. The zentorch_sdpa_attn op builds both. |
| 29 | |
| 30 | dtype must match the op's ISA gate (bf16: AVX512-BF16, fp32: AVX512); |
| 31 | fp16 needs AVX512-FP16, which torch exposes no query for, so it stays |
| 32 | native. |
| 33 | """ |
| 34 | if attn_type not in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): |
| 35 | return False |
| 36 | if not has_zentorch_op(["zentorch_sdpa_attn"]): |
| 37 | return False |
| 38 | if dtype == torch.bfloat16: |
| 39 | return torch.cpu._is_avx512_bf16_supported() |
| 40 | if dtype == torch.float32: |
| 41 | return torch.cpu._is_avx512_supported() |
| 42 | return False |
| 43 | |
| 44 | |
| 45 | def zentorch_sdpa_attn( |
| 46 | query: torch.Tensor, |
| 47 | key: torch.Tensor, |
| 48 | value: torch.Tensor, |
| 49 | output: torch.Tensor, |
| 50 | attn_metadata: CPUAttentionMetadata, |
| 51 | scale: float, |
| 52 | sliding_window: int = -1, |
| 53 | alibi_slopes: torch.Tensor | None = None, |
| 54 | ) -> torch.Tensor: |
| 55 | """Run encoder / encoder-only attention with zentorch_sdpa. |
| 56 | |
| 57 | Encoder attention reads no KV cache, so the packed |
| 58 | [num_tokens, num_heads, head_size] query/key/value are attended in place |
| 59 | and written straight into `output`. |
| 60 | |
| 61 | Args: |
| 62 | query: Query of shape [num_tokens, num_heads, head_size]. |
| 63 | key: Key of shape [num_tokens, num_kv_heads, head_size]. |
| 64 | value: Value of shape [num_tokens, num_kv_heads, head_size]. |
| 65 | output: Pre-allocated output, same shape as `query`. |
| 66 | attn_metadata: Metadata carrying the per-sequence token offsets and |
| 67 | whether attention is causal. |
| 68 | scale: Softmax scale. |
| 69 | sliding_window: Encoder window size, or -1 for full bidirectional |
| 70 | attention. Handed to the op, which masks token i to |
| 71 | [i-(W-1), i+(W-1)]. |
| 72 | alibi_slopes: Per-query-head ALiBi slopes, or None. Handed to the op, |
| 73 | which applies them as an additive per-head bias. |
| 74 | |
| 75 | Returns: |
| 76 | `output`, filled in place. |
| 77 | |
| 78 | """ |
| 79 | # zentorch_sdpa_attn builds the ALiBi mask in query.dtype. A bf16 per-head |
| 80 | # mask is wrong when the query head count is not a power of two, so run |
| 81 | # that path in fp32 and write bf16 back. |
| 82 | if alibi_slopes is not None and query.dtype == torch.bfloat16: |
| 83 | out32 = torch.empty_like(query, dtype=torch.float32) |
| 84 | zentorch_sdpa_attn( |
| 85 | query.float(), |
| 86 | key.float(), |
| 87 | value.float(), |
| 88 | out32, |
| 89 | attn_metadata, |
| 90 | scale, |
| 91 | sliding_window, |
| 92 | alibi_slopes, |
| 93 | ) |
| 94 | output.copy_(out32.to(dtype=output.dtype)) |
| 95 | return output |
| 96 | |
| 97 | # The op counts window tokens on each side of the diagonal; vLLM's window |
| 98 | # includes the diagonal itself. |
| 99 | window = -1 if sliding_window in (None, -1) else sliding_window - 1 |
| 100 | |
| 101 | # zentorch_sdpa takes no enable_gqa argument but derives the KV head count |
| 102 | # from the key tensor and maps each query head onto its KV head, so GQA/MQA |
| 103 | # needs no expansion here. |
| 104 | query = query.movedim(0, query.dim() - 2) |
| 105 | key = key.movedim(0, key.dim() - 2) |
| 106 | value = value.movedim(0, value.dim() - 2) |
| 107 | |
| 108 | start_loc = attn_metadata.query_start_loc.numpy() |
| 109 | seq_lens = start_loc[1:] - start_loc[:-1] |
| 110 | |
| 111 | # Encoder sequences are packed along the token dimension. When they all |
| 112 | # have the same length, recover a dense BHSD batch and invoke zentorch_sdpa |
| 113 | # once for the whole scheduler batch instead of once per sequence. |
| 114 | if len(seq_lens) > 0 and (seq_lens == seq_lens[0]).all(): |
| 115 | batch_size = len(seq_lens) |
| 116 | seq_len = int(seq_lens[0]) |
| 117 | |
| 118 | def _packed_hsd_to_bhsd(tensor: torch.Tensor) -> torch.Tensor: |
| 119 | return tensor.unflatten(1, (batch_size, seq_len)).permute(1, 0, 2, 3) |
| 120 | |
| 121 | torch.ops.zentorch.zentorch_sdpa_attn( |
| 122 | _packed_hsd_to_bhsd(query), |
| 123 | _packed_hsd_to_bhsd(key), |
| 124 | _packed_hsd_to_bhsd(value), |
| 125 | # The op writes [B, H, S, D]; hand it the matching view of the |
| 126 | # packed buffer so it stores there instead of into a temporary. |
| 127 | output.unflatten(0, (batch_size, seq_len)).permute(0, 2, 1, 3), |
| 128 | scale=scale, |
| 129 | is_causal=attn_metadata.causal, |
| 130 | alibi_slopes=alibi_slopes, |
| 131 | left_window_size=window, |
| 132 | right_window_size=window, |
| 133 | ) |
| 134 | return output |
| 135 | |
| 136 | for start_q, end_q in zip(start_loc[:-1], start_loc[1:], strict=True): |
| 137 | torch.ops.zentorch.zentorch_sdpa_attn( |
| 138 | query[None, :, start_q:end_q, :], |
| 139 | key[None, :, start_q:end_q, :], |
| 140 | value[None, :, start_q:end_q, :], |
| 141 | output[start_q:end_q, :, :].movedim(0, 1).unsqueeze(0), |
| 142 | scale=scale, |
| 143 | is_causal=attn_metadata.causal, |
| 144 | alibi_slopes=alibi_slopes, |
| 145 | left_window_size=window, |
| 146 | right_window_size=window, |
| 147 | ) |
| 148 | return output |