iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Attention][CPU] Run Zen CPU encoder attention on zentorch SDPA (#54508)

Browse Source 
Signed-off-by: priyansh jain <priyansh.jain2@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Priyansh Jain
2026-09-24 02:05:44 -07:00
committed by GitHub
co-authored by Cursor Andreas Karatzas
parent 14f98cbe5d
commit 5963795ec7
5 changed files with 319 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
setup.py
+1 -1
View File
@@ -1529,7 +1529,7 @@ setup(
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": ["zentorch==2.13.0.0"],
"zen": ["zentorch==2.13.0.1"],
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.3.3"],
tests/kernels/attention/test_cpu_attn.py
+133
View File
@@ -3,13 +3,20 @@
import functools
import math
from types import SimpleNamespace
import pytest
import torch
from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
from vllm.platforms import CpuArchEnum, current_platform
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.attention.backend import AttentionType
from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
from vllm.v1.attention.backends.zentorch_sdpa import (
should_use_zentorch_sdpa,
zentorch_sdpa_attn,
)
if not current_platform.is_cpu():
pytest.skip("skipping CPU-only tests", allow_module_level=True)
@@ -47,6 +54,24 @@ _FP8_RTOL = 0.1
ENCODER_SEQ_LENS = [
[1, 678, 2367, 145, 4162, 36, 7812],
]
# Uniform (dense zentorch_sdpa call) and ragged (per-sequence loop).
ZEN_ENCODER_SEQ_LENS = [
[128, 128, 128],
[1, 67, 233, 5],
]
_skip_no_zentorch_sdpa = pytest.mark.skipif(
not current_platform.is_zen_cpu() or not has_zentorch_op(["zentorch_sdpa_attn"]),
reason="zentorch_sdpa requires a Zen CPU with the op registered",
)
# should_use_zentorch_sdpa additionally gates each dtype on its ISA.
_skip_no_avx512_bf16 = pytest.mark.skipif(
not torch.cpu._is_avx512_bf16_supported(),
reason="zentorch_sdpa bfloat16 requires AVX512-BF16",
)
_skip_no_avx512 = pytest.mark.skipif(
not torch.cpu._is_avx512_supported(),
reason="zentorch_sdpa float32 requires AVX512",
)
def get_attn_isa(
@@ -205,6 +230,7 @@ def ref_varlen_encoder_attn(
seq_lens: list[int],
scale: float,
sliding_window: int | None = None,
alibi_slopes: torch.Tensor | None = None,
) -> torch.Tensor:
num_seqs = len(seq_lens)
dtype = query.dtype
@@ -232,6 +258,11 @@ def ref_varlen_encoder_attn(
else:
mask = empty_mask.logical_not()
if alibi_slopes is not None:
q_pos = torch.arange(0, seq_len)[None, :, None]
kv_pos = torch.arange(0, seq_len)[None, None, :]
attn += -alibi_slopes[:, None, None].float() * (q_pos - kv_pos)
attn.masked_fill_(mask, float("-inf"))
attn = torch.softmax(attn, dim=-1)
out = torch.einsum("hqk,khd->qhd", attn, v).to(dtype=dtype)
@@ -772,6 +803,108 @@ def test_varlen_encoder_attention_amx(
)
@torch.inference_mode()
def varlen_encoder_zentorch_sdpa(
seq_lens: list[int],
num_heads: tuple[int, int],
head_size: int,
sliding_window: int | None,
dtype: torch.dtype,
use_alibi: bool,
) -> None:
set_random_seed(0)
num_query_heads, num_kv_heads = num_heads
assert num_query_heads % num_kv_heads == 0
scale = head_size**-0.5
token_num = sum(seq_lens)
query = tensor_cache(
token_num * num_query_heads * head_size, dtype, tag="query"
).view(token_num, num_query_heads, head_size)
key_value = tensor_cache(
2 * token_num * num_kv_heads * head_size, dtype, tag="kv"
).view(2, token_num, num_kv_heads, head_size)
key, value = key_value.unbind(0)
query_start_loc = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
torch.cumsum(torch.tensor(seq_lens, dtype=torch.int32), 0, out=query_start_loc[1:])
alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None
output = torch.empty_like(query)
zentorch_sdpa_attn(
query,
key,
value,
output,
SimpleNamespace(query_start_loc=query_start_loc, causal=False),
scale,
sliding_window if sliding_window is not None else -1,
alibi_slopes,
)
ref_output = ref_varlen_encoder_attn(
query=query,
key=key,
value=value,
seq_lens=seq_lens,
scale=scale,
sliding_window=sliding_window,
alibi_slopes=alibi_slopes,
)
torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
_ZENTORCH_SDPA_DTYPES = [
pytest.param(torch.bfloat16, marks=_skip_no_avx512_bf16),
pytest.param(torch.float32, marks=_skip_no_avx512),
]
@_skip_no_zentorch_sdpa
@pytest.mark.parametrize("seq_lens", ZEN_ENCODER_SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", [128])
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
@pytest.mark.parametrize("dtype", _ZENTORCH_SDPA_DTYPES)
@pytest.mark.parametrize("use_alibi", [False, True])
def test_zentorch_sdpa_attn(
seq_lens: list[int],
num_heads: tuple[int, int],
head_size: int,
sliding_window: int | None,
dtype: torch.dtype,
use_alibi: bool,
) -> None:
varlen_encoder_zentorch_sdpa(
seq_lens=seq_lens,
num_heads=num_heads,
head_size=head_size,
sliding_window=sliding_window,
dtype=dtype,
use_alibi=use_alibi,
)
@_skip_no_zentorch_sdpa
@pytest.mark.parametrize(
"attn_type", [AttentionType.ENCODER_ONLY, AttentionType.ENCODER]
)
@pytest.mark.parametrize("dtype", _ZENTORCH_SDPA_DTYPES)
def test_should_use_zentorch_sdpa(attn_type: str, dtype: torch.dtype) -> None:
assert should_use_zentorch_sdpa(attn_type, dtype)
@_skip_no_zentorch_sdpa
@pytest.mark.parametrize(
"attn_type,dtype",
[
(AttentionType.DECODER, torch.bfloat16),
(AttentionType.ENCODER_ONLY, torch.float16),
],
)
def test_should_not_use_zentorch_sdpa(attn_type: str, dtype: torch.dtype) -> None:
assert not should_use_zentorch_sdpa(attn_type, dtype)
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
vllm/v1/attention/backends/cpu_attn.py
+26
View File
@@ -13,6 +13,7 @@ from vllm import _custom_ops as ops
from vllm import envs
from vllm.config import (
VllmConfig,
get_current_vllm_config,
get_layers_from_vllm_config,
)
from vllm.logger import init_logger
@@ -31,6 +32,10 @@ from vllm.v1.attention.backend import (
from vllm.v1.attention.backends.utils import (
get_num_attention_heads_from_layers,
)
from vllm.v1.attention.backends.zentorch_sdpa import (
should_use_zentorch_sdpa,
zentorch_sdpa_attn,
)
from vllm.v1.kv_cache_interface import (
AttentionSpec,
CrossAttentionSpec,
@@ -334,6 +339,12 @@ class CPUAttentionBackendImpl(AttentionImpl):
"heads in the layer"
)
vllm_config = get_current_vllm_config()
self.use_zentorch_sdpa = should_use_zentorch_sdpa(
attn_type,
vllm_config.model_config.dtype,
)
def forward(
self,
layer: AttentionLayer,
@@ -383,6 +394,21 @@ class CPUAttentionBackendImpl(AttentionImpl):
AttentionType.ENCODER,
)
if is_encoder_attention:
if self.use_zentorch_sdpa:
# Encoder attention never reads the KV cache back, so the
# zentorch path attends the packed QKV directly instead of
# staging it through the scratch encoder cache.
zentorch_sdpa_attn(
query[:num_actual_tokens],
key[:num_actual_tokens],
value[:num_actual_tokens],
output[:num_actual_tokens],
attn_metadata,
self.scale,
self.sliding_window,
self.alibi_slopes,
)
return output
# For encoder attention,
kv_cache = attn_metadata.encoder_cache
vllm/v1/attention/backends/zentorch_sdpa.py
+148
View File
@@ -0,0 +1,148 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Encoder attention through the zentorch SDPA kernel on Zen CPUs."""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
from vllm.v1.attention.backend import AttentionType
if TYPE_CHECKING:
from vllm.v1.attention.backends.cpu_attn import CPUAttentionMetadata
__all__ = ["should_use_zentorch_sdpa", "zentorch_sdpa_attn"]
def should_use_zentorch_sdpa(
attn_type: str,
dtype: torch.dtype,
) -> bool:
"""True when encoder attention should dispatch to zentorch_sdpa.
Full, sliding-window, and ALiBi encoder layers all dispatch. A window is
a single-head [1, 1, S, S] additive mask; ALiBi is a per-query-head
[1, H, S, S] bias. The zentorch_sdpa_attn op builds both.
dtype must match the op's ISA gate (bf16: AVX512-BF16, fp32: AVX512);
fp16 needs AVX512-FP16, which torch exposes no query for, so it stays
native.
"""
if attn_type not in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
return False
if not has_zentorch_op(["zentorch_sdpa_attn"]):
return False
if dtype == torch.bfloat16:
return torch.cpu._is_avx512_bf16_supported()
if dtype == torch.float32:
return torch.cpu._is_avx512_supported()
return False
def zentorch_sdpa_attn(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
output: torch.Tensor,
attn_metadata: CPUAttentionMetadata,
scale: float,
sliding_window: int = -1,
alibi_slopes: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run encoder / encoder-only attention with zentorch_sdpa.
Encoder attention reads no KV cache, so the packed
[num_tokens, num_heads, head_size] query/key/value are attended in place
and written straight into `output`.
Args:
query: Query of shape [num_tokens, num_heads, head_size].
key: Key of shape [num_tokens, num_kv_heads, head_size].
value: Value of shape [num_tokens, num_kv_heads, head_size].
output: Pre-allocated output, same shape as `query`.
attn_metadata: Metadata carrying the per-sequence token offsets and
whether attention is causal.
scale: Softmax scale.
sliding_window: Encoder window size, or -1 for full bidirectional
attention. Handed to the op, which masks token i to
[i-(W-1), i+(W-1)].
alibi_slopes: Per-query-head ALiBi slopes, or None. Handed to the op,
which applies them as an additive per-head bias.
Returns:
`output`, filled in place.
"""
# zentorch_sdpa_attn builds the ALiBi mask in query.dtype. A bf16 per-head
# mask is wrong when the query head count is not a power of two, so run
# that path in fp32 and write bf16 back.
if alibi_slopes is not None and query.dtype == torch.bfloat16:
out32 = torch.empty_like(query, dtype=torch.float32)
zentorch_sdpa_attn(
query.float(),
key.float(),
value.float(),
out32,
attn_metadata,
scale,
sliding_window,
alibi_slopes,
)
output.copy_(out32.to(dtype=output.dtype))
return output
# The op counts window tokens on each side of the diagonal; vLLM's window
# includes the diagonal itself.
window = -1 if sliding_window in (None, -1) else sliding_window - 1
# zentorch_sdpa takes no enable_gqa argument but derives the KV head count
# from the key tensor and maps each query head onto its KV head, so GQA/MQA
# needs no expansion here.
query = query.movedim(0, query.dim() - 2)
key = key.movedim(0, key.dim() - 2)
value = value.movedim(0, value.dim() - 2)
start_loc = attn_metadata.query_start_loc.numpy()
seq_lens = start_loc[1:] - start_loc[:-1]
# Encoder sequences are packed along the token dimension. When they all
# have the same length, recover a dense BHSD batch and invoke zentorch_sdpa
# once for the whole scheduler batch instead of once per sequence.
if len(seq_lens) > 0 and (seq_lens == seq_lens[0]).all():
batch_size = len(seq_lens)
seq_len = int(seq_lens[0])
def _packed_hsd_to_bhsd(tensor: torch.Tensor) -> torch.Tensor:
return tensor.unflatten(1, (batch_size, seq_len)).permute(1, 0, 2, 3)
torch.ops.zentorch.zentorch_sdpa_attn(
_packed_hsd_to_bhsd(query),
_packed_hsd_to_bhsd(key),
_packed_hsd_to_bhsd(value),
# The op writes [B, H, S, D]; hand it the matching view of the
# packed buffer so it stores there instead of into a temporary.
output.unflatten(0, (batch_size, seq_len)).permute(0, 2, 1, 3),
scale=scale,
is_causal=attn_metadata.causal,
alibi_slopes=alibi_slopes,
left_window_size=window,
right_window_size=window,
)
return output
for start_q, end_q in zip(start_loc[:-1], start_loc[1:], strict=True):
torch.ops.zentorch.zentorch_sdpa_attn(
query[None, :, start_q:end_q, :],
key[None, :, start_q:end_q, :],
value[None, :, start_q:end_q, :],
output[start_q:end_q, :, :].movedim(0, 1).unsqueeze(0),
scale=scale,
is_causal=attn_metadata.causal,
alibi_slopes=alibi_slopes,
left_window_size=window,
right_window_size=window,
)
return output
vllm/v1/attention/ops/vit_attn_wrappers.py
+11 -3
View File
@@ -241,9 +241,17 @@ def apply_sdpa(
(batch_size x seq_len x num_heads x head_size)
"""
q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
output = F.scaled_dot_product_attention(
q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
)
if current_platform.is_zen_cpu() and hasattr(torch.ops.zentorch, "zentorch_sdpa"):
# Schema: (query, key, value, dropout_p=0., is_causal=False, *,
# attn_mask=None, scale=None) -> (attention, logsumexp).
# No enable_gqa: GQA is inferred from key.size(1) vs query.size(1).
output, _ = torch.ops.zentorch.zentorch_sdpa(
q, k, v, dropout_p=0.0, is_causal=False, scale=scale
)
else:
output = F.scaled_dot_product_attention(
q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
)
output = einops.rearrange(output, "b h s d -> b s h d ")
return output