A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 15:00 UTC
[KV Offload] Offload all KV blocks when doing prefill in P/D (#40346)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Signed-off-by: omerpaz95 <73347585+omerpaz95@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com>
omerpaz95 committed
Apr 26, 2026 at 15:06 UTC
b39c266dae8cd7aee31f667c973e9698ed0b2361
3 files changed
+71
-2
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+56
@@ -478,3 +478,59 @@ class TestSlidingWindowLookup:
478
sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX)
479
is None
480
)
481
+
482
+
483
+@pytest.mark.parametrize("async_scheduling", [True, False])
484
+def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool):
485
+ """With do_remote_decode=True, after loading prefix blocks from CPU,
486
+ all blocks must be re-stored — not just the newly computed ones.
487
+
488
+ This supports P/D disaggregation where the prefill instance offloads the
489
+ complete KV cache so a remote decode node can consume it."""
490
+ offloaded_block_size = 12
491
+ gpu_block_size = 4
492
+ num_gpu_blocks = 100
493
+
494
+ runner = request_runner(
495
+ offloaded_block_size=offloaded_block_size,
496
+ gpu_block_size=gpu_block_size,
497
+ num_gpu_blocks=num_gpu_blocks,
498
+ async_scheduling=async_scheduling,
499
+ )
500
+
501
+ # Store 1 offloaded block (3 GPU blocks) via a normal request.
502
+ runner.new_request(token_ids=[0] * offloaded_block_size)
503
+ runner.manager.prepare_store.side_effect = (
504
+ lambda keys, req_context: generate_store_output(keys)
505
+ )
506
+ runner.run(
507
+ decoded_tokens=[EOS_TOKEN_ID],
508
+ expected_stored_gpu_block_indexes=(0, 1, 2),
509
+ )
510
+
511
+ # Reset GPU prefix cache so the next request must load from CPU.
512
+ runner.scheduler.reset_prefix_cache()
513
+
514
+ # New request with do_remote_decode=True and 2 offloaded blocks.
515
+ # The first offloaded block matches what we stored in CPU.
516
+ runner.new_request(
517
+ token_ids=[0] * offloaded_block_size * 2,
518
+ kv_transfer_params={"do_remote_decode": True},
519
+ )
520
+ runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
521
+ runner.manager.prepare_store.side_effect = (
522
+ lambda keys, req_context: generate_store_output(keys)
523
+ )
524
+
525
+ # Load the first offloaded block from CPU.
526
+ runner.run(
527
+ decoded_tokens=[0],
528
+ expected_loaded_gpu_block_indexes=(0, 1, 2),
529
+ )
530
+
531
+ # Store must include ALL 6 GPU blocks (both the loaded prefix and
532
+ # the newly computed block), not just the 3 new ones.
533
+ runner.run(
534
+ decoded_tokens=[EOS_TOKEN_ID],
535
+ expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
536
+ )
tests/v1/kv_connector/unit/offloading_connector/utils.py
+7
-1
@@ -270,7 +270,11 @@ class RequestRunner:
270
slot_mapping={},
271
)
272
273
- def new_request(self, token_ids: list[int]):
273
+ def new_request(
274
+ self,
275
+ token_ids: list[int],
276
+ kv_transfer_params: dict | None = None,
277
+ ):
278
self.req_id += 1
279
280
sampling_params = SamplingParams(max_tokens=1000)
@@ -283,6 +287,8 @@ class RequestRunner:
287
pooling_params=None,
288
block_hasher=self._block_hasher,
289
)
290
+ if kv_transfer_params is not None:
291
+ req.kv_transfer_params = kv_transfer_params
292
293
self.scheduler.add_request(req)
294
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+8
-1
@@ -314,6 +314,9 @@ class OffloadingConnectorScheduler:
314
num_locally_computed_tokens = req_status.num_locally_computed_tokens
315
num_cached_tokens = num_locally_computed_tokens + num_external_tokens
316
317
+ params = req_status.req_context.kv_transfer_params
318
+ do_remote_decode = params is not None and params.get("do_remote_decode")
319
+
320
keys_to_load: list[OffloadKey] = []
321
dst_block_ids: list[int] = []
322
# per group
@@ -360,7 +363,11 @@ class OffloadingConnectorScheduler:
363
group_sizes.append(num_pending_gpu_blocks)
364
block_indices.append(num_locally_computed_gpu_blocks)
365
363
- group_state.next_stored_block_idx = num_blocks
366
+ if not do_remote_decode:
367
+ # For P/D prefill requests (do_remote_decode=True), we do
368
+ # NOT skip saving the hit prefix, as we need to stream the
369
+ # entire KV cache so a remote decode node can consume it.
370
+ group_state.next_stored_block_idx = num_blocks
371
372
src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
373
dst_spec = GPULoadStoreSpec(