[CPU] Keep the hybrid block size across the engine core's config rerun (#61805)

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py

@@ -8,6 +8,7 @@ import pytest
 
 from tests.models.utils import check_logprobs_close
 from vllm import LLM, SamplingParams
+from vllm.inputs import TokensPrompt
 from vllm.platforms import current_platform
 
 pytestmark = pytest.mark.cpu_model
@@ -112,3 +113,24 @@ def test_prefix_cache_hit_matches_cold_cache(full_prefill_refs):
         name_0="cold_cache",
         name_1="warm_cache",
     )
+
+
+def test_prefix_cache_hit_reaches_last_full_block():
+    """A repeated prompt must hit every full block of the hybrid block size.
+
+    The workers raise the block size to fit the GDN state pages, and the
+    engine core reruns its config checks after its handshake; the raised size
+    must survive them, or the prefill's state checkpoint misses the block
+    boundaries and nothing is reused. One chunk holds the whole prompt.
+    """
+    llm = _make_llm(enable_prefix_caching=True, max_num_batched_tokens=2048)
+    block_size = llm.llm_engine.vllm_config.cache_config.block_size
+    # Ends mid-block: only the last full block's state can be reused.
+    num_tokens = block_size + block_size // 2
+    assert num_tokens <= 2048
+    prompt = TokensPrompt(prompt_token_ids=[1000 + i for i in range(num_tokens)])
+    sp = SamplingParams(max_tokens=1)
+    hits = [llm.generate(prompt, sp)[0].num_cached_tokens for _ in range(2)]
+    del llm
+
+    assert hits == [0, block_size]

vllm/platforms/cpu.py

@@ -289,10 +289,9 @@ class CpuPlatform(Platform):
                     cache_config.block_size,
                 )
             cache_config.block_size = 256
-        elif not cache_config.user_specified_block_size:
-            cache_config.block_size = 128
-
-        if not reference_cpu_mla_enabled and cache_config.block_size % 32 != 0:
+        elif (
+            cache_config.user_specified_block_size and cache_config.block_size % 32 != 0
+        ):
             logger.warning(
                 "CPU backend prefers block_size is multiples of 32, "
                 "otherwise the performance is not optimized."
@@ -532,15 +531,24 @@ class CpuPlatform(Platform):
     @classmethod
     def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
         model_config = vllm_config.model_config
-        if model_config is None or not model_config.is_hybrid:
+        if model_config is None:
             return
-
-        # reconcile attention and mamba page sizes
         backend_classes = cls._find_non_ssm_backends(vllm_config)
         if not backend_classes:
             return
 
-        pre_block_size = vllm_config.cache_config.block_size
+        # Not in check_and_update_config(), which reruns with every
+        # VllmConfig.__post_init__ and would undo the hybrid alignment below.
+        cache_config = vllm_config.cache_config
+        if not cache_config.user_specified_block_size:
+            cache_config.block_size = cls._preferred_block_size_for_backends(
+                backend_classes, 128, vllm_config
+            )
+        if not model_config.is_hybrid:
+            return
+
+        # reconcile attention and mamba page sizes
+        pre_block_size = cache_config.block_size
         cls._align_hybrid_block_size(vllm_config, backend_classes[0])
         cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)