From 6f7df92a8e6cdc74a725b8f10b4d0b48ba2b37ef Mon Sep 17 00:00:00 2001 From: maobaolong Date: Fri, 28 Aug 2026 12:36:55 +0800 Subject: [PATCH] [CPU][MLA] Fix prefill backend selection so MLA runs end-to-end on CPU (#51471) Signed-off-by: baoloongmao Co-authored-by: Codex --- tests/v1/attention/test_cpu_mla_backend.py | 166 +++++++++++++++++- .../v1/attention/test_mla_prefill_selector.py | 27 ++- vllm/_custom_ops.py | 25 +++ .../layers/attention/mla_attention.py | 12 ++ .../backends/mla/prefill/cpu_native.py | 61 ------- .../backends/mla/prefill/cpu_sdpa.py | 137 +++++++++++++++ .../backends/mla/prefill/registry.py | 4 +- .../backends/mla/prefill/selector.py | 6 +- vllm/v1/attention/ops/merge_attn_states.py | 73 ++++++++ 9 files changed, 442 insertions(+), 69 deletions(-) delete mode 100644 vllm/v1/attention/backends/mla/prefill/cpu_native.py create mode 100644 vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py diff --git a/tests/v1/attention/test_cpu_mla_backend.py b/tests/v1/attention/test_cpu_mla_backend.py index 18576ef076..fc3867099e 100644 --- a/tests/v1/attention/test_cpu_mla_backend.py +++ b/tests/v1/attention/test_cpu_mla_backend.py @@ -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) diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index bb03aa496a..d5856ab6bd 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -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 ) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 550bf94c43..e54cba02b0 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index fad2854861..20b5e6f679 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -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, diff --git a/vllm/v1/attention/backends/mla/prefill/cpu_native.py b/vllm/v1/attention/backends/mla/prefill/cpu_native.py deleted file mode 100644 index 52ba8b7e24..0000000000 --- a/vllm/v1/attention/backends/mla/prefill/cpu_native.py +++ /dev/null @@ -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." - ) diff --git a/vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py b/vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py new file mode 100644 index 0000000000..d5c9afc0cc --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/cpu_sdpa.py @@ -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 diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index b521f4ff1e..e7e00a0d01 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -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 diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index 5cdfe3c6ca..c7e935fa75 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -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." ) diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py index 20c7503e9d..6791e1f5fd 100644 --- a/vllm/v1/attention/ops/merge_attn_states.py +++ b/vllm/v1/attention/ops/merge_attn_states.py @@ -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,