[CPU][BugFix] Remove redundant kv cache write (#50437)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: Claude <noreply@anthropic.com>
Fadi Arafeh committed
Jul 31, 2026 at 08:28 UTC
10e6b400150c8d2cbedad54260def4871d464667
3 files changed
+45
-12
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+3
-2
@@ -56,10 +56,11 @@ function cpu_tests() {
56
set -e
57
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m"
58
59
- # Run model tests
59
+ # Test encoder-decoder and encoder-only models
60
docker exec cpu-test bash -c "
61
set -e
62
- pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model"
62
+ pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model
63
+ pytest -x -v -s 'tests/models/language/pooling/test_embedding.py::test_models[sentence-transformers/all-MiniLM-L12-v2]'"
64
65
# Run quantized model tests
66
docker exec cpu-test bash -c "
tests/kernels/attention/test_cpu_attn.py
+36
@@ -687,6 +687,42 @@ def test_varlen_encoder_attention_vec(
687
)
688
689
690
+@pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS)
691
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
692
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
693
+@pytest.mark.parametrize(
694
+ "block_size",
695
+ [
696
+ 128,
697
+ ],
698
+)
699
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
700
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
701
+@pytest.mark.parametrize("isa", ["neon"])
702
+@pytest.mark.skipif(
703
+ current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
704
+ reason="Not an Arm CPU.",
705
+)
706
+def test_varlen_encoder_attention_neon(
707
+ seq_lens: list[int],
708
+ num_heads: tuple[int, int],
709
+ head_size: int,
710
+ sliding_window: int | None,
711
+ dtype: torch.dtype,
712
+ block_size: int,
713
+ isa: str,
714
+) -> None:
715
+ varlen_encoder_attention(
716
+ seq_lens=seq_lens,
717
+ num_heads=num_heads,
718
+ head_size=head_size,
719
+ sliding_window=sliding_window,
720
+ dtype=dtype,
721
+ block_size=block_size,
722
+ isa=isa,
723
+ )
724
+
725
+
726
@pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS)
727
@pytest.mark.parametrize("num_heads", NUM_HEADS)
728
@pytest.mark.parametrize("head_size", HEAD_SIZES)
vllm/v1/attention/backends/cpu_attn.py
+6
-10
@@ -350,8 +350,11 @@ class CPUAttentionBackendImpl(AttentionImpl):
350
351
num_actual_tokens = attn_metadata.num_actual_tokens
352
353
- # For encoder attention
354
- if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
353
+ is_encoder_attention = self.attn_type in (
354
+ AttentionType.ENCODER_ONLY,
355
+ AttentionType.ENCODER,
356
+ )
357
+ if is_encoder_attention:
358
# For encoder attention,
359
kv_cache = attn_metadata.encoder_cache
360
@@ -362,14 +365,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
365
kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
366
key_cache, value_cache = kv_cache.chunk(2, dim=2)
367
365
- # key and value may be None in the case of cross attention. They are
366
- # calculated once based on the output from the encoder and then cached
367
- # in KV cache.
368
- if (
369
- self.kv_sharing_target_layer_name is None
370
- and key is not None
371
- and value is not None
372
- ):
368
+ if is_encoder_attention:
369
ops.cpu_attn_reshape_and_cache(
370
key,
371
value,