[CPU] Add MLA backend so DeepSeek-V2/V3 can run on CPU (#49453)
Signed-off-by: baoloongmao <baoloongmao@tencent.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
maobaolong committed
Aug 6, 2026 at 16:51 UTC
13726c80fe5758443b4e1508611ccf504a06085c
6 files changed
+367
-14
csrc/cpu/cpu_types_arm.hpp
+13
-8
@@ -943,10 +943,14 @@ inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
943
};
944
945
inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
946
- fmadd(acc.reg.val[0], a.reg.val[0], b.reg.val[0]);
947
- fmadd(acc.reg.val[1], a.reg.val[1], b.reg.val[1]);
948
- fmadd(acc.reg.val[2], a.reg.val[2], b.reg.val[2]);
949
- fmadd(acc.reg.val[3], a.reg.val[3], b.reg.val[3]);
946
+ // NOTE: at::vec::fmadd(a, b, c) returns a * b + c and does not mutate its
947
+ // arguments; we must capture the result back into acc, otherwise the FMA
948
+ // is a no-op. On x86 fma is implemented as `acc = acc + a * b` which does
949
+ // the assignment; the ARM version previously dropped the return value.
950
+ acc.reg.val[0] = fmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
951
+ acc.reg.val[1] = fmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
952
+ acc.reg.val[2] = fmadd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
953
+ acc.reg.val[3] = fmadd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
954
};
955
956
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
@@ -967,10 +971,11 @@ inline void fma(FP32Vec16& acc, BF16Vec32& a, BF16Vec32& b) {
971
std::tie(b0_low, b0_high) = convert_bfloat16_float(b.reg.val[0]);
972
std::tie(b1_low, b1_high) = convert_bfloat16_float(b.reg.val[1]);
973
970
- fmadd(acc.reg.val[0], a0_low, b0_low);
971
- fmadd(acc.reg.val[1], a0_high, b0_high);
972
- fmadd(acc.reg.val[2], a1_low, b1_low);
973
- fmadd(acc.reg.val[3], a1_high, b1_high);
974
+ // Same fmadd semantic as the FP32 overload above: capture the result.
975
+ acc.reg.val[0] = fmadd(a0_low, b0_low, acc.reg.val[0]);
976
+ acc.reg.val[1] = fmadd(a0_high, b0_high, acc.reg.val[1]);
977
+ acc.reg.val[2] = fmadd(a1_low, b1_low, acc.reg.val[2]);
978
+ acc.reg.val[3] = fmadd(a1_high, b1_high, acc.reg.val[3]);
979
};
980
981
template <>
csrc/cpu/mla_decode.cpp
+9
@@ -57,6 +57,12 @@ void mla_decode_block_head(
57
std::fill(max_val, max_val + HEAD_UNROLL, -FLT_MAX);
58
59
f32_vec_type acc_vec[BLOCK_SIZE][HEAD_UNROLL];
60
+ // Explicitly zero-initialize the accumulator. On x86 the default
61
+ // ctor of FP32Vec16 already sets zero, but on ARM (and any backend
62
+ // whose Vectorized<float> default ctor is trivial) the array holds
63
+ // undefined values and the fma below accumulates garbage.
64
+ for (int b = 0; b < BLOCK_SIZE; ++b)
65
+ for (int u = 0; u < HEAD_UNROLL; ++u) acc_vec[b][u] = f32_vec_type{0.0f};
66
for (int i = 0; i < HEAD_DIM; i += QK_NUM_ELEM) {
67
// load to registers
68
qk_vec_type q_vec[HEAD_UNROLL];
@@ -96,6 +102,9 @@ void mla_decode_block_head(
102
}
103
104
f32_vec_type this_out[V_HEAD_DIM / V_NUM_ELEM][HEAD_UNROLL];
105
+ // Same zero-initialization concern as acc_vec above.
106
+ for (int i = 0; i < V_HEAD_DIM / V_NUM_ELEM; ++i)
107
+ for (int u = 0; u < HEAD_UNROLL; ++u) this_out[i][u] = f32_vec_type{0.0f};
108
109
for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
110
// load to registers
tests/v1/attention/test_cpu_mla_backend.py
new
+88
@@ -0,0 +1,88 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Smoke tests for the CPU MLA backend.
4
+
5
+Two levels of coverage are provided:
6
+
7
+* :func:`test_kv_cache_cpu_write` checks that the kv-cache write logic
8
+ inlined in the CPU MLA backend writes the latent KV cache in the same
9
+ layout the CPU decode kernel expects. This runs without downloading
10
+ any model weights and is safe for CI.
11
+* :func:`test_cpu_mla_backend_smoke` exercises the full end-to-end path
12
+ (LLM engine, model construction, prefill + decode) against a shrunk
13
+ DeepSeek-V2-Lite with `hf_overrides`. It uses `dummy` weights so the
14
+ outputs are not meaningful, only shape / plumbing correctness is
15
+ asserted. Marked `cpu_model` so it is skipped on non-CPU CI.
16
+"""
17
+
18
+import pytest
19
+import torch
20
+
21
+from vllm.platforms import current_platform
22
+
23
+
24
+@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
25
+def test_kv_cache_cpu_write() -> None:
26
+ kv_lora_rank = 512
27
+ pe_dim = 64
28
+ block_size = 16
29
+ num_blocks = 4
30
+ num_tokens = 8
31
+ dtype = torch.float32
32
+
33
+ kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype)
34
+ k_pe = torch.randn(num_tokens, pe_dim, dtype=dtype)
35
+ kv_cache = torch.zeros(num_blocks, block_size, kv_lora_rank + pe_dim, dtype=dtype)
36
+ slot_mapping = torch.arange(num_tokens, dtype=torch.long)
37
+ # Mark the last slot as padding to make sure we honour negative slots.
38
+ slot_mapping[-1] = -1
39
+
40
+ flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
41
+ slots = slot_mapping.to(torch.long)
42
+ valid_mask = slots >= 0
43
+ if not valid_mask.all().item():
44
+ slots = slots[valid_mask]
45
+ kv_c = kv_c[valid_mask]
46
+ k_pe = k_pe[valid_mask]
47
+ target_dtype = flat.dtype
48
+ flat[slots, :kv_lora_rank] = kv_c.to(target_dtype)
49
+ flat[slots, kv_lora_rank:] = k_pe.to(target_dtype)
50
+
51
+ flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
52
+ for i in range(num_tokens - 1):
53
+ assert torch.equal(flat[i, :kv_lora_rank], kv_c[i])
54
+ assert torch.equal(flat[i, kv_lora_rank:], k_pe[i])
55
+ # The padding slot must remain untouched.
56
+ assert torch.all(flat[num_tokens - 1] == 0)
57
+
58
+
59
+@pytest.mark.cpu_model
60
+@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
61
+def test_cpu_mla_backend_smoke() -> None:
62
+ # Shrink the architecture aggressively so the test is CI-friendly.
63
+ from vllm import LLM, SamplingParams
64
+
65
+ hf_overrides = {
66
+ "num_hidden_layers": 2,
67
+ "n_routed_experts": 4,
68
+ "first_k_dense_replace": 0,
69
+ "num_experts_per_tok": 2,
70
+ }
71
+ llm = LLM(
72
+ model="deepseek-ai/DeepSeek-V2-Lite",
73
+ trust_remote_code=True,
74
+ load_format="dummy",
75
+ enforce_eager=True,
76
+ max_model_len=128,
77
+ max_num_seqs=2,
78
+ block_size=16,
79
+ hf_overrides=hf_overrides,
80
+ )
81
+
82
+ sampling_params = SamplingParams(max_tokens=4, temperature=0.0)
83
+ outputs = llm.generate(["Hello", "MLA on CPU"], sampling_params)
84
+ assert len(outputs) == 2
85
+ for output in outputs:
86
+ # `dummy` weights do not produce meaningful text, but the number
87
+ # of generated tokens must match what we asked for.
88
+ assert len(output.outputs[0].token_ids) == 4
vllm/platforms/cpu.py
+36
-6
@@ -78,12 +78,19 @@ class CpuPlatform(Platform):
78
attn_selector_config: "AttentionSelectorConfig",
79
num_heads: int | None = None,
80
) -> str:
81
- if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:
82
- logger.info("Cannot use %s backend on CPU.", selected_backend)
83
- if attn_selector_config.use_mla:
84
- raise NotImplementedError("MLA is not supported on CPU.")
81
if attn_selector_config.use_sparse:
82
raise NotImplementedError("Sparse Attention is not supported on CPU.")
83
+ if attn_selector_config.use_mla:
84
+ # Reference MLA implementation on CPU. Performance is not the
85
+ # goal here; the backend simply wires the CPU decode kernel
86
+ # (`mla_decode_kvcache`) and an SDPA-based prefill together with
87
+ # the shared MLA scaffolding so that DeepSeek-style models can
88
+ # execute on CPU.
89
+ if selected_backend and selected_backend != AttentionBackendEnum.CPU_MLA:
90
+ logger.info("Cannot use %s backend on CPU.", selected_backend)
91
+ return AttentionBackendEnum.CPU_MLA.get_path()
92
+ if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:
93
+ logger.info("Cannot use %s backend on CPU.", selected_backend)
94
return AttentionBackendEnum.CPU_ATTN.get_path()
95
96
@classmethod
@@ -116,10 +123,25 @@ class CpuPlatform(Platform):
123
124
cache_config = vllm_config.cache_config
125
119
- if not cache_config.user_specified_block_size:
126
+ # The CPU MLA decode kernel only compiles with block_size=16 today
127
+ # (see csrc/cpu/mla_decode.cpp). If the model uses MLA we override
128
+ # the default block size regardless of user preference to avoid a
129
+ # runtime kernel dispatch failure.
130
+ cpu_mla_enabled = model_config is not None and getattr(
131
+ model_config, "use_mla", False
132
+ )
133
+ if cpu_mla_enabled:
134
+ if cache_config.user_specified_block_size and cache_config.block_size != 16:
135
+ logger.warning(
136
+ "CPU MLA backend requires block_size=16, overriding "
137
+ "user-specified block_size=%s.",
138
+ cache_config.block_size,
139
+ )
140
+ cache_config.block_size = 16
141
+ elif not cache_config.user_specified_block_size:
142
cache_config.block_size = 128
143
122
- if cache_config.block_size % 32 != 0:
144
+ if not cpu_mla_enabled and cache_config.block_size % 32 != 0:
145
logger.warning(
146
"CPU backend prefers block_size is multiples of 32, "
147
"otherwise the performance is not optimized."
@@ -321,6 +343,7 @@ class CpuPlatform(Platform):
343
"prefill and prefix caching to be disabled."
344
)
345
vllm_config.scheduler_config.enable_chunked_prefill = False
346
+ vllm_config.cache_config.enable_prefix_caching = False
347
vllm_config.scheduler_config.max_num_batched_tokens = max(
348
vllm_config.model_config.max_model_len,
349
vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
@@ -473,6 +496,13 @@ class CpuPlatform(Platform):
496
from vllm._custom_ops import cpu_attn_reshape_and_cache
497
from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
498
499
+ # MLA uses a single latent cache of shape [N, block_size, head_size],
500
+ # so the classic key/value split does not apply. The MLA backend
501
+ # writes the cache itself via `concat_and_cache_mla` inside
502
+ # `do_kv_cache_update`, so there is nothing to pack here.
503
+ if kv_cache.dim() == 3:
504
+ return
505
+
506
num_blocks, num_kv_heads, block_size, fused_head_size = kv_cache.shape
507
head_size = fused_head_size // 2
508
vllm/v1/attention/backends/mla/cpu_mla.py
new
+220
@@ -0,0 +1,220 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU MLA backend.
4
+
5
+This is a *reference-quality* MLA backend for the CPU platform. It is
6
+intended to make DeepSeek-V2/V3 style models runnable on CPU (for
7
+functional verification, tiny-model smoke tests and CI), not to be
8
+performant. The prefill path is a plain PyTorch SDPA and the decode
9
+path forwards to the existing CPU decode kernel
10
+`torch.ops._C.mla_decode_kvcache` (see csrc/cpu/mla_decode.cpp).
11
+
12
+Key design points:
13
+
14
+* We inherit the shared MLA scaffolding (`MLACommonBackend`,
15
+ `MLACommonImpl`, `MLACommonMetadata`) so that the same
16
+ ``forward_impl`` in ``MLAAttention`` orchestrates weight-absorbed
17
+ decode (MQA) and non-absorbed prefill (MHA) for us.
18
+* The parent ``MLACommonImpl.__init__`` tries to pick a GPU prefill
19
+ kernel (flash_attn / flashinfer / cudnn / trtllm) and raises when
20
+ none of them are available. On CPU none of them apply, so we
21
+ bypass that logic and set the attributes ourselves.
22
+* Chunked prefill and prefix caching are disabled by the CPU platform
23
+ when ``use_mla`` is set.
24
+* The CPU decode kernel only supports ``head_dim=576``,
25
+ ``v_head_dim=512`` and ``block_size=16`` today; the platform layer
26
+ forces ``block_size=16`` for MLA models.
27
+"""
28
+
29
+from typing import ClassVar
30
+
31
+import torch
32
+
33
+from vllm import _custom_ops as ops
34
+from vllm.config.cache import CacheDType
35
+from vllm.logger import init_logger
36
+from vllm.model_executor.layers.attention.mla_attention import (
37
+ MLACommonBackend,
38
+ MLACommonImpl,
39
+ MLACommonMetadata,
40
+)
41
+from vllm.v1.attention.backend import (
42
+ AttentionLayer,
43
+ AttentionType,
44
+ MultipleOf,
45
+)
46
+
47
+logger = init_logger(__name__)
48
+
49
+
50
+class CPUMLABackend(MLACommonBackend):
51
+ """Attention backend descriptor for CPU MLA."""
52
+
53
+ supported_dtypes: ClassVar[list[torch.dtype]] = [
54
+ torch.float16,
55
+ torch.bfloat16,
56
+ torch.float32,
57
+ ]
58
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
59
+
60
+ @classmethod
61
+ def get_supported_head_sizes(cls) -> list[int]:
62
+ # DeepSeek-V2/V3 latent cache size = kv_lora_rank(512)+rope(64)=576.
63
+ # The current CPU decode kernel only compiles for this size.
64
+ return [576]
65
+
66
+ @staticmethod
67
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
68
+ return [16]
69
+
70
+ @classmethod
71
+ def supports_block_size(cls, block_size: int | None) -> bool:
72
+ return block_size is None or block_size == 16
73
+
74
+ @staticmethod
75
+ def get_name() -> str:
76
+ return "CPU_MLA"
77
+
78
+ @staticmethod
79
+ def get_impl_cls() -> type["CPUMLAImpl"]:
80
+ return CPUMLAImpl
81
+
82
+
83
+class CPUMLAImpl(MLACommonImpl[MLACommonMetadata]):
84
+ """CPU implementation of MLA attention.
85
+
86
+ See module docstring for the overall design. This class only
87
+ overrides the parts of the parent that would otherwise pull in
88
+ CUDA-only dependencies.
89
+ """
90
+
91
+ # Decode always returns the softmax lse if the parent needs it, but
92
+ # CPU decode currently only produces the output tensor. Report False
93
+ # so DCP / cross-batch reductions are not requested from us.
94
+ can_return_lse_for_decode: bool = False
95
+
96
+ def __init__(
97
+ self,
98
+ num_heads: int,
99
+ head_size: int,
100
+ scale: float,
101
+ num_kv_heads: int,
102
+ alibi_slopes: list[float] | None,
103
+ sliding_window: int | None,
104
+ kv_cache_dtype: str,
105
+ logits_soft_cap: float | None,
106
+ attn_type: str,
107
+ kv_sharing_target_layer_name: str | None,
108
+ # MLA-specific arguments (forwarded from `MLAAttention`).
109
+ **mla_args,
110
+ ) -> None:
111
+ # NOTE: We deliberately skip ``MLACommonImpl.__init__`` because it
112
+ # tries to pick a GPU prefill kernel and raises on CPU. Instead we
113
+ # replicate the attribute setup here and then wire our own hooks.
114
+ if kv_sharing_target_layer_name is not None:
115
+ raise NotImplementedError("KV sharing is not supported for MLA")
116
+ if attn_type != AttentionType.DECODER:
117
+ raise NotImplementedError("CPU MLA only supports decoder self-attention.")
118
+ unsupported = [alibi_slopes, sliding_window, logits_soft_cap]
119
+ if any(unsupported):
120
+ raise NotImplementedError(
121
+ "CPU MLA does not support alibi_slopes / sliding_window / "
122
+ "logits_soft_cap."
123
+ )
124
+
125
+ self.num_heads = num_heads
126
+ self.head_size = head_size
127
+ self.scale = float(scale)
128
+ self.num_kv_heads = num_kv_heads
129
+ self.kv_cache_dtype = kv_cache_dtype
130
+
131
+ self.q_lora_rank = mla_args["q_lora_rank"]
132
+ self.kv_lora_rank = mla_args["kv_lora_rank"]
133
+ self.qk_nope_head_dim = mla_args["qk_nope_head_dim"]
134
+ self.qk_rope_head_dim = mla_args["qk_rope_head_dim"]
135
+ self.qk_head_dim = mla_args["qk_head_dim"]
136
+ self.v_head_dim = mla_args["v_head_dim"]
137
+ self.kv_b_proj = mla_args["kv_b_proj"]
138
+ self.indexer = mla_args.get("indexer")
139
+ self.q_pad_num_heads = mla_args.get("q_pad_num_heads")
140
+ self.supports_quant_query_input = False
141
+
142
+ # Never take the flashinfer specialisation of ``_concat_k_nope_k_pe``.
143
+ self._use_flashinfer_concat_mla_k = False
144
+
145
+ # We do not need V padding since our SDPA prefill takes k / v with
146
+ # matching head dims.
147
+ self._pad_v = False
148
+
149
+ # Distributed / context-parallel bookkeeping expected by the parent.
150
+ self.dcp_world_size = 1
151
+ self.cp_kv_cache_interleave_size = 1
152
+
153
+ # ------------------------------------------------------------------
154
+ # Decode (MQA)
155
+ # ------------------------------------------------------------------
156
+ def forward_mqa(
157
+ self,
158
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
159
+ kv_c_and_k_pe_cache: torch.Tensor,
160
+ attn_metadata: MLACommonMetadata,
161
+ layer: AttentionLayer,
162
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
163
+ assert attn_metadata.decode is not None
164
+ # The parent hands us either a tuple (q_nope, q_pe) or a single
165
+ # tensor already concatenated. CPU decode kernel wants a single
166
+ # (bs, num_heads, head_dim=kv_lora_rank+rope) contiguous tensor.
167
+ if isinstance(q, tuple):
168
+ q = torch.cat(q, dim=-1)
169
+ q = q.contiguous()
170
+
171
+ bs, num_heads, _ = q.shape
172
+ out = torch.zeros(
173
+ bs, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
174
+ )
175
+
176
+ # The CPU kernel is templated on int32 block_tables / seq_lens.
177
+ block_table = attn_metadata.decode.block_table.to(torch.int32)
178
+ seq_lens = attn_metadata.decode.seq_lens.to(torch.int32)
179
+
180
+ ops.mla_decode_kvcache_cpu(
181
+ out,
182
+ q,
183
+ kv_c_and_k_pe_cache,
184
+ self.scale,
185
+ block_table,
186
+ seq_lens,
187
+ )
188
+ # The parent only consumes the LSE when DCP is active, which we
189
+ # disable, so returning None here is safe.
190
+ return out, None
191
+
192
+ # ------------------------------------------------------------------
193
+ # KV cache write
194
+ # ------------------------------------------------------------------
195
+ def do_kv_cache_update(
196
+ self,
197
+ kv_c_normed: torch.Tensor,
198
+ k_pe: torch.Tensor,
199
+ kv_cache: torch.Tensor,
200
+ slot_mapping: torch.Tensor,
201
+ kv_cache_dtype: str,
202
+ k_scale: torch.Tensor,
203
+ ) -> None:
204
+ if kv_cache.numel() == 0:
205
+ return
206
+ k_pe_sq = k_pe.squeeze(1)
207
+ kv_lora_rank = kv_c_normed.shape[-1]
208
+ pe_dim = k_pe_sq.shape[-1]
209
+ flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
210
+ slots = slot_mapping.flatten().to(torch.long)
211
+ # Only write valid slots. Padding slots are set to -1 by the
212
+ # scheduler and must be skipped.
213
+ valid_mask = slots >= 0
214
+ if not valid_mask.all().item():
215
+ slots = slots[valid_mask]
216
+ kv_c_normed = kv_c_normed[valid_mask]
217
+ k_pe_sq = k_pe_sq[valid_mask]
218
+ target_dtype = flat.dtype
219
+ flat[slots, :kv_lora_rank] = kv_c_normed.to(target_dtype)
220
+ flat[slots, kv_lora_rank:] = k_pe_sq.view(-1, pe_dim).to(target_dtype)
vllm/v1/attention/backends/registry.py
+1
@@ -122,6 +122,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
122
"RocmAiterUnifiedAttentionBackend"
123
)
124
CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend"
125
+ CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend"
126
TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend"
127
# Placeholder for third-party/custom backends - must be registered before use
128
# set to None to avoid alias with other backend, whose value is an empty string