[Bugfix][KV Offload] Reject misaligned chunked-local CPU suffix hits

docs/features/kv_offloading_usage.md

@@ -478,6 +478,13 @@ selects disk slots and is checked against `disk_capacity_bytes`. Omitting
 `num_cpu_blocks` preserves automatic per-worker sizing. This option applies
 to `SimpleCPUOffloadConnector`, separately from `OffloadingConnector` above.
 
+## Chunked-local attention with `SimpleCPUOffloadConnector`
+
+For chunked-local attention, CPU cache reuse requires the GPU-local prefix to
+end at an attention-chunk boundary for every chunked-local cache group. Otherwise,
+the connector keeps the GPU prefix and recomputes the remaining tokens. Requests
+with no GPU-local prefix, or with a chunk-aligned prefix, can still reuse CPU KV.
+
 ## Further Reading
 
 - [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).

tests/v1/simple_kv_offload/test_integration.py

@@ -5,8 +5,11 @@
 import time
 
 import pytest
+import torch
+from transformers import Llama4ForCausalLM, Llama4TextConfig
 
-from vllm import LLM, SamplingParams, TokensPrompt
+from tests.models.utils import check_logprobs_close
+from vllm import LLM, SamplingParams, TokensPrompt, envs
 from vllm.config import KVTransferConfig
 from vllm.platforms import current_platform
 
@@ -26,6 +29,90 @@
 ]
 
 
+def test_chunked_local_partial_gpu_hit(vllm_runner, tmp_path, monkeypatch):
+    """A CPU suffix after a short GPU hit must preserve the cold prediction."""
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    envs.disable_envs_cache()
+    config = Llama4TextConfig(
+        vocab_size=512,
+        hidden_size=128,
+        intermediate_size=256,
+        intermediate_size_mlp=256,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        head_dim=64,
+        max_position_embeddings=256,
+        attention_chunk_size=64,
+        no_rope_layers=[1, 0],
+        interleave_moe_layer_step=100,
+        num_local_experts=2,
+        use_qk_norm=False,
+        initializer_range=0.08,
+    )
+    with torch.random.fork_rng(devices=[]):
+        torch.manual_seed(1234)
+        model = Llama4ForCausalLM(config)
+        model.save_pretrained(tmp_path)
+        del model
+
+    with vllm_runner(
+        str(tmp_path),
+        skip_tokenizer_init=True,
+        dtype="float16",
+        enforce_eager=True,
+        compilation_config={"mode": 0},
+        attention_config={"backend": "TRITON_ATTN"},
+        max_model_len=256,
+        max_num_batched_tokens=256,
+        max_num_seqs=1,
+        block_size=16,
+        num_gpu_blocks_override=64,
+        kv_cache_memory_bytes=16 << 20,
+        enable_prefix_caching=True,
+        disable_hybrid_kv_cache_manager=False,
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="SimpleCPUOffloadConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={"cpu_bytes_to_use": 16 << 20},
+        ),
+    ) as runner:
+        llm = runner.get_llm()
+        tokens = [3 + (i * 37) % 500 for i in range(81)]
+        prompt = TokensPrompt(prompt_token_ids=tokens)
+        params = SamplingParams(
+            temperature=0, max_tokens=1, logprobs=5, ignore_eos=True, detokenize=False
+        )
+        cold = llm.generate(prompt, params, use_tqdm=False)[0].outputs[0]
+
+        def reset_gpu_cache():
+            deadline = time.monotonic() + 30
+            while not llm.reset_prefix_cache():
+                assert time.monotonic() < deadline, "CPU stores did not finish"
+                llm.llm_engine.step()
+
+        reset_gpu_cache()
+        cpu_hit = llm.generate(prompt, params, use_tqdm=False)[0]
+        assert cpu_hit.num_cached_tokens == 80
+        reset_gpu_cache()
+        short_params = params.clone()
+        short_params.skip_reading_prefix_cache = True
+        llm.generate(
+            TokensPrompt(prompt_token_ids=tokens[:17]), short_params, use_tqdm=False
+        )
+        mixed_hit = llm.generate(prompt, params, use_tqdm=False)[0]
+        for name, output in (("CPU hit", cpu_hit), ("mixed hit", mixed_hit)):
+            cached = output.outputs[0]
+            check_logprobs_close(
+                outputs_0_lst=[(list(cold.token_ids), "", cold.logprobs)],
+                outputs_1_lst=[(list(cached.token_ids), "", cached.logprobs)],
+                name_0="cold",
+                name_1=name,
+                always_check_logprobs=True,
+            )
+        assert mixed_hit.num_cached_tokens == 16
+
+
 def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
     kv_transfer_config = KVTransferConfig(
         kv_connector="SimpleCPUOffloadConnector",

tests/v1/simple_kv_offload/test_scheduler.py

@@ -45,6 +45,7 @@
     register_all_kvcache_specs,
 )
 from vllm.v1.kv_cache_interface import (
+    ChunkedLocalAttentionSpec,
     CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
@@ -1789,6 +1790,40 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None:
         )
 
 
+@pytest.mark.parametrize(
+    "local_tokens,cached_chunk_blocks,expected_external_tokens",
+    [(0, 4, 64), (0, 5, 80), (16, 4, 0), (16, 5, 0), (64, 4, 0), (64, 5, 16)],
+)
+def test_chunked_local_cpu_lookup_requires_chunk_aligned_prefix(
+    local_tokens: int, cached_chunk_blocks: int, expected_external_tokens: int
+) -> None:
+    """Recompute suffixes whose relative and absolute chunk boundaries differ."""
+    config = _make_kv_cache_config(num_blocks=32, num_groups=2)
+    config.kv_cache_groups[1].kv_cache_spec = ChunkedLocalAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+        attention_chunk_size=64,
+    )
+    fix = make_scheduler(num_cpu_blocks=32, kv_cache_config=config)
+    scheduler = fix.scheduler
+    request = make_request(num_blocks=5)
+    pool = scheduler.cpu_block_pool
+    cached = _allocate_gpu_blocks(pool, request, 5, group_id=0)
+    cached += _allocate_gpu_blocks(pool, request, cached_chunk_blocks, group_id=1)
+    pool.free_blocks(cached)
+    free_before_lookup = pool.get_num_free_blocks()
+
+    assert scheduler.get_num_new_matched_tokens(request, local_tokens) == (
+        expected_external_tokens,
+        expected_external_tokens > 0,
+    )
+
+    scheduler.request_finished(request, [])
+    assert pool.get_num_free_blocks() == free_before_lookup
+
+
 # ---------------------------------------------------------------------------
 # Test 11: TOCTOU between Phase A and Phase B (regression for #39702)
 # ---------------------------------------------------------------------------

vllm/v1/simple_kv_offload/manager.py

@@ -38,6 +38,7 @@
 )
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import (
+    ChunkedLocalAttentionSpec,
     FullAttentionSpec,
     MambaSpec,
     SlidingWindowSpec,
@@ -385,6 +386,15 @@ def get_num_new_matched_tokens(
             )
             return 0, False
 
+        # Suffix-relative chunk boundaries must match absolute token positions.
+        for group_id in self.prefix_cacheable_group_ids:
+            spec = self.cpu_kv_cache_config.kv_cache_groups[group_id].kv_cache_spec
+            if (
+                isinstance(spec, ChunkedLocalAttentionSpec)
+                and num_computed_tokens % spec.attention_chunk_size
+            ):
+                return 0, False
+
         num_skipped_hashes = num_computed_tokens // self.hash_block_size
         remaining_hashes = request.block_hashes[num_skipped_hashes:]