[Attention][CPU] Use zentorch SDPA for CPU MLA prefill (#54967)
Signed-off-by: Rakul Chauhan <rakul.chauhan@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Rakul-Chauhan committed
Sep 27, 2026 at 14:00 UTC
24c9772d19251dbbf70fef75119546827c540c63
5 files changed
+222
-3
tests/v1/attention/test_cpu_mla_backend.py
+67
-1
@@ -25,6 +25,9 @@ from vllm.v1.attention.backends.mla.prefill.cpu_sdpa import (
25
from vllm.v1.attention.backends.mla.prefill.selector import (
26
get_mla_prefill_backend,
27
)
28
+from vllm.v1.attention.backends.mla.prefill.zen_cpu_sdpa import (
29
+ ZenCPUSDPAMLAPrefillBackend,
30
+)
31
32
33
def _reference_ragged_attention(
@@ -157,8 +160,10 @@ def test_cpu_mla_backend_smoke(tmp_path) -> None:
160
def test_cpu_mla_prefill_backend_selected() -> None:
161
# On CPU (no device capability) the MLA prefill backend must be the
162
# SDPA-based CPU backend, not flash-attn which is unavailable on CPU.
163
+ # Accelerator-specific subclasses (e.g. the zentorch-backed Zen CPU
164
+ # backend) are also valid selections.
165
backend_cls = get_mla_prefill_backend(None)
161
- assert backend_cls is CPUSDPAMLAPrefillBackend
166
+ assert issubclass(backend_cls, CPUSDPAMLAPrefillBackend)
167
168
169
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
@@ -248,3 +253,64 @@ def test_cpu_mla_prefill_context_chunk() -> None:
253
assert lse.float().allclose(ref_lse.float(), atol=2e-2)
254
assert out.shape == (5, 4, 16)
255
assert lse.shape == (4, 5)
256
+
257
+
258
+@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
259
+@pytest.mark.skipif(
260
+ not ZenCPUSDPAMLAPrefillBackend.is_available(),
261
+ reason="requires a Zen CPU with the zentorch_sdpa op registered",
262
+)
263
+def test_zen_cpu_mla_prefill_new_tokens() -> None:
264
+ # Runs the zentorch kernel itself. The head dims are deliberately
265
+ # asymmetric (v_head_dim < qk_head_dim, as in real MLA) so that the V
266
+ # padding and output slicing are exercised; the generic-CPU cases above
267
+ # use equal dims and never reach that code.
268
+ qk_head_dim = 24
269
+ v_head_dim = 16
270
+ backend = ZenCPUSDPAMLAPrefillBackend(
271
+ num_heads=4,
272
+ scale=0.25,
273
+ kv_lora_rank=16,
274
+ qk_nope_head_dim=16,
275
+ qk_rope_head_dim=8,
276
+ v_head_dim=v_head_dim,
277
+ vllm_config=None,
278
+ )
279
+
280
+ class _PrefillMeta:
281
+ query_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32)
282
+
283
+ backend.prepare_metadata(_PrefillMeta())
284
+ q = torch.randn(6, 4, qk_head_dim, dtype=torch.bfloat16)
285
+ k = torch.randn(6, 4, qk_head_dim, dtype=torch.bfloat16)
286
+ v = torch.randn(6, 4, v_head_dim, dtype=torch.bfloat16)
287
+
288
+ out = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=False)
289
+ ref, _ = _reference_ragged_attention(
290
+ q,
291
+ k,
292
+ v,
293
+ _PrefillMeta.query_start_loc,
294
+ _PrefillMeta.query_start_loc,
295
+ scale=0.25,
296
+ causal=True,
297
+ )
298
+ assert out.shape == (6, 4, v_head_dim)
299
+ assert out.float().allclose(ref.float(), atol=2e-2)
300
+
301
+ # zentorch_sdpa does not return a usable LSE, so that path must fall back
302
+ # to the fp32 reference and match the generic backend exactly.
303
+ generic = CPUSDPAMLAPrefillBackend(
304
+ num_heads=4,
305
+ scale=0.25,
306
+ kv_lora_rank=16,
307
+ qk_nope_head_dim=16,
308
+ qk_rope_head_dim=8,
309
+ v_head_dim=v_head_dim,
310
+ vllm_config=None,
311
+ )
312
+ generic.prepare_metadata(_PrefillMeta())
313
+ zen_out, zen_lse = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
314
+ gen_out, gen_lse = generic.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
315
+ assert torch.equal(zen_out, gen_out)
316
+ assert torch.equal(zen_lse, gen_lse)
tests/v1/attention/test_mla_prefill_selector.py
+33
@@ -60,10 +60,43 @@ class TestGetMLAPrefillBackend:
60
61
def test_cpu_uses_sdpa_prefill(self):
62
vllm_config = _make_vllm_config()
63
+ zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class()
64
+
65
+ # A Zen CPU with zentorch selects the Zen backend, so expect whichever
66
+ # one this host offers.
67
+ expected = (
68
+ zen_cls if zen_cls.is_available() else MLAPrefillBackendEnum.CPU.get_class()
69
+ )
70
71
with patch("vllm.platforms.current_platform") as mock_platform:
72
mock_platform.is_cpu.return_value = True
73
74
+ backend = get_mla_prefill_backend(vllm_config)
75
+ assert backend is expected
76
+
77
+ def test_zen_cpu_prefers_zentorch_prefill(self):
78
+ vllm_config = _make_vllm_config()
79
+ zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class()
80
+
81
+ with (
82
+ patch("vllm.platforms.current_platform") as mock_platform,
83
+ patch.object(zen_cls, "is_available", return_value=True),
84
+ ):
85
+ mock_platform.is_cpu.return_value = True
86
+
87
+ backend = get_mla_prefill_backend(vllm_config)
88
+ assert backend is zen_cls
89
+
90
+ def test_zen_cpu_falls_back_when_zentorch_unavailable(self):
91
+ vllm_config = _make_vllm_config()
92
+ zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class()
93
+
94
+ with (
95
+ patch("vllm.platforms.current_platform") as mock_platform,
96
+ patch.object(zen_cls, "is_available", return_value=False),
97
+ ):
98
+ mock_platform.is_cpu.return_value = True
99
+
100
backend = get_mla_prefill_backend(vllm_config)
101
assert backend is MLAPrefillBackendEnum.CPU.get_class()
102
vllm/v1/attention/backends/mla/prefill/registry.py
+4
@@ -53,6 +53,10 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
53
"AiterFlashAttnPrefillBackend"
54
)
55
CPU = "vllm.v1.attention.backends.mla.prefill.cpu_sdpa.CPUSDPAMLAPrefillBackend"
56
+ ZEN_CPU = (
57
+ "vllm.v1.attention.backends.mla.prefill.zen_cpu_sdpa."
58
+ "ZenCPUSDPAMLAPrefillBackend"
59
+ )
60
# Placeholder for third-party/custom backends - must be registered before use
61
# set to None to avoid alias with other backend, whose value is an empty string
62
CUSTOM = None
vllm/v1/attention/backends/mla/prefill/selector.py
+12
-2
@@ -110,8 +110,18 @@ def get_mla_prefill_backend(
110
from vllm.platforms import current_platform
111
112
if current_platform.is_cpu():
113
- logger.info_once("Using CPU SDPA MLA prefill backend.")
114
- return MLAPrefillBackendEnum.CPU.get_class()
113
+ # CPUs have no compute capability, so the capability-driven priority
114
+ # path below does not apply. Prefer an accelerator-specific CPU backend
115
+ # when its kernels are present, else the generic SDPA one.
116
+ for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU):
117
+ try:
118
+ cpu_backend_cls = backend_enum.get_class()
119
+ except ImportError:
120
+ continue
121
+ if cpu_backend_cls.is_available():
122
+ logger.info_once("Using %s MLA prefill backend.", backend_enum.name)
123
+ return cpu_backend_cls
124
+ raise ValueError("No valid CPU MLA prefill backend found.")
125
126
device_capability = current_platform.get_device_capability()
127
if device_capability is None:
vllm/v1/attention/backends/mla/prefill/zen_cpu_sdpa.py
new
+106
@@ -0,0 +1,106 @@
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)