From 5963795ec7d7beaf34a05a213cce113c6508a8d4 Mon Sep 17 00:00:00 2001 From: Priyansh Jain <167848587+Priyjain-amd@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:35:44 +0530 Subject: [PATCH] [Attention][CPU] Run Zen CPU encoder attention on zentorch SDPA (#54508) Signed-off-by: priyansh jain Co-authored-by: Cursor Co-authored-by: Andreas Karatzas --- setup.py | 2 +- tests/kernels/attention/test_cpu_attn.py | 133 ++++++++++++++++++ vllm/v1/attention/backends/cpu_attn.py | 26 ++++ vllm/v1/attention/backends/zentorch_sdpa.py | 148 ++++++++++++++++++++ vllm/v1/attention/ops/vit_attn_wrappers.py | 14 +- 5 files changed, 319 insertions(+), 4 deletions(-) create mode 100644 vllm/v1/attention/backends/zentorch_sdpa.py diff --git a/setup.py b/setup.py index ed9616195f..67f192cdd3 100644 --- a/setup.py +++ b/setup.py @@ -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"], diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index 3046e2d2ac..f79b1b15d2 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -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) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index b4f25f4194..46b01cc551 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -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 diff --git a/vllm/v1/attention/backends/zentorch_sdpa.py b/vllm/v1/attention/backends/zentorch_sdpa.py new file mode 100644 index 0000000000..cafb524b0e --- /dev/null +++ b/vllm/v1/attention/backends/zentorch_sdpa.py @@ -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 diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index b90eb58e9b..dccc670ce5 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -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