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

[CPU][MLA] Fix prefill backend selection so MLA runs end-to-end on CPU (#51471)

Browse Source 
Signed-off-by: baoloongmao <baoloongmao@tencent.com>
Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
maobaolong
2026-08-28 12:36:55 +08:00
committed by GitHub
co-authored by Codex
parent e6bfe03ad7
commit 6f7df92a8e
9 changed files with 442 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_cpu_mla_backend.py
+164 -2
View File
@@ -19,6 +19,51 @@ import pytest
import torch
from vllm.platforms import current_platform
from vllm.v1.attention.backends.mla.prefill.cpu_sdpa import (
CPUSDPAMLAPrefillBackend,
)
from vllm.v1.attention.backends.mla.prefill.selector import (
get_mla_prefill_backend,
)
def _reference_ragged_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q_cu_seq_lens: torch.Tensor,
kv_cu_seq_lens: torch.Tensor,
*,
scale: float,
causal: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
outputs: list[torch.Tensor] = []
lses: list[torch.Tensor] = []
num_reqs = q_cu_seq_lens.numel() - 1
for req_idx in range(num_reqs):
q_start = int(q_cu_seq_lens[req_idx])
q_end = int(q_cu_seq_lens[req_idx + 1])
kv_start = int(kv_cu_seq_lens[req_idx])
kv_end = int(kv_cu_seq_lens[req_idx + 1])
q_r = q[q_start:q_end].transpose(0, 1).float()
k_r = k[kv_start:kv_end].transpose(0, 1).float()
v_r = v[kv_start:kv_end].transpose(0, 1).float()
scores = torch.matmul(q_r, k_r.transpose(-2, -1)) * scale
if causal:
causal_mask = torch.ones(
(q_end - q_start, kv_end - kv_start),
dtype=torch.bool,
device=scores.device,
).triu(1)
scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf"))
probs = torch.softmax(scores, dim=-1)
outputs.append(torch.matmul(probs, v_r).transpose(0, 1).to(v.dtype))
lses.append(torch.logsumexp(scores, dim=-1))
return torch.cat(outputs, dim=0), torch.cat(lses, dim=1)
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
@@ -58,7 +103,7 @@ def test_kv_cache_cpu_write() -> None:
@pytest.mark.cpu_model
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
def test_cpu_mla_backend_smoke() -> None:
def test_cpu_mla_backend_smoke(tmp_path) -> None:
# Shrink the architecture aggressively so the test is CI-friendly.
from vllm import LLM, SamplingParams
@@ -70,19 +115,136 @@ def test_cpu_mla_backend_smoke() -> None:
}
llm = LLM(
model="deepseek-ai/DeepSeek-V2-Lite",
skip_tokenizer_init=True,
trust_remote_code=True,
load_format="dummy",
enforce_eager=True,
max_model_len=128,
max_num_seqs=2,
block_size=16,
gpu_memory_utilization=0.25,
hf_overrides=hf_overrides,
kv_transfer_config={
"kv_connector": "ExampleConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"shared_storage_path": str(tmp_path),
},
},
)
sampling_params = SamplingParams(max_tokens=4, temperature=0.0)
outputs = llm.generate(["Hello", "MLA on CPU"], sampling_params)
outputs = llm.generate([[1, 2, 3, 4], [5, 6, 7]], sampling_params, use_tqdm=False)
assert len(outputs) == 2
for output in outputs:
# `dummy` weights do not produce meaningful text, but the number
# of generated tokens must match what we asked for.
assert len(output.outputs[0].token_ids) == 4
prefix_hit_prompt = [1] * 20
outputs = llm.generate([prefix_hit_prompt], sampling_params, use_tqdm=False)
assert len(outputs[0].outputs[0].token_ids) == 4
assert llm.reset_prefix_cache(reset_connector=False)
# ExampleConnector persists KV externally at block granularity. With a
# 20-token prompt and block_size=16, the second run reloads 16 cached
# tokens and forces the 4-token suffix through MLA prefill with context.
outputs = llm.generate([prefix_hit_prompt], sampling_params, use_tqdm=False)
assert len(outputs[0].outputs[0].token_ids) == 4
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
def test_cpu_mla_prefill_backend_selected() -> None:
# On CPU (no device capability) the MLA prefill backend must be the
# SDPA-based CPU backend, not flash-attn which is unavailable on CPU.
backend_cls = get_mla_prefill_backend(None)
assert backend_cls is CPUSDPAMLAPrefillBackend
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
def test_cpu_mla_prefill_new_tokens() -> None:
# The SDPA prefill must exactly match a reference causal attention over
# the ragged (varlen) q/k/v layout, with per-request padding handled.
backend = CPUSDPAMLAPrefillBackend(
num_heads=4,
scale=0.25,
kv_lora_rank=16,
qk_nope_head_dim=16,
qk_rope_head_dim=0,
v_head_dim=16,
vllm_config=None,
)
class _PrefillMeta:
query_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32)
backend.prepare_metadata(_PrefillMeta())
q = torch.randn(6, 4, 16, dtype=torch.bfloat16)
k = torch.randn(6, 4, 16, dtype=torch.bfloat16)
v = torch.randn(6, 4, 16, dtype=torch.bfloat16)
out = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=False)
ref, _ = _reference_ragged_attention(
q,
k,
v,
_PrefillMeta.query_start_loc,
_PrefillMeta.query_start_loc,
scale=0.25,
causal=True,
)
assert out.float().allclose(ref.float(), atol=2e-2)
assert out.shape == (6, 4, 16)
out_lse, lse = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
ref_out, ref_lse = _reference_ragged_attention(
q,
k,
v,
_PrefillMeta.query_start_loc,
_PrefillMeta.query_start_loc,
scale=0.25,
causal=True,
)
assert out_lse.float().allclose(ref_out.float(), atol=2e-2)
assert lse.float().allclose(ref_lse.float(), atol=2e-2)
assert lse.shape == (4, 6)
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
def test_cpu_mla_prefill_context_chunk() -> None:
backend = CPUSDPAMLAPrefillBackend(
num_heads=4,
scale=0.25,
kv_lora_rank=16,
qk_nope_head_dim=16,
qk_rope_head_dim=0,
v_head_dim=16,
vllm_config=None,
)
class _PrefillMeta:
query_start_loc = torch.tensor([0, 5], dtype=torch.int32)
class _Chunk:
query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32)
cu_seq_lens = torch.tensor([0, 3, 7], dtype=torch.int32)
backend.prepare_metadata(_PrefillMeta())
q = torch.randn(5, 4, 16, dtype=torch.bfloat16)
k = torch.randn(7, 4, 16, dtype=torch.bfloat16)
v = torch.randn(7, 4, 16, dtype=torch.bfloat16)
out, lse = backend.run_prefill_context_chunk(_Chunk(), q, k, v)
ref_out, ref_lse = _reference_ragged_attention(
q,
k,
v,
_Chunk.query_start_loc,
_Chunk.cu_seq_lens,
scale=0.25,
causal=False,
)
assert out.float().allclose(ref_out.float(), atol=2e-2)
assert lse.float().allclose(ref_lse.float(), atol=2e-2)
assert out.shape == (5, 4, 16)
assert lse.shape == (4, 5)
tests/v1/attention/test_mla_prefill_selector.py
+26 -1
View File
@@ -58,10 +58,31 @@ def _make_vllm_config(
class TestGetMLAPrefillBackend:
"""Tests for get_mla_prefill_backend (public API)."""
def test_no_device_capability_returns_flash_attn(self):
def test_cpu_uses_sdpa_prefill(self):
vllm_config = _make_vllm_config()
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.is_cpu.return_value = True
backend = get_mla_prefill_backend(vllm_config)
assert backend is MLAPrefillBackendEnum.CPU.get_class()
def test_no_device_capability_returns_flash_attn(self):
vllm_config = _make_vllm_config()
class FlashAttnBackend:
@staticmethod
def get_name():
return "FLASH_ATTN"
with (
patch("vllm.platforms.current_platform") as mock_platform,
patch.object(
MLAPrefillBackendEnum.FLASH_ATTN,
"get_class",
return_value=FlashAttnBackend,
),
):
mock_platform.get_device_capability.return_value = None
mock_platform.is_cpu.return_value = False
@@ -80,6 +101,7 @@ class TestGetMLAPrefillBackend:
)
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.is_cpu.return_value = False
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
@@ -98,6 +120,7 @@ class TestGetMLAPrefillBackend:
)
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.is_cpu.return_value = False
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
@@ -111,6 +134,7 @@ class TestGetMLAPrefillBackend:
)
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.is_cpu.return_value = False
mock_platform.get_device_capability.return_value = DeviceCapability(
major=10, minor=0
)
@@ -153,6 +177,7 @@ class TestGetMLAPrefillBackend:
return_value=3,
),
):
mock_platform.is_cpu.return_value = False
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
vllm/_custom_ops.py
+25
View File
@@ -174,6 +174,31 @@ def mla_decode_kvcache_cpu(
torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens)
def gather_mla_context_cache_cpu(
*,
src_cache: torch.Tensor,
dst: torch.Tensor,
block_table: torch.Tensor,
starts: torch.Tensor,
cu_seq_lens: torch.Tensor,
) -> None:
page_size = src_cache.shape[1]
flat_cache = src_cache.view(-1, src_cache.shape[-1])
seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]
out_start = 0
for req_idx, (start, seq_len) in enumerate(zip(starts.tolist(), seq_lens.tolist())):
if seq_len <= 0:
continue
positions = torch.arange(start, start + seq_len, device=block_table.device)
block_ids = block_table[req_idx, positions // page_size].to(torch.long)
offsets = positions % page_size
slots = block_ids * page_size + offsets
out_end = out_start + seq_len
dst[out_start:out_end].copy_(flat_cache[slots])
out_start = out_end
# merge attn states ops
def merge_attn_states(
output: torch.Tensor,
vllm/model_executor/layers/attention/mla_attention.py
+12
View File
@@ -2785,6 +2785,18 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
batch_size=chunk.num_requests,
seq_starts=chunk.starts,
)
elif current_platform.is_cpu():
assert not is_quantized_kv_cache(self.kv_cache_dtype), (
"CPU MLA context gather fallback only supports "
f"non-quantized KV cache, got {self.kv_cache_dtype}"
)
ops.gather_mla_context_cache_cpu(
src_cache=kv_c_and_k_pe_cache,
dst=workspace[:toks],
block_table=block_table,
starts=chunk.starts,
cu_seq_lens=chunk.cu_seq_lens,
)
elif not use_fp8_prefill:
ops.gather_and_maybe_dequant_cache(
src_cache=kv_c_and_k_pe_cache,
vllm/v1/attention/backends/mla/prefill/cpu_native.py
-61
View File
@@ -1,61 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Inert MLA prefill backend for the CPU platform.
CPU's MLA attention impl (``CPUMLAImpl``) fully overrides ``forward_mha`` with a
kernel that attends directly against the paged latent KV cache (covering both
fresh prefill and cached-prefix continuation in one pass), so it never calls
into a pluggable ``MLAPrefillBackend``'s ``run_prefill_new_tokens``/
``run_prefill_context_chunk``. This class exists only to satisfy
``MLACommonMetadataBuilder``'s structural requirement that
``MLAAttention.prefill_backend`` be a real, cloneable object.
"""
from typing import TYPE_CHECKING
import torch
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
if TYPE_CHECKING:
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonPrefillMetadata,
)
class CPUNativeMLAPrefillBackend(MLAPrefillBackend):
"""Placeholder MLA prefill backend for CPU; never actually invoked."""
@staticmethod
def get_name() -> str:
return "CPU_NATIVE"
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
out: torch.Tensor | None = None,
output_scale: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
raise AssertionError(
"CPUNativeMLAPrefillBackend.run_prefill_new_tokens is unreachable: "
"CPUMLAImpl.forward_mha fully overrides the dense-MHA prefill path "
"and never calls the pluggable prefill backend."
)
def run_prefill_context_chunk(
self,
chunk: "MLACommonPrefillMetadata.ContextChunk",
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
raise AssertionError(
"CPUNativeMLAPrefillBackend.run_prefill_context_chunk is "
"unreachable: CPUMLAImpl.forward_mha fully overrides the "
"dense-MHA prefill path and never calls the pluggable prefill "
"backend."
)
vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py
+137
View File
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU MLA prefill backend implemented with dense per-request attention.
This backend is selected on CPU devices, where flash-attn is unavailable. It
computes MLA prefill attention by slicing the ragged (varlen) q/k/v tensors
per request and evaluating dense attention in PyTorch.
"""
from typing import TYPE_CHECKING
import torch
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
if TYPE_CHECKING:
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonPrefillMetadata,
)
class CPUSDPAMLAPrefillBackend(MLAPrefillBackend):
"""MLA prefill backend for CPU using dense PyTorch attention.
flash-attn is not installable on CPU, so this backend provides an
equivalent prefill implementation using per-request dense attention.
"""
@staticmethod
def get_name() -> str:
return "CPU_SDPA_MLA"
@classmethod
def is_available(cls) -> bool:
# SDPA is always available on CPU.
return True
def supports_out(self) -> bool:
return True
def _ragged_sdpa(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q_cu_seq_lens: torch.Tensor,
kv_cu_seq_lens: torch.Tensor,
causal: bool,
return_softmax_lse: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
# q/k/v: [num_tokens, num_heads, head_dim] (ragged, concatenated).
# q_cu_seq_lens/kv_cu_seq_lens: [num_reqs + 1] row offsets for each
# request's query and key/value tokens.
num_reqs = q_cu_seq_lens.numel() - 1
assert num_reqs == kv_cu_seq_lens.numel() - 1
out_chunks: list[torch.Tensor] = []
lse_chunks: list[torch.Tensor] = []
for r in range(num_reqs):
q_start = int(q_cu_seq_lens[r])
q_end = int(q_cu_seq_lens[r + 1])
kv_start = int(kv_cu_seq_lens[r])
kv_end = int(kv_cu_seq_lens[r + 1])
q_r = q[q_start:q_end].transpose(0, 1).float() # [NH, Sq, Hq]
k_r = k[kv_start:kv_end].transpose(0, 1).float() # [NH, Sk, Hq]
v_r = v[kv_start:kv_end].transpose(0, 1).float() # [NH, Sk, Hv]
attn_scores = torch.matmul(q_r, k_r.transpose(-2, -1)) * self.scale
if causal:
causal_mask = torch.ones(
(q_end - q_start, kv_end - kv_start),
dtype=torch.bool,
device=attn_scores.device,
).triu(1)
attn_scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf"))
attn_probs = torch.softmax(attn_scores, dim=-1, dtype=torch.float32)
out_r = torch.matmul(attn_probs, v_r).transpose(0, 1).to(v.dtype)
out_chunks.append(out_r) # [Sq, NH, Hv]
if return_softmax_lse:
lse_chunks.append(torch.logsumexp(attn_scores, dim=-1))
out = torch.cat(out_chunks, dim=0)
if not return_softmax_lse:
return out
lse = torch.cat(lse_chunks, dim=1)
return out, lse
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
out: torch.Tensor | None = None,
output_scale: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
assert output_scale is None, (
"CPU MLA prefill backend does not support fused FP8 output"
)
out_t = self._ragged_sdpa(
q,
k,
v,
self._prefill_metadata.query_start_loc,
self._prefill_metadata.query_start_loc,
causal=True,
return_softmax_lse=return_softmax_lse,
)
if out is not None and not return_softmax_lse:
assert isinstance(out_t, torch.Tensor)
out.copy_(out_t)
return out
return out_t
def run_prefill_context_chunk(
self,
chunk: "MLACommonPrefillMetadata.ContextChunk",
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
assert out is None, "CPU MLA context prefill does not support in-place output"
result = self._ragged_sdpa(
q,
k,
v,
chunk.query_start_loc,
chunk.cu_seq_lens,
causal=False,
return_softmax_lse=True,
)
assert isinstance(result, tuple)
out, lse = result
return out, lse
vllm/v1/attention/backends/mla/prefill/registry.py
+1 -3
View File
@@ -52,9 +52,7 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
"vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
"AiterFlashAttnPrefillBackend"
)
CPU_NATIVE = (
"vllm.v1.attention.backends.mla.prefill.cpu_native.CPUNativeMLAPrefillBackend"
)
CPU = "vllm.v1.attention.backends.mla.prefill.cpu_sdpa.CPUSDPAMLAPrefillBackend"
# Placeholder for third-party/custom backends - must be registered before use
# set to None to avoid alias with other backend, whose value is an empty string
CUSTOM = None
vllm/v1/attention/backends/mla/prefill/selector.py
+4 -2
View File
@@ -107,10 +107,12 @@ def get_mla_prefill_backend(
"""
from vllm.platforms import current_platform
if current_platform.is_cpu():
logger.info_once("Using CPU SDPA MLA prefill backend.")
return MLAPrefillBackendEnum.CPU.get_class()
device_capability = current_platform.get_device_capability()
if device_capability is None:
if current_platform.is_cpu():
return MLAPrefillBackendEnum.CPU_NATIVE.get_class()
logger.info_once(
"Device capability not available, using FlashAttention MLA prefill backend."
)
vllm/v1/attention/ops/merge_attn_states.py
+73
View File
@@ -6,6 +6,68 @@ import torch
from vllm.platforms import current_platform
def _merge_attn_states_torch(
output: torch.Tensor,
prefix_output: torch.Tensor,
prefix_lse: torch.Tensor,
suffix_output: torch.Tensor,
suffix_lse: torch.Tensor,
output_lse: torch.Tensor | None = None,
prefill_tokens_with_context: int | None = None,
output_scale: torch.Tensor | None = None,
) -> None:
assert output_scale is None, "CPU fallback does not support FP8 output"
num_tokens = output.shape[0]
if prefill_tokens_with_context is None:
prefill_tokens_with_context = num_tokens
prefix_lse = torch.where(
torch.isinf(prefix_lse) & (prefix_lse > 0),
torch.full_like(prefix_lse, float("-inf")),
prefix_lse,
)
suffix_lse = torch.where(
torch.isinf(suffix_lse) & (suffix_lse > 0),
torch.full_like(suffix_lse, float("-inf")),
suffix_lse,
)
if prefill_tokens_with_context < num_tokens:
output[prefill_tokens_with_context:].copy_(
suffix_output[prefill_tokens_with_context:]
)
if output_lse is not None:
output_lse[:, prefill_tokens_with_context:].copy_(
suffix_lse[:, prefill_tokens_with_context:]
)
if prefill_tokens_with_context == 0:
return
prefix_lse = prefix_lse[:, :prefill_tokens_with_context].float()
suffix_lse = suffix_lse[:, :prefill_tokens_with_context].float()
max_lse = torch.maximum(prefix_lse, suffix_lse)
prefix_se = torch.exp(prefix_lse - max_lse)
suffix_se = torch.exp(suffix_lse - max_lse)
out_se = prefix_se + suffix_se
both_empty = max_lse == float("-inf")
prefix_scale = (prefix_se / out_se).transpose(0, 1).unsqueeze(-1)
suffix_scale = (suffix_se / out_se).transpose(0, 1).unsqueeze(-1)
merged = (
prefix_output[:prefill_tokens_with_context].float() * prefix_scale
+ suffix_output[:prefill_tokens_with_context].float() * suffix_scale
)
merged.masked_fill_(both_empty.transpose(0, 1).unsqueeze(-1), 0.0)
output[:prefill_tokens_with_context].copy_(merged.to(output.dtype))
if output_lse is not None:
merged_lse = torch.log(out_se) + max_lse
merged_lse.masked_fill_(both_empty, float("-inf"))
output_lse[:, :prefill_tokens_with_context].copy_(merged_lse)
def merge_attn_states(
output: torch.Tensor,
prefix_output: torch.Tensor,
@@ -94,6 +156,17 @@ def merge_attn_states(
prefill_tokens_with_context,
output_scale,
)
elif current_platform.is_cpu():
return _merge_attn_states_torch(
output,
prefix_output,
prefix_lse,
suffix_output,
suffix_lse,
output_lse,
prefill_tokens_with_context,
output_scale,
)
else:
from vllm.v1.attention.ops.triton_merge_attn_states import (
merge_attn_states,