From 10e6b400150c8d2cbedad54260def4871d464667 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Fri, 31 Jul 2026 08:28:44 +0100 Subject: [PATCH] [CPU][BugFix] Remove redundant kv cache write (#50437) Signed-off-by: Fadi Arafeh Signed-off-by: jiang1.li Co-authored-by: jiang1.li Co-authored-by: Claude --- .../scripts/hardware_ci/run-cpu-test-arm.sh | 5 +-- tests/kernels/attention/test_cpu_attn.py | 36 +++++++++++++++++++ vllm/v1/attention/backends/cpu_attn.py | 16 ++++----- 3 files changed, 45 insertions(+), 12 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 09396a9697..7c8ee86eed 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -56,10 +56,11 @@ function cpu_tests() { set -e python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m" - # Run model tests + # Test encoder-decoder and encoder-only models docker exec cpu-test bash -c " set -e - pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model" + pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model + pytest -x -v -s 'tests/models/language/pooling/test_embedding.py::test_models[sentence-transformers/all-MiniLM-L12-v2]'" # Run quantized model tests docker exec cpu-test bash -c " diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index 3eb660aadb..ff8e6db92b 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -687,6 +687,42 @@ def test_varlen_encoder_attention_vec( ) +@pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize( + "block_size", + [ + 128, + ], +) +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("isa", ["neon"]) +@pytest.mark.skipif( + current_platform.get_cpu_architecture() != CpuArchEnum.ARM, + reason="Not an Arm CPU.", +) +def test_varlen_encoder_attention_neon( + seq_lens: list[int], + num_heads: tuple[int, int], + head_size: int, + sliding_window: int | None, + dtype: torch.dtype, + block_size: int, + isa: str, +) -> None: + varlen_encoder_attention( + seq_lens=seq_lens, + num_heads=num_heads, + head_size=head_size, + sliding_window=sliding_window, + dtype=dtype, + block_size=block_size, + isa=isa, + ) + + @pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index ea63f39522..a1b71f036c 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -350,8 +350,11 @@ class CPUAttentionBackendImpl(AttentionImpl): num_actual_tokens = attn_metadata.num_actual_tokens - # For encoder attention - if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): + is_encoder_attention = self.attn_type in ( + AttentionType.ENCODER_ONLY, + AttentionType.ENCODER, + ) + if is_encoder_attention: # For encoder attention, kv_cache = attn_metadata.encoder_cache @@ -362,14 +365,7 @@ class CPUAttentionBackendImpl(AttentionImpl): kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1)) key_cache, value_cache = kv_cache.chunk(2, dim=2) - # key and value may be None in the case of cross attention. They are - # calculated once based on the output from the encoder and then cached - # in KV cache. - if ( - self.kv_sharing_target_layer_name is None - and key is not None - and value is not None - ): + if is_encoder_attention: ops.cpu_attn_reshape_and_cache( key, value,