[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>

.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 "

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)

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,