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 10, 09:00 UTC
[BugFix] Wait for compute before offloading KV to CPU (#31341)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Jan 11, 2026 at 00:25 UTC
2a4dbe24eadcb8e0354e47f608b53399aec52c43
3 files changed
+44
-30
tests/v1/kv_connector/unit/test_offloading_connector.py
+19
-13
@@ -172,6 +172,7 @@ class RequestRunner:
172
173
self.pending_loads_count: int = 0
174
self.pending_stores_count: int = 0
175
+ self.unsubmitted_stores_count = 0
176
177
self.completed_loads: list[TransferSummary] = []
178
self.completed_stores: list[TransferSummary] = []
@@ -279,7 +280,9 @@ class RequestRunner:
280
assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)
281
282
self.pending_loads_count += len(kv_connector_metadata.reqs_to_load)
282
- self.pending_stores_count += len(kv_connector_metadata.reqs_to_store)
283
+
284
+ self.pending_stores_count += self.unsubmitted_stores_count
285
+ self.unsubmitted_stores_count = len(kv_connector_metadata.reqs_to_store)
286
287
self.worker_connector.bind_connector_metadata(kv_connector_metadata)
288
self.worker_connector.start_load_kv(self._dummy_ctx)
@@ -414,10 +417,13 @@ def test_offloading_connector(request_runner):
417
runner.manager.prepare_store.side_effect = (
418
lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
419
)
417
- runner.run(decoded_tokens=[0], expected_stored_gpu_block_indexes=(3, 4, 5))
420
+ runner.run(decoded_tokens=[0])
421
422
# add block missing 1 token -> no offload
420
- runner.run(decoded_tokens=[0] * (offloaded_block_size - 1))
423
+ runner.run(
424
+ decoded_tokens=[0] * (offloaded_block_size - 1),
425
+ expected_stored_gpu_block_indexes=(3, 4, 5),
426
+ )
427
runner.manager.prepare_store.assert_not_called()
428
429
# +1 token -> single block, fail prepare_store
@@ -435,23 +441,20 @@ def test_offloading_connector(request_runner):
441
runner.manager.prepare_store.side_effect = (
442
lambda block_hashes: generate_store_output(block_hashes)
443
)
438
- runner.run(
439
- decoded_tokens=[0] * offloaded_block_size,
440
- expected_stored_gpu_block_indexes=(15, 16, 17),
441
- )
444
+ runner.run(decoded_tokens=[0] * offloaded_block_size)
445
runner.manager.touch.assert_called()
446
block_hashes1 = list(runner.manager.touch.call_args.args[0])
447
assert len(block_hashes1) == 6
448
449
# terminate request
447
- runner.run(decoded_tokens=[EOS_TOKEN_ID])
450
+ runner.run(
451
+ decoded_tokens=[EOS_TOKEN_ID],
452
+ expected_stored_gpu_block_indexes=(15, 16, 17),
453
+ )
454
455
# create a new request differing only on the last token
456
runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
451
- runner.run(
452
- decoded_tokens=[0],
453
- expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
454
- )
457
+ runner.run(decoded_tokens=[0])
458
runner.manager.touch.assert_called()
459
block_hashes2 = list(runner.manager.touch.call_args.args[0])
460
assert len(block_hashes2) == 6
@@ -461,7 +464,10 @@ def test_offloading_connector(request_runner):
464
assert block_hashes1[5] != block_hashes2[5]
465
466
# terminate request
464
- runner.run(decoded_tokens=[EOS_TOKEN_ID])
467
+ runner.run(
468
+ decoded_tokens=[EOS_TOKEN_ID],
469
+ expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
470
+ )
471
472
# full_block_tokens - num_computed_tokens < offloaded_block_size
473
runner.new_request(
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+17
-6
@@ -78,7 +78,7 @@ class OffloadingConnector(KVConnectorBase_V1):
78
def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
79
assert self.connector_worker is not None
80
assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
81
- self.connector_worker.start_load_kv(self._connector_metadata)
81
+ self.connector_worker.start_kv_transfers(self._connector_metadata)
82
83
def wait_for_layer_load(self, layer_name: str) -> None:
84
pass
@@ -95,7 +95,7 @@ class OffloadingConnector(KVConnectorBase_V1):
95
def wait_for_save(self):
96
assert self.connector_worker is not None
97
assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
98
- self.connector_worker.start_store_kv(self._connector_metadata)
98
+ self.connector_worker.prepare_store_kv(self._connector_metadata)
99
100
def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
101
assert self.connector_worker is not None
@@ -427,6 +427,8 @@ class OffloadingConnectorWorker:
427
self._load_job: dict[ReqId, int] = {}
428
# req_id -> set(active job IDs)
429
self._store_jobs = defaultdict[ReqId, set[int]](set)
430
+ # list of store jobs pending submission (job_id, transfer_spec)
431
+ self._unsubmitted_store_jobs: list[tuple[int, TransferSpec]] = []
432
433
self._finished_reqs_waiting_for_store: set[ReqId] = set()
434
@@ -464,20 +466,29 @@ class OffloadingConnectorWorker:
466
attn_backends = {cross_layer_name: attn_backend}
467
self._register_handlers(kv_caches, attn_backends)
468
467
- def start_load_kv(self, metadata: OffloadingConnectorMetadata):
469
+ def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):
470
+ for job_id, transfer_spec in self._unsubmitted_store_jobs:
471
+ success = self.worker.transfer_async(job_id, transfer_spec)
472
+ assert success
473
+ self._unsubmitted_store_jobs.clear()
474
+
475
for req_id, transfer_spec in metadata.reqs_to_load.items():
476
job_id = self._generate_job_id()
477
self._jobs[job_id] = (req_id, False)
478
assert req_id not in self._load_job
479
self._load_job[req_id] = job_id
473
- assert self.worker.transfer_async(job_id, transfer_spec)
480
+ success = self.worker.transfer_async(job_id, transfer_spec)
481
+ assert success
482
475
- def start_store_kv(self, metadata: OffloadingConnectorMetadata):
483
+ def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
484
for req_id, transfer_spec in metadata.reqs_to_store.items():
485
job_id = self._generate_job_id()
486
self._jobs[job_id] = (req_id, True)
487
self._store_jobs[req_id].add(job_id)
480
- assert self.worker.transfer_async(job_id, transfer_spec)
488
+ # NOTE(orozery): defer the store to the beginning of the next engine step,
489
+ # so that offloading starts AFTER transfers related to token sampling,
490
+ # thereby avoiding delays to token generation due to offloading.
491
+ self._unsubmitted_store_jobs.append((job_id, transfer_spec))
492
493
def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
494
"""
vllm/v1/kv_offload/worker/cpu_gpu.py
+8
-11
@@ -68,7 +68,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
68
kv_dim_before_num_blocks: list[bool],
69
src_block_size_factor: int,
70
dst_block_size_factor: int,
71
- priority: int,
71
):
72
"""
73
Initialize a SingleDirectionOffloadingHandler.
@@ -85,8 +84,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
84
per KV block in a source tensor.
85
dst_block_size_factor: The number of kernel blocks
86
per KV block in a destination tensor.
88
- priority: The priority of the backing CUDA streams.
89
- Lower numbers indicate higher priority.
87
"""
88
assert len(src_tensors) == len(dst_tensors) == len(kv_dim_before_num_blocks)
89
@@ -95,7 +92,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
92
self.kv_dim_before_num_blocks: list[bool] = kv_dim_before_num_blocks
93
self.src_block_size_factor: int = src_block_size_factor
94
self.dst_block_size_factor: int = dst_block_size_factor
98
- self.priority = priority
95
+
96
+ assert len(src_tensors) > 0
97
+ self.gpu_to_cpu: bool = self.src_tensors[0].is_cuda
98
99
# queue of transfers (job_id, stream, event)
100
self._transfers: deque[tuple[int, torch.cuda.Stream, torch.Event]] = deque()
@@ -130,12 +129,12 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
129
expand_block_ids(dst_blocks, self.dst_block_size_factor, src_to_dst[:, 1])
130
src_to_dst_tensor = torch.from_numpy(src_to_dst)
131
133
- stream = (
134
- self._stream_pool.pop()
135
- if self._stream_pool
136
- else torch.cuda.Stream(priority=self.priority)
137
- )
132
+ stream = self._stream_pool.pop() if self._stream_pool else torch.cuda.Stream()
133
event = self._event_pool.pop() if self._event_pool else torch.Event()
134
+
135
+ if self.gpu_to_cpu:
136
+ # wait for model computation to finish before offloading
137
+ stream.wait_stream(torch.cuda.current_stream())
138
if self._transfers:
139
_, _, last_event = self._transfers[-1]
140
# assure job will start only after the previous one completes
@@ -267,7 +266,6 @@ class CpuGpuOffloadingHandlers:
266
kv_dim_before_num_blocks=kv_dim_before_num_blocks,
267
src_block_size_factor=gpu_block_size_factor,
268
dst_block_size_factor=cpu_block_size_factor,
270
- priority=1,
269
)
270
271
self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler(
@@ -276,5 +274,4 @@ class CpuGpuOffloadingHandlers:
274
kv_dim_before_num_blocks=kv_dim_before_num_blocks,
275
src_block_size_factor=cpu_block_size_factor,
276
dst_block_size_factor=gpu_block_size_factor,
279
- priority=-1,
277
)