From dcd07967caa2537c68aeeaddca4315b1319512b1 Mon Sep 17 00:00:00 2001 From: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:06:03 +0800 Subject: [PATCH] [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 --- .../test_cpu_linear_attn_chunked_prefix.py | 22 +++++++++++++++++ vllm/platforms/cpu.py | 24 ++++++++++++------- 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py b/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py index 1b6a48b7d2..d56df03dbc 100644 --- a/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py +++ b/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] diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 542a7a405d..91ff96ab33 100644 --- a/vllm/platforms/cpu.py +++ b/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)