[Attention][CPU] Run Zen CPU encoder attention on zentorch SDPA (#54508)
Signed-off-by: priyansh jain <priyansh.jain2@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Priyansh Jain committed
Sep 24, 2026 at 14:35 UTC
5963795ec7d7beaf34a05a213cce113c6508a8d4
5 files changed
+319
-4
setup.py
+1
-1
@@ -1529,7 +1529,7 @@ setup(
1529
install_requires=get_requirements(),
1530
extras_require={
1531
# AMD Zen CPU optimizations via zentorch
1532
- "zen": ["zentorch==2.13.0.0"],
1532
+ "zen": ["zentorch==2.13.0.1"],
1533
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
1534
"tensorizer": ["tensorizer==2.10.1"],
1535
"fastsafetensors": ["fastsafetensors >= 0.3.3"],
tests/kernels/attention/test_cpu_attn.py
+133
@@ -3,13 +3,20 @@
3
4
import functools
5
import math
6
+from types import SimpleNamespace
7
8
import pytest
9
import torch
10
11
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
12
from vllm.platforms import CpuArchEnum, current_platform
13
from vllm.utils.torch_utils import set_random_seed
14
+from vllm.v1.attention.backend import AttentionType
15
from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
16
+from vllm.v1.attention.backends.zentorch_sdpa import (
17
+ should_use_zentorch_sdpa,
18
+ zentorch_sdpa_attn,
19
+)
20
21
if not current_platform.is_cpu():
22
pytest.skip("skipping CPU-only tests", allow_module_level=True)
@@ -47,6 +54,24 @@ _FP8_RTOL = 0.1
54
ENCODER_SEQ_LENS = [
55
[1, 678, 2367, 145, 4162, 36, 7812],
56
]
57
+# Uniform (dense zentorch_sdpa call) and ragged (per-sequence loop).
58
+ZEN_ENCODER_SEQ_LENS = [
59
+ [128, 128, 128],
60
+ [1, 67, 233, 5],
61
+]
62
+_skip_no_zentorch_sdpa = pytest.mark.skipif(
63
+ not current_platform.is_zen_cpu() or not has_zentorch_op(["zentorch_sdpa_attn"]),
64
+ reason="zentorch_sdpa requires a Zen CPU with the op registered",
65
+)
66
+# should_use_zentorch_sdpa additionally gates each dtype on its ISA.
67
+_skip_no_avx512_bf16 = pytest.mark.skipif(
68
+ not torch.cpu._is_avx512_bf16_supported(),
69
+ reason="zentorch_sdpa bfloat16 requires AVX512-BF16",
70
+)
71
+_skip_no_avx512 = pytest.mark.skipif(
72
+ not torch.cpu._is_avx512_supported(),
73
+ reason="zentorch_sdpa float32 requires AVX512",
74
+)
75
76
77
def get_attn_isa(
@@ -205,6 +230,7 @@ def ref_varlen_encoder_attn(
230
seq_lens: list[int],
231
scale: float,
232
sliding_window: int | None = None,
233
+ alibi_slopes: torch.Tensor | None = None,
234
) -> torch.Tensor:
235
num_seqs = len(seq_lens)
236
dtype = query.dtype
@@ -232,6 +258,11 @@ def ref_varlen_encoder_attn(
258
else:
259
mask = empty_mask.logical_not()
260
261
+ if alibi_slopes is not None:
262
+ q_pos = torch.arange(0, seq_len)[None, :, None]
263
+ kv_pos = torch.arange(0, seq_len)[None, None, :]
264
+ attn += -alibi_slopes[:, None, None].float() * (q_pos - kv_pos)
265
+
266
attn.masked_fill_(mask, float("-inf"))
267
attn = torch.softmax(attn, dim=-1)
268
out = torch.einsum("hqk,khd->qhd", attn, v).to(dtype=dtype)
@@ -772,6 +803,108 @@ def test_varlen_encoder_attention_amx(
803
)
804
805
806
+@torch.inference_mode()
807
+def varlen_encoder_zentorch_sdpa(
808
+ seq_lens: list[int],
809
+ num_heads: tuple[int, int],
810
+ head_size: int,
811
+ sliding_window: int | None,
812
+ dtype: torch.dtype,
813
+ use_alibi: bool,
814
+) -> None:
815
+ set_random_seed(0)
816
+ num_query_heads, num_kv_heads = num_heads
817
+ assert num_query_heads % num_kv_heads == 0
818
+ scale = head_size**-0.5
819
+ token_num = sum(seq_lens)
820
+
821
+ query = tensor_cache(
822
+ token_num * num_query_heads * head_size, dtype, tag="query"
823
+ ).view(token_num, num_query_heads, head_size)
824
+ key_value = tensor_cache(
825
+ 2 * token_num * num_kv_heads * head_size, dtype, tag="kv"
826
+ ).view(2, token_num, num_kv_heads, head_size)
827
+ key, value = key_value.unbind(0)
828
+
829
+ query_start_loc = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
830
+ torch.cumsum(torch.tensor(seq_lens, dtype=torch.int32), 0, out=query_start_loc[1:])
831
+ alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None
832
+ output = torch.empty_like(query)
833
+
834
+ zentorch_sdpa_attn(
835
+ query,
836
+ key,
837
+ value,
838
+ output,
839
+ SimpleNamespace(query_start_loc=query_start_loc, causal=False),
840
+ scale,
841
+ sliding_window if sliding_window is not None else -1,
842
+ alibi_slopes,
843
+ )
844
+ ref_output = ref_varlen_encoder_attn(
845
+ query=query,
846
+ key=key,
847
+ value=value,
848
+ seq_lens=seq_lens,
849
+ scale=scale,
850
+ sliding_window=sliding_window,
851
+ alibi_slopes=alibi_slopes,
852
+ )
853
+ torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
854
+
855
+
856
+_ZENTORCH_SDPA_DTYPES = [
857
+ pytest.param(torch.bfloat16, marks=_skip_no_avx512_bf16),
858
+ pytest.param(torch.float32, marks=_skip_no_avx512),
859
+]
860
+
861
+
862
+@_skip_no_zentorch_sdpa
863
+@pytest.mark.parametrize("seq_lens", ZEN_ENCODER_SEQ_LENS)
864
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
865
+@pytest.mark.parametrize("head_size", [128])
866
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
867
+@pytest.mark.parametrize("dtype", _ZENTORCH_SDPA_DTYPES)
868
+@pytest.mark.parametrize("use_alibi", [False, True])
869
+def test_zentorch_sdpa_attn(
870
+ seq_lens: list[int],
871
+ num_heads: tuple[int, int],
872
+ head_size: int,
873
+ sliding_window: int | None,
874
+ dtype: torch.dtype,
875
+ use_alibi: bool,
876
+) -> None:
877
+ varlen_encoder_zentorch_sdpa(
878
+ seq_lens=seq_lens,
879
+ num_heads=num_heads,
880
+ head_size=head_size,
881
+ sliding_window=sliding_window,
882
+ dtype=dtype,
883
+ use_alibi=use_alibi,
884
+ )
885
+
886
+
887
+@_skip_no_zentorch_sdpa
888
+@pytest.mark.parametrize(
889
+ "attn_type", [AttentionType.ENCODER_ONLY, AttentionType.ENCODER]
890
+)
891
+@pytest.mark.parametrize("dtype", _ZENTORCH_SDPA_DTYPES)
892
+def test_should_use_zentorch_sdpa(attn_type: str, dtype: torch.dtype) -> None:
893
+ assert should_use_zentorch_sdpa(attn_type, dtype)
894
+
895
+
896
+@_skip_no_zentorch_sdpa
897
+@pytest.mark.parametrize(
898
+ "attn_type,dtype",
899
+ [
900
+ (AttentionType.DECODER, torch.bfloat16),
901
+ (AttentionType.ENCODER_ONLY, torch.float16),
902
+ ],
903
+)
904
+def test_should_not_use_zentorch_sdpa(attn_type: str, dtype: torch.dtype) -> None:
905
+ assert not should_use_zentorch_sdpa(attn_type, dtype)
906
+
907
+
908
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
909
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
910
@pytest.mark.parametrize("num_heads", NUM_HEADS)
vllm/v1/attention/backends/cpu_attn.py
+26
@@ -13,6 +13,7 @@ from vllm import _custom_ops as ops
13
from vllm import envs
14
from vllm.config import (
15
VllmConfig,
16
+ get_current_vllm_config,
17
get_layers_from_vllm_config,
18
)
19
from vllm.logger import init_logger
@@ -31,6 +32,10 @@ from vllm.v1.attention.backend import (
32
from vllm.v1.attention.backends.utils import (
33
get_num_attention_heads_from_layers,
34
)
35
+from vllm.v1.attention.backends.zentorch_sdpa import (
36
+ should_use_zentorch_sdpa,
37
+ zentorch_sdpa_attn,
38
+)
39
from vllm.v1.kv_cache_interface import (
40
AttentionSpec,
41
CrossAttentionSpec,
@@ -334,6 +339,12 @@ class CPUAttentionBackendImpl(AttentionImpl):
339
"heads in the layer"
340
)
341
342
+ vllm_config = get_current_vllm_config()
343
+ self.use_zentorch_sdpa = should_use_zentorch_sdpa(
344
+ attn_type,
345
+ vllm_config.model_config.dtype,
346
+ )
347
+
348
def forward(
349
self,
350
layer: AttentionLayer,
@@ -383,6 +394,21 @@ class CPUAttentionBackendImpl(AttentionImpl):
394
AttentionType.ENCODER,
395
)
396
if is_encoder_attention:
397
+ if self.use_zentorch_sdpa:
398
+ # Encoder attention never reads the KV cache back, so the
399
+ # zentorch path attends the packed QKV directly instead of
400
+ # staging it through the scratch encoder cache.
401
+ zentorch_sdpa_attn(
402
+ query[:num_actual_tokens],
403
+ key[:num_actual_tokens],
404
+ value[:num_actual_tokens],
405
+ output[:num_actual_tokens],
406
+ attn_metadata,
407
+ self.scale,
408
+ self.sliding_window,
409
+ self.alibi_slopes,
410
+ )
411
+ return output
412
# For encoder attention,
413
kv_cache = attn_metadata.encoder_cache
414
vllm/v1/attention/backends/zentorch_sdpa.py
new
+148
@@ -0,0 +1,148 @@
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
vllm/v1/attention/ops/vit_attn_wrappers.py
+11
-3
@@ -241,9 +241,17 @@ def apply_sdpa(
241
(batch_size x seq_len x num_heads x head_size)
242
"""
243
q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
244
- output = F.scaled_dot_product_attention(
245
- q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
246
- )
244
+ if current_platform.is_zen_cpu() and hasattr(torch.ops.zentorch, "zentorch_sdpa"):
245
+ # Schema: (query, key, value, dropout_p=0., is_causal=False, *,
246
+ # attn_mask=None, scale=None) -> (attention, logsumexp).
247
+ # No enable_gqa: GQA is inferred from key.size(1) vs query.size(1).
248
+ output, _ = torch.ops.zentorch.zentorch_sdpa(
249
+ q, k, v, dropout_p=0.0, is_causal=False, scale=scale
250
+ )
251
+ else:
252
+ output = F.scaled_dot_product_attention(
253
+ q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
254
+ )
255
output = einops.rearrange(output, "b h s d -> b s h d ")
256
return output
257