From 4cc0d9e24baeb302db5279847dd396a213605c9c Mon Sep 17 00:00:00 2001 From: CharlesWu <296364977+charxwu@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:24:03 +0800 Subject: [PATCH 1/2] [Bugfix][KV Offload] Reject misaligned chunked-local CPU suffix hits Retain the GPU prefix and recompute the suffix when CPU lookup would use chunk boundaries relative to a non-chunk-aligned local prefix. Add real scheduler and LLM regressions and document the reuse restriction. Co-authored-by: Codex Signed-off-by: CharlesWu <296364977+charxwu@users.noreply.github.com> --- docs/features/kv_offloading_usage.md | 7 ++ .../v1/simple_kv_offload/test_integration.py | 89 ++++++++++++++++++- tests/v1/simple_kv_offload/test_scheduler.py | 35 ++++++++ vllm/v1/simple_kv_offload/manager.py | 10 +++ 4 files changed, 140 insertions(+), 1 deletion(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 60d5fa72f99c..804044b802de 100644 --- a/docs/features/kv_offloading_usage.md +++ b/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). diff --git a/tests/v1/simple_kv_offload/test_integration.py b/tests/v1/simple_kv_offload/test_integration.py index ea6caa421565..c9119b7bff23 100644 --- a/tests/v1/simple_kv_offload/test_integration.py +++ b/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", diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 4eff332ce056..d56ab77797f9 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/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) # --------------------------------------------------------------------------- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 5d23d6cab1ff..62a022d9e0f2 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/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:] From 75e5d3ba281239737f7ca7b7867731d72f574cd2 Mon Sep 17 00:00:00 2001 From: CharlesWu <296364977+charxwu@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:41:02 +0800 Subject: [PATCH 2/2] ci: retrigger documentation build after precheck approval Trigger the PR documentation webhook now that the pre-run gate passes. No source or documentation content changes. Co-authored-by: Codex Signed-off-by: CharlesWu <296364977+charxwu@users.noreply.github.com>