@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6f7df92a8e

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

Signed-off-by: baoloongmao <baoloongmao@tencent.com> Co-authored-by: Codex <codex@openai.com>

maobaolong committed Aug 28, 2026 at 12:36 UTC 6f7df92a8e6cdc74a725b8f10b4d0b48ba2b37ef
9 files changed +442 -69
tests/v1/attention/test_cpu_mla_backend.py
+164 -2
@@ -19,6 +19,51 @@ import pytest
19 import torch
20
21 from vllm.platforms import current_platform
22 +from vllm.v1.attention.backends.mla.prefill.cpu_sdpa import (
23 + CPUSDPAMLAPrefillBackend,
24 +)
25 +from vllm.v1.attention.backends.mla.prefill.selector import (
26 + get_mla_prefill_backend,
27 +)
28 +
29 +
30 +def _reference_ragged_attention(
31 + q: torch.Tensor,
32 + k: torch.Tensor,
33 + v: torch.Tensor,
34 + q_cu_seq_lens: torch.Tensor,
35 + kv_cu_seq_lens: torch.Tensor,
36 + *,
37 + scale: float,
38 + causal: bool,
39 +) -> tuple[torch.Tensor, torch.Tensor]:
40 + outputs: list[torch.Tensor] = []
41 + lses: list[torch.Tensor] = []
42 + num_reqs = q_cu_seq_lens.numel() - 1
43 + for req_idx in range(num_reqs):
44 + q_start = int(q_cu_seq_lens[req_idx])
45 + q_end = int(q_cu_seq_lens[req_idx + 1])
46 + kv_start = int(kv_cu_seq_lens[req_idx])
47 + kv_end = int(kv_cu_seq_lens[req_idx + 1])
48 +
49 + q_r = q[q_start:q_end].transpose(0, 1).float()
50 + k_r = k[kv_start:kv_end].transpose(0, 1).float()
51 + v_r = v[kv_start:kv_end].transpose(0, 1).float()
52 +
53 + scores = torch.matmul(q_r, k_r.transpose(-2, -1)) * scale
54 + if causal:
55 + causal_mask = torch.ones(
56 + (q_end - q_start, kv_end - kv_start),
57 + dtype=torch.bool,
58 + device=scores.device,
59 + ).triu(1)
60 + scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf"))
61 +
62 + probs = torch.softmax(scores, dim=-1)
63 + outputs.append(torch.matmul(probs, v_r).transpose(0, 1).to(v.dtype))
64 + lses.append(torch.logsumexp(scores, dim=-1))
65 +
66 + return torch.cat(outputs, dim=0), torch.cat(lses, dim=1)
67
68
69 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
@@ -58,7 +103,7 @@ def test_kv_cache_cpu_write() -> None:
103
104 @pytest.mark.cpu_model
105 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
61 -def test_cpu_mla_backend_smoke() -> None:
106 +def test_cpu_mla_backend_smoke(tmp_path) -> None:
107 # Shrink the architecture aggressively so the test is CI-friendly.
108 from vllm import LLM, SamplingParams
109
@@ -70,19 +115,136 @@ def test_cpu_mla_backend_smoke() -> None:
115 }
116 llm = LLM(
117 model="deepseek-ai/DeepSeek-V2-Lite",
118 + skip_tokenizer_init=True,
119 trust_remote_code=True,
120 load_format="dummy",
121 enforce_eager=True,
122 max_model_len=128,
123 max_num_seqs=2,
124 block_size=16,
125 + gpu_memory_utilization=0.25,
126 hf_overrides=hf_overrides,
127 + kv_transfer_config={
128 + "kv_connector": "ExampleConnector",
129 + "kv_role": "kv_both",
130 + "kv_connector_extra_config": {
131 + "shared_storage_path": str(tmp_path),
132 + },
133 + },
134 )
135
136 sampling_params = SamplingParams(max_tokens=4, temperature=0.0)
83 - outputs = llm.generate(["Hello", "MLA on CPU"], sampling_params)
137 + outputs = llm.generate([[1, 2, 3, 4], [5, 6, 7]], sampling_params, use_tqdm=False)
138 assert len(outputs) == 2
139 for output in outputs:
140 # `dummy` weights do not produce meaningful text, but the number
141 # of generated tokens must match what we asked for.
142 assert len(output.outputs[0].token_ids) == 4
143 +
144 + prefix_hit_prompt = [1] * 20
145 + outputs = llm.generate([prefix_hit_prompt], sampling_params, use_tqdm=False)
146 + assert len(outputs[0].outputs[0].token_ids) == 4
147 + assert llm.reset_prefix_cache(reset_connector=False)
148 +
149 + # ExampleConnector persists KV externally at block granularity. With a
150 + # 20-token prompt and block_size=16, the second run reloads 16 cached
151 + # tokens and forces the 4-token suffix through MLA prefill with context.
152 + outputs = llm.generate([prefix_hit_prompt], sampling_params, use_tqdm=False)
153 + assert len(outputs[0].outputs[0].token_ids) == 4
154 +
155 +
156 +@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
157 +def test_cpu_mla_prefill_backend_selected() -> None:
158 + # On CPU (no device capability) the MLA prefill backend must be the
159 + # SDPA-based CPU backend, not flash-attn which is unavailable on CPU.
160 + backend_cls = get_mla_prefill_backend(None)
161 + assert backend_cls is CPUSDPAMLAPrefillBackend
162 +
163 +
164 +@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
165 +def test_cpu_mla_prefill_new_tokens() -> None:
166 + # The SDPA prefill must exactly match a reference causal attention over
167 + # the ragged (varlen) q/k/v layout, with per-request padding handled.
168 + backend = CPUSDPAMLAPrefillBackend(
169 + num_heads=4,
170 + scale=0.25,
171 + kv_lora_rank=16,
172 + qk_nope_head_dim=16,
173 + qk_rope_head_dim=0,
174 + v_head_dim=16,
175 + vllm_config=None,
176 + )
177 +
178 + class _PrefillMeta:
179 + query_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32)
180 +
181 + backend.prepare_metadata(_PrefillMeta())
182 + q = torch.randn(6, 4, 16, dtype=torch.bfloat16)
183 + k = torch.randn(6, 4, 16, dtype=torch.bfloat16)
184 + v = torch.randn(6, 4, 16, dtype=torch.bfloat16)
185 + out = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=False)
186 + ref, _ = _reference_ragged_attention(
187 + q,
188 + k,
189 + v,
190 + _PrefillMeta.query_start_loc,
191 + _PrefillMeta.query_start_loc,
192 + scale=0.25,
193 + causal=True,
194 + )
195 + assert out.float().allclose(ref.float(), atol=2e-2)
196 + assert out.shape == (6, 4, 16)
197 +
198 + out_lse, lse = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
199 + ref_out, ref_lse = _reference_ragged_attention(
200 + q,
201 + k,
202 + v,
203 + _PrefillMeta.query_start_loc,
204 + _PrefillMeta.query_start_loc,
205 + scale=0.25,
206 + causal=True,
207 + )
208 + assert out_lse.float().allclose(ref_out.float(), atol=2e-2)
209 + assert lse.float().allclose(ref_lse.float(), atol=2e-2)
210 + assert lse.shape == (4, 6)
211 +
212 +
213 +@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
214 +def test_cpu_mla_prefill_context_chunk() -> None:
215 + backend = CPUSDPAMLAPrefillBackend(
216 + num_heads=4,
217 + scale=0.25,
218 + kv_lora_rank=16,
219 + qk_nope_head_dim=16,
220 + qk_rope_head_dim=0,
221 + v_head_dim=16,
222 + vllm_config=None,
223 + )
224 +
225 + class _PrefillMeta:
226 + query_start_loc = torch.tensor([0, 5], dtype=torch.int32)
227 +
228 + class _Chunk:
229 + query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32)
230 + cu_seq_lens = torch.tensor([0, 3, 7], dtype=torch.int32)
231 +
232 + backend.prepare_metadata(_PrefillMeta())
233 + q = torch.randn(5, 4, 16, dtype=torch.bfloat16)
234 + k = torch.randn(7, 4, 16, dtype=torch.bfloat16)
235 + v = torch.randn(7, 4, 16, dtype=torch.bfloat16)
236 +
237 + out, lse = backend.run_prefill_context_chunk(_Chunk(), q, k, v)
238 + ref_out, ref_lse = _reference_ragged_attention(
239 + q,
240 + k,
241 + v,
242 + _Chunk.query_start_loc,
243 + _Chunk.cu_seq_lens,
244 + scale=0.25,
245 + causal=False,
246 + )
247 + assert out.float().allclose(ref_out.float(), atol=2e-2)
248 + assert lse.float().allclose(ref_lse.float(), atol=2e-2)
249 + assert out.shape == (5, 4, 16)
250 + assert lse.shape == (4, 5)
tests/v1/attention/test_mla_prefill_selector.py
+26 -1
@@ -58,10 +58,31 @@ def _make_vllm_config(
58 class TestGetMLAPrefillBackend:
59 """Tests for get_mla_prefill_backend (public API)."""
60
61 - def test_no_device_capability_returns_flash_attn(self):
61 + def test_cpu_uses_sdpa_prefill(self):
62 vllm_config = _make_vllm_config()
63
64 with patch("vllm.platforms.current_platform") as mock_platform:
65 + mock_platform.is_cpu.return_value = True
66 +
67 + backend = get_mla_prefill_backend(vllm_config)
68 + assert backend is MLAPrefillBackendEnum.CPU.get_class()
69 +
70 + def test_no_device_capability_returns_flash_attn(self):
71 + vllm_config = _make_vllm_config()
72 +
73 + class FlashAttnBackend:
74 + @staticmethod
75 + def get_name():
76 + return "FLASH_ATTN"
77 +
78 + with (
79 + patch("vllm.platforms.current_platform") as mock_platform,
80 + patch.object(
81 + MLAPrefillBackendEnum.FLASH_ATTN,
82 + "get_class",
83 + return_value=FlashAttnBackend,
84 + ),
85 + ):
86 mock_platform.get_device_capability.return_value = None
87 mock_platform.is_cpu.return_value = False
88
@@ -80,6 +101,7 @@ class TestGetMLAPrefillBackend:
101 )
102
103 with patch("vllm.platforms.current_platform") as mock_platform:
104 + mock_platform.is_cpu.return_value = False
105 mock_platform.get_device_capability.return_value = DeviceCapability(
106 major=9, minor=0
107 )
@@ -98,6 +120,7 @@ class TestGetMLAPrefillBackend:
120 )
121
122 with patch("vllm.platforms.current_platform") as mock_platform:
123 + mock_platform.is_cpu.return_value = False
124 mock_platform.get_device_capability.return_value = DeviceCapability(
125 major=9, minor=0
126 )
@@ -111,6 +134,7 @@ class TestGetMLAPrefillBackend:
134 )
135
136 with patch("vllm.platforms.current_platform") as mock_platform:
137 + mock_platform.is_cpu.return_value = False
138 mock_platform.get_device_capability.return_value = DeviceCapability(
139 major=10, minor=0
140 )
@@ -153,6 +177,7 @@ class TestGetMLAPrefillBackend:
177 return_value=3,
178 ),
179 ):
180 + mock_platform.is_cpu.return_value = False
181 mock_platform.get_device_capability.return_value = DeviceCapability(
182 major=9, minor=0
183 )
vllm/_custom_ops.py
+25
@@ -174,6 +174,31 @@ def mla_decode_kvcache_cpu(
174 torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens)
175
176
177 +def gather_mla_context_cache_cpu(
178 + *,
179 + src_cache: torch.Tensor,
180 + dst: torch.Tensor,
181 + block_table: torch.Tensor,
182 + starts: torch.Tensor,
183 + cu_seq_lens: torch.Tensor,
184 +) -> None:
185 + page_size = src_cache.shape[1]
186 + flat_cache = src_cache.view(-1, src_cache.shape[-1])
187 + seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]
188 +
189 + out_start = 0
190 + for req_idx, (start, seq_len) in enumerate(zip(starts.tolist(), seq_lens.tolist())):
191 + if seq_len <= 0:
192 + continue
193 + positions = torch.arange(start, start + seq_len, device=block_table.device)
194 + block_ids = block_table[req_idx, positions // page_size].to(torch.long)
195 + offsets = positions % page_size
196 + slots = block_ids * page_size + offsets
197 + out_end = out_start + seq_len
198 + dst[out_start:out_end].copy_(flat_cache[slots])
199 + out_start = out_end
200 +
201 +
202 # merge attn states ops
203 def merge_attn_states(
204 output: torch.Tensor,
vllm/model_executor/layers/attention/mla_attention.py
+12
@@ -2785,6 +2785,18 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2785 batch_size=chunk.num_requests,
2786 seq_starts=chunk.starts,
2787 )
2788 + elif current_platform.is_cpu():
2789 + assert not is_quantized_kv_cache(self.kv_cache_dtype), (
2790 + "CPU MLA context gather fallback only supports "
2791 + f"non-quantized KV cache, got {self.kv_cache_dtype}"
2792 + )
2793 + ops.gather_mla_context_cache_cpu(
2794 + src_cache=kv_c_and_k_pe_cache,
2795 + dst=workspace[:toks],
2796 + block_table=block_table,
2797 + starts=chunk.starts,
2798 + cu_seq_lens=chunk.cu_seq_lens,
2799 + )
2800 elif not use_fp8_prefill:
2801 ops.gather_and_maybe_dequant_cache(
2802 src_cache=kv_c_and_k_pe_cache,
vllm/v1/attention/backends/mla/prefill/cpu_native.py deleted
-61
@@ -1,61 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Inert MLA prefill backend for the CPU platform.
4 -
5 -CPU's MLA attention impl (``CPUMLAImpl``) fully overrides ``forward_mha`` with a
6 -kernel that attends directly against the paged latent KV cache (covering both
7 -fresh prefill and cached-prefix continuation in one pass), so it never calls
8 -into a pluggable ``MLAPrefillBackend``'s ``run_prefill_new_tokens``/
9 -``run_prefill_context_chunk``. This class exists only to satisfy
10 -``MLACommonMetadataBuilder``'s structural requirement that
11 -``MLAAttention.prefill_backend`` be a real, cloneable object.
12 -"""
13 -
14 -from typing import TYPE_CHECKING
15 -
16 -import torch
17 -
18 -from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
19 -
20 -if TYPE_CHECKING:
21 - from vllm.model_executor.layers.attention.mla_attention import (
22 - MLACommonPrefillMetadata,
23 - )
24 -
25 -
26 -class CPUNativeMLAPrefillBackend(MLAPrefillBackend):
27 - """Placeholder MLA prefill backend for CPU; never actually invoked."""
28 -
29 - @staticmethod
30 - def get_name() -> str:
31 - return "CPU_NATIVE"
32 -
33 - def run_prefill_new_tokens(
34 - self,
35 - q: torch.Tensor,
36 - k: torch.Tensor,
37 - v: torch.Tensor,
38 - return_softmax_lse: bool,
39 - out: torch.Tensor | None = None,
40 - output_scale: torch.Tensor | None = None,
41 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
42 - raise AssertionError(
43 - "CPUNativeMLAPrefillBackend.run_prefill_new_tokens is unreachable: "
44 - "CPUMLAImpl.forward_mha fully overrides the dense-MHA prefill path "
45 - "and never calls the pluggable prefill backend."
46 - )
47 -
48 - def run_prefill_context_chunk(
49 - self,
50 - chunk: "MLACommonPrefillMetadata.ContextChunk",
51 - q: torch.Tensor,
52 - k: torch.Tensor,
53 - v: torch.Tensor,
54 - out: torch.Tensor | None = None,
55 - ) -> tuple[torch.Tensor, torch.Tensor]:
56 - raise AssertionError(
57 - "CPUNativeMLAPrefillBackend.run_prefill_context_chunk is "
58 - "unreachable: CPUMLAImpl.forward_mha fully overrides the "
59 - "dense-MHA prefill path and never calls the pluggable prefill "
60 - "backend."
61 - )
vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py new
+137
@@ -0,0 +1,137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CPU MLA prefill backend implemented with dense per-request attention.
4 +
5 +This backend is selected on CPU devices, where flash-attn is unavailable. It
6 +computes MLA prefill attention by slicing the ragged (varlen) q/k/v tensors
7 +per request and evaluating dense attention in PyTorch.
8 +"""
9 +
10 +from typing import TYPE_CHECKING
11 +
12 +import torch
13 +
14 +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
15 +
16 +if TYPE_CHECKING:
17 + from vllm.model_executor.layers.attention.mla_attention import (
18 + MLACommonPrefillMetadata,
19 + )
20 +
21 +
22 +class CPUSDPAMLAPrefillBackend(MLAPrefillBackend):
23 + """MLA prefill backend for CPU using dense PyTorch attention.
24 +
25 + flash-attn is not installable on CPU, so this backend provides an
26 + equivalent prefill implementation using per-request dense attention.
27 + """
28 +
29 + @staticmethod
30 + def get_name() -> str:
31 + return "CPU_SDPA_MLA"
32 +
33 + @classmethod
34 + def is_available(cls) -> bool:
35 + # SDPA is always available on CPU.
36 + return True
37 +
38 + def supports_out(self) -> bool:
39 + return True
40 +
41 + def _ragged_sdpa(
42 + self,
43 + q: torch.Tensor,
44 + k: torch.Tensor,
45 + v: torch.Tensor,
46 + q_cu_seq_lens: torch.Tensor,
47 + kv_cu_seq_lens: torch.Tensor,
48 + causal: bool,
49 + return_softmax_lse: bool,
50 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
51 + # q/k/v: [num_tokens, num_heads, head_dim] (ragged, concatenated).
52 + # q_cu_seq_lens/kv_cu_seq_lens: [num_reqs + 1] row offsets for each
53 + # request's query and key/value tokens.
54 + num_reqs = q_cu_seq_lens.numel() - 1
55 + assert num_reqs == kv_cu_seq_lens.numel() - 1
56 +
57 + out_chunks: list[torch.Tensor] = []
58 + lse_chunks: list[torch.Tensor] = []
59 + for r in range(num_reqs):
60 + q_start = int(q_cu_seq_lens[r])
61 + q_end = int(q_cu_seq_lens[r + 1])
62 + kv_start = int(kv_cu_seq_lens[r])
63 + kv_end = int(kv_cu_seq_lens[r + 1])
64 +
65 + q_r = q[q_start:q_end].transpose(0, 1).float() # [NH, Sq, Hq]
66 + k_r = k[kv_start:kv_end].transpose(0, 1).float() # [NH, Sk, Hq]
67 + v_r = v[kv_start:kv_end].transpose(0, 1).float() # [NH, Sk, Hv]
68 +
69 + attn_scores = torch.matmul(q_r, k_r.transpose(-2, -1)) * self.scale
70 + if causal:
71 + causal_mask = torch.ones(
72 + (q_end - q_start, kv_end - kv_start),
73 + dtype=torch.bool,
74 + device=attn_scores.device,
75 + ).triu(1)
76 + attn_scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf"))
77 +
78 + attn_probs = torch.softmax(attn_scores, dim=-1, dtype=torch.float32)
79 + out_r = torch.matmul(attn_probs, v_r).transpose(0, 1).to(v.dtype)
80 + out_chunks.append(out_r) # [Sq, NH, Hv]
81 + if return_softmax_lse:
82 + lse_chunks.append(torch.logsumexp(attn_scores, dim=-1))
83 +
84 + out = torch.cat(out_chunks, dim=0)
85 + if not return_softmax_lse:
86 + return out
87 + lse = torch.cat(lse_chunks, dim=1)
88 + return out, lse
89 +
90 + def run_prefill_new_tokens(
91 + self,
92 + q: torch.Tensor,
93 + k: torch.Tensor,
94 + v: torch.Tensor,
95 + return_softmax_lse: bool,
96 + out: torch.Tensor | None = None,
97 + output_scale: torch.Tensor | None = None,
98 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
99 + assert output_scale is None, (
100 + "CPU MLA prefill backend does not support fused FP8 output"
101 + )
102 + out_t = self._ragged_sdpa(
103 + q,
104 + k,
105 + v,
106 + self._prefill_metadata.query_start_loc,
107 + self._prefill_metadata.query_start_loc,
108 + causal=True,
109 + return_softmax_lse=return_softmax_lse,
110 + )
111 + if out is not None and not return_softmax_lse:
112 + assert isinstance(out_t, torch.Tensor)
113 + out.copy_(out_t)
114 + return out
115 + return out_t
116 +
117 + def run_prefill_context_chunk(
118 + self,
119 + chunk: "MLACommonPrefillMetadata.ContextChunk",
120 + q: torch.Tensor,
121 + k: torch.Tensor,
122 + v: torch.Tensor,
123 + out: torch.Tensor | None = None,
124 + ) -> tuple[torch.Tensor, torch.Tensor]:
125 + assert out is None, "CPU MLA context prefill does not support in-place output"
126 + result = self._ragged_sdpa(
127 + q,
128 + k,
129 + v,
130 + chunk.query_start_loc,
131 + chunk.cu_seq_lens,
132 + causal=False,
133 + return_softmax_lse=True,
134 + )
135 + assert isinstance(result, tuple)
136 + out, lse = result
137 + return out, lse
vllm/v1/attention/backends/mla/prefill/registry.py
+1 -3
@@ -52,9 +52,7 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
52 "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
53 "AiterFlashAttnPrefillBackend"
54 )
55 - CPU_NATIVE = (
56 - "vllm.v1.attention.backends.mla.prefill.cpu_native.CPUNativeMLAPrefillBackend"
57 - )
55 + CPU = "vllm.v1.attention.backends.mla.prefill.cpu_sdpa.CPUSDPAMLAPrefillBackend"
56 # Placeholder for third-party/custom backends - must be registered before use
57 # set to None to avoid alias with other backend, whose value is an empty string
58 CUSTOM = None
vllm/v1/attention/backends/mla/prefill/selector.py
+4 -2
@@ -107,10 +107,12 @@ def get_mla_prefill_backend(
107 """
108 from vllm.platforms import current_platform
109
110 + if current_platform.is_cpu():
111 + logger.info_once("Using CPU SDPA MLA prefill backend.")
112 + return MLAPrefillBackendEnum.CPU.get_class()
113 +
114 device_capability = current_platform.get_device_capability()
115 if device_capability is None:
112 - if current_platform.is_cpu():
113 - return MLAPrefillBackendEnum.CPU_NATIVE.get_class()
116 logger.info_once(
117 "Device capability not available, using FlashAttention MLA prefill backend."
118 )
vllm/v1/attention/ops/merge_attn_states.py
+73
@@ -6,6 +6,68 @@ import torch
6 from vllm.platforms import current_platform
7
8
9 +def _merge_attn_states_torch(
10 + output: torch.Tensor,
11 + prefix_output: torch.Tensor,
12 + prefix_lse: torch.Tensor,
13 + suffix_output: torch.Tensor,
14 + suffix_lse: torch.Tensor,
15 + output_lse: torch.Tensor | None = None,
16 + prefill_tokens_with_context: int | None = None,
17 + output_scale: torch.Tensor | None = None,
18 +) -> None:
19 + assert output_scale is None, "CPU fallback does not support FP8 output"
20 +
21 + num_tokens = output.shape[0]
22 + if prefill_tokens_with_context is None:
23 + prefill_tokens_with_context = num_tokens
24 +
25 + prefix_lse = torch.where(
26 + torch.isinf(prefix_lse) & (prefix_lse > 0),
27 + torch.full_like(prefix_lse, float("-inf")),
28 + prefix_lse,
29 + )
30 + suffix_lse = torch.where(
31 + torch.isinf(suffix_lse) & (suffix_lse > 0),
32 + torch.full_like(suffix_lse, float("-inf")),
33 + suffix_lse,
34 + )
35 +
36 + if prefill_tokens_with_context < num_tokens:
37 + output[prefill_tokens_with_context:].copy_(
38 + suffix_output[prefill_tokens_with_context:]
39 + )
40 + if output_lse is not None:
41 + output_lse[:, prefill_tokens_with_context:].copy_(
42 + suffix_lse[:, prefill_tokens_with_context:]
43 + )
44 +
45 + if prefill_tokens_with_context == 0:
46 + return
47 +
48 + prefix_lse = prefix_lse[:, :prefill_tokens_with_context].float()
49 + suffix_lse = suffix_lse[:, :prefill_tokens_with_context].float()
50 + max_lse = torch.maximum(prefix_lse, suffix_lse)
51 + prefix_se = torch.exp(prefix_lse - max_lse)
52 + suffix_se = torch.exp(suffix_lse - max_lse)
53 + out_se = prefix_se + suffix_se
54 + both_empty = max_lse == float("-inf")
55 +
56 + prefix_scale = (prefix_se / out_se).transpose(0, 1).unsqueeze(-1)
57 + suffix_scale = (suffix_se / out_se).transpose(0, 1).unsqueeze(-1)
58 + merged = (
59 + prefix_output[:prefill_tokens_with_context].float() * prefix_scale
60 + + suffix_output[:prefill_tokens_with_context].float() * suffix_scale
61 + )
62 + merged.masked_fill_(both_empty.transpose(0, 1).unsqueeze(-1), 0.0)
63 + output[:prefill_tokens_with_context].copy_(merged.to(output.dtype))
64 +
65 + if output_lse is not None:
66 + merged_lse = torch.log(out_se) + max_lse
67 + merged_lse.masked_fill_(both_empty, float("-inf"))
68 + output_lse[:, :prefill_tokens_with_context].copy_(merged_lse)
69 +
70 +
71 def merge_attn_states(
72 output: torch.Tensor,
73 prefix_output: torch.Tensor,
@@ -94,6 +156,17 @@ def merge_attn_states(
156 prefill_tokens_with_context,
157 output_scale,
158 )
159 + elif current_platform.is_cpu():
160 + return _merge_attn_states_torch(
161 + output,
162 + prefix_output,
163 + prefix_lse,
164 + suffix_output,
165 + suffix_lse,
166 + output_lse,
167 + prefill_tokens_with_context,
168 + output_scale,
169 + )
170 else:
171 from vllm.v1.attention.ops.triton_merge_attn_states import (
172 merge_attn_states,