[Bugfix][KV Offload] Respect prefix-cache bypass in SimpleCPUOffload (#54998)
Signed-off-by: Andy Lo <andy@mistral.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Andy Lo committed
Sep 8, 2026 at 14:14 UTC
5ebce2391e772994a92f903bdb97b4631a05e851
2 files changed
+56
tests/v1/simple_kv_offload/test_scheduler.py
+53
@@ -430,6 +430,59 @@ def test_eager_store_and_load_roundtrip() -> None:
430
assert len(meta2.load_cpu_blocks) == len(meta2.load_gpu_blocks)
431
432
433
+def test_prompt_logprobs_skip_cpu_cache_lookup() -> None:
434
+ """Prompt logprobs require every prompt token to be recomputed."""
435
+ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
436
+ sched = fix.scheduler
437
+
438
+ num_blocks = 2
439
+ req = make_request(num_blocks=num_blocks)
440
+ kv_blocks = _alloc_and_register(fix, req, num_blocks)
441
+ sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
442
+ sched_out = make_scheduler_output(
443
+ {req.request_id: num_blocks * BLOCK_SIZE},
444
+ new_reqs={req.request_id: kv_blocks.get_block_ids()},
445
+ )
446
+ meta = sched.build_connector_meta(sched_out)
447
+ simulate_store_completion(sched, meta.store_event)
448
+
449
+ retry_request_id = "req-prompt-logprobs"
450
+ cached_request = Request(
451
+ request_id=retry_request_id,
452
+ prompt_token_ids=req.prompt_token_ids,
453
+ sampling_params=req.sampling_params,
454
+ pooling_params=None,
455
+ mm_features=None,
456
+ block_hasher=req._block_hasher,
457
+ )
458
+ assert sched.get_num_new_matched_tokens(cached_request, num_computed_tokens=0) == (
459
+ num_blocks * BLOCK_SIZE,
460
+ True,
461
+ )
462
+ pending_hit = sched._pending_cpu_hits[retry_request_id]
463
+ pinned_blocks = [
464
+ block for group in pending_hit[0] for block in group if not block.is_null
465
+ ]
466
+ assert pinned_blocks
467
+ assert all(block.ref_cnt == 1 for block in pinned_blocks)
468
+
469
+ prompt_logprobs_request = Request(
470
+ request_id=retry_request_id,
471
+ prompt_token_ids=req.prompt_token_ids,
472
+ sampling_params=SamplingParams(max_tokens=1, prompt_logprobs=1),
473
+ pooling_params=None,
474
+ mm_features=None,
475
+ block_hasher=req._block_hasher,
476
+ )
477
+ assert prompt_logprobs_request.skip_reading_prefix_cache
478
+
479
+ assert sched.get_num_new_matched_tokens(
480
+ prompt_logprobs_request, num_computed_tokens=0
481
+ ) == (0, False)
482
+ assert retry_request_id not in sched._pending_cpu_hits
483
+ assert all(block.ref_cnt == 0 for block in pinned_blocks)
484
+
485
+
486
def test_eager_store_preserves_secondary_block_hashes() -> None:
487
"""CPU copies retain fine-grained hashes owned by the GPU block."""
488
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
vllm/v1/simple_kv_offload/manager.py
+3
@@ -293,6 +293,9 @@ class SimpleCPUOffloadScheduler:
293
if stale := self._pending_cpu_hits.pop(request.request_id, None):
294
self._free_pending_cpu_hit(stale)
295
296
+ if request.skip_reading_prefix_cache:
297
+ return 0, False
298
+
299
num_skipped_hashes = num_computed_tokens // self.hash_block_size
300
remaining_hashes = request.block_hashes[num_skipped_hashes:]
301