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, 13:45 UTC
[KV Offload] Per-job store completion for CPU offloading connector (#39186)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Itay Etelis committed
Apr 29, 2026 at 08:52 UTC
8a8c9b564ef015c76cf398200b8f0891e6e51bb8
8 files changed
+430
-145
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+134
@@ -232,6 +232,9 @@ def test_request_preemption(request_runner, async_scheduling: bool):
232
expected_stored_gpu_block_indexes=(9, 10, 11),
233
)
234
235
+ # All stores completed before request_finished -> fence index empty.
236
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
237
+
238
239
@pytest.mark.parametrize("async_scheduling", [True, False])
240
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
@@ -292,6 +295,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
295
# second request will use the GPU prefix cache
296
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
297
298
+ # Fence index drained: stores completed before request_finished ran.
299
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
300
+
301
302
@pytest.mark.parametrize("async_scheduling", [True, False])
303
def test_abort_loading_requests(request_runner, async_scheduling: bool):
@@ -534,3 +540,131 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
540
decoded_tokens=[EOS_TOKEN_ID],
541
expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
542
)
543
+ # All stores completed before request_finished -> fence index empty.
544
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
545
+
546
+
547
+# ---------------------------------------------------------------------------
548
+# Tests for the per-job-store-completion design and fence invariants.
549
+# ---------------------------------------------------------------------------
550
+
551
+
552
+def test_loads_do_not_populate_fence_index(request_runner):
553
+ """Loads don't populate _block_id_to_pending_jobs (protected by
554
+ delay_free_blocks while in flight)."""
555
+ runner = request_runner(
556
+ offloaded_block_size=12,
557
+ gpu_block_size=4,
558
+ num_gpu_blocks=100,
559
+ async_scheduling=False,
560
+ )
561
+ runner.new_request(token_ids=[0] * 12)
562
+ runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
563
+ runner.run(decoded_tokens=[], complete_transfers=False)
564
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
565
+
566
+
567
+def test_fence_at_update_state_after_alloc(request_runner):
568
+ """A load reusing a finished request's pending-store block triggers
569
+ a flush via update_state_after_alloc's fence.
570
+
571
+ num_gpu_blocks=2 forces the BlockPool to give req2 the same block
572
+ req1 just freed.
573
+ """
574
+ runner = request_runner(
575
+ offloaded_block_size=4,
576
+ gpu_block_size=4,
577
+ num_gpu_blocks=2,
578
+ async_scheduling=False,
579
+ )
580
+
581
+ runner.new_request(token_ids=[0] * 4)
582
+ runner.manager.prepare_store.side_effect = (
583
+ lambda keys, req_context: generate_store_output(keys)
584
+ )
585
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
586
+ assert runner.connector_scheduler._block_id_to_pending_jobs
587
+
588
+ runner.scheduler.reset_prefix_cache()
589
+ runner.new_request(token_ids=[0] * 4)
590
+ runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
591
+ runner.manager.prepare_store.side_effect = (
592
+ lambda keys, req_context: generate_store_output([])
593
+ )
594
+ runner.run(
595
+ decoded_tokens=[],
596
+ complete_transfers=False,
597
+ expected_stored_gpu_block_indexes=(0,),
598
+ expected_flushed_gpu_block_indexes=(0,),
599
+ )
600
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
601
+
602
+
603
+def test_fence_at_build_store_jobs(request_runner):
604
+ """A new prefill (no load -> update_state_after_alloc returns early)
605
+ reusing a finished request's pending-store block is flushed by
606
+ _build_store_jobs's fence."""
607
+ runner = request_runner(
608
+ offloaded_block_size=4,
609
+ gpu_block_size=4,
610
+ num_gpu_blocks=2,
611
+ async_scheduling=False,
612
+ )
613
+
614
+ runner.new_request(token_ids=[0] * 4)
615
+ runner.manager.prepare_store.side_effect = (
616
+ lambda keys, req_context: generate_store_output(keys)
617
+ )
618
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
619
+ assert runner.connector_scheduler._block_id_to_pending_jobs
620
+
621
+ runner.scheduler.reset_prefix_cache()
622
+ runner.new_request(token_ids=[1] * 4)
623
+ runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
624
+ runner.manager.prepare_store.side_effect = (
625
+ lambda keys, req_context: generate_store_output([])
626
+ )
627
+ runner.run(
628
+ decoded_tokens=[EOS_TOKEN_ID],
629
+ expected_stored_gpu_block_indexes=(0,),
630
+ expected_flushed_gpu_block_indexes=(0,),
631
+ )
632
+ assert runner.connector_scheduler._block_id_to_pending_jobs == {}
633
+
634
+
635
+@pytest.mark.parametrize("async_scheduling", [True, False])
636
+def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
637
+ """complete_store fires per-job, not deferred to request finish.
638
+ Each call carries only that store's keys."""
639
+ offloaded_block_size = 12
640
+ runner = request_runner(
641
+ offloaded_block_size=offloaded_block_size,
642
+ gpu_block_size=4,
643
+ num_gpu_blocks=100,
644
+ async_scheduling=async_scheduling,
645
+ )
646
+ runner.new_request(token_ids=[0] * offloaded_block_size)
647
+ runner.manager.prepare_store.side_effect = (
648
+ lambda keys, req_context: generate_store_output(keys)
649
+ )
650
+
651
+ # First store: fires when block 0 is fully populated.
652
+ runner.run(decoded_tokens=[0, 0], expected_stored_gpu_block_indexes=(0, 1, 2))
653
+ assert runner.manager.complete_store.call_count == 1
654
+ first_call_keys = set(runner.manager.complete_store.call_args.args[0])
655
+ assert len(first_call_keys) == 1
656
+ runner.manager.complete_store.reset_mock()
657
+
658
+ # Second store: fires when block 1 is fully populated, with different keys.
659
+ runner.run(
660
+ decoded_tokens=[0] * (offloaded_block_size + 1),
661
+ expected_stored_gpu_block_indexes=(3, 4, 5),
662
+ )
663
+ assert runner.manager.complete_store.call_count == 1
664
+ second_call_keys = set(runner.manager.complete_store.call_args.args[0])
665
+ assert first_call_keys != second_call_keys
666
+ runner.manager.complete_store.reset_mock()
667
+
668
+ # Finish: no store pending -> no further call.
669
+ runner.run(decoded_tokens=[EOS_TOKEN_ID])
670
+ assert runner.manager.complete_store.call_count == 0
tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py
new
+32
@@ -0,0 +1,32 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
7
+ OffloadingWorkerMetadata,
8
+)
9
+
10
+pytestmark = pytest.mark.cpu_test
11
+
12
+
13
+def test_aggregate_sums_counts():
14
+ meta1 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1})
15
+ meta2 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1})
16
+ result = meta1.aggregate(meta2)
17
+ assert result.completed_jobs == {42: 2, 7: 2}
18
+
19
+
20
+def test_aggregate_disjoint_jobs():
21
+ meta1 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1})
22
+ meta2 = OffloadingWorkerMetadata(completed_jobs={43: 1, 8: 1})
23
+ result = meta1.aggregate(meta2)
24
+ assert result.completed_jobs == {42: 1, 7: 1, 43: 1, 8: 1}
25
+
26
+
27
+def test_aggregate_multiple_workers():
28
+ meta1 = OffloadingWorkerMetadata(completed_jobs={42: 1, 43: 1, 7: 1})
29
+ meta2 = OffloadingWorkerMetadata(completed_jobs={42: 1, 7: 1, 8: 1})
30
+ meta3 = OffloadingWorkerMetadata(completed_jobs={42: 1, 43: 1, 8: 1})
31
+ result = meta1.aggregate(meta2).aggregate(meta3)
32
+ assert result.completed_jobs == {42: 3, 43: 2, 7: 2, 8: 2}
tests/v1/kv_connector/unit/offloading_connector/utils.py
+13
-21
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import copy
3
from collections.abc import Iterable, Iterator
4
from dataclasses import dataclass
5
from typing import Any
@@ -19,6 +18,7 @@ from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
18
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
19
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
20
OffloadingConnectorMetadata,
21
+ OffloadingWorkerMetadata,
22
)
23
from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
24
OffloadingConnector,
@@ -51,7 +51,6 @@ from vllm.v1.kv_offload.worker.worker import (
51
TransferResult,
52
TransferSpec,
53
)
54
-from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
54
from vllm.v1.request import Request
55
from vllm.v1.structured_output import StructuredOutputManager
56
@@ -369,7 +368,12 @@ class RequestRunner:
368
prev_scheduler_output = None
369
prev_model_runner_output = None
370
while True:
372
- assert self.scheduler.requests
371
+ # Strict-always-False frees the request immediately on EOS, but
372
+ # the worker may still have a deferred store queued. In production
373
+ # the next request's step drains it; in single-request tests we
374
+ # must keep stepping until the scheduler sees no in-flight jobs.
375
+ if not self.scheduler.requests and not self.connector_scheduler._jobs:
376
+ break
377
378
scheduler_output = self.scheduler.schedule()
379
self._update_gpu_block_idx()
@@ -392,6 +396,10 @@ class RequestRunner:
396
finished_sending, finished_recving = self.worker_connector.get_finished(
397
scheduler_output.finished_req_ids
398
)
399
+ worker_meta = (
400
+ self.worker_connector.build_connector_worker_meta()
401
+ or OffloadingWorkerMetadata()
402
+ )
403
404
self.worker_connector.clear_connector_metadata()
405
@@ -400,6 +408,7 @@ class RequestRunner:
408
finished_sending=finished_sending,
409
finished_recving=finished_recving,
410
token_id=token_id or 0,
411
+ kv_connector_worker_meta=worker_meta,
412
)
413
414
prev_token_id = token_id
@@ -420,7 +429,7 @@ class RequestRunner:
429
if (
430
prev_token_id == EOS_TOKEN_ID
431
and prev_token_id != token_id
423
- and self.scheduler.requests
432
+ and (self.scheduler.requests or self.connector_scheduler._jobs)
433
):
434
# continue for one more step to allow offloading to kick off
435
continue
@@ -435,26 +444,9 @@ class RequestRunner:
444
445
self._parse_transfers()
446
438
- # run one more step to update finished stored
447
if EOS_TOKEN_ID in decoded_tokens:
448
assert not self.scheduler.running
449
442
- while self.scheduler.requests:
443
- scheduler_output = self.scheduler.schedule()
444
-
445
- finished_sending, finished_recving = self.worker_connector.get_finished(
446
- scheduler_output.finished_req_ids
447
- )
448
-
449
- assert not finished_recving
450
-
451
- model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
452
- model_runner_output.kv_connector_output = KVConnectorOutput(
453
- finished_sending=finished_sending
454
- )
455
-
456
- self.scheduler.update_from_output(scheduler_output, model_runner_output)
457
-
450
def run(
451
self,
452
decoded_tokens: list[int],
tests/v1/kv_connector/unit/utils.py
+4
@@ -24,6 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
24
KVConnectorBase_V1,
25
KVConnectorMetadata,
26
KVConnectorRole,
27
+ KVConnectorWorkerMetadata,
28
)
29
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa
30
ExampleConnector,
@@ -249,6 +250,7 @@ def create_model_runner_output(
250
invalid_block_ids: set[int] | None = None,
251
use_eos: bool = False,
252
token_id: int = 0,
253
+ kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None,
254
) -> ModelRunnerOutput:
255
"""Make dummy model runner output for testing."""
256
@@ -266,11 +268,13 @@ def create_model_runner_output(
268
finished_sending is None
269
and finished_recving is None
270
and invalid_block_ids is None
271
+ and kv_connector_worker_meta is None
272
)
273
else KVConnectorOutput(
274
finished_sending=finished_sending,
275
finished_recving=finished_recving,
276
invalid_block_ids=invalid_block_ids or set(),
277
+ kv_connector_worker_meta=kv_connector_worker_meta,
278
)
279
)
280
vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py
+50
-5
@@ -1,15 +1,60 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
3
+from dataclasses import dataclass, field
4
5
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
5
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
6
+ KVConnectorMetadata,
7
+ KVConnectorWorkerMetadata,
8
+)
9
from vllm.v1.kv_offload.worker.worker import TransferSpec
10
11
ReqId = str
12
13
14
+@dataclass
15
+class TransferJob:
16
+ """A transfer job bundling request context with transfer spec.
17
+
18
+ Used for both loads and stores, keyed by scheduler-assigned job ID.
19
+ The worker reports the job ID back when the transfer finishes,
20
+ and the scheduler processes the completion.
21
+ """
22
+
23
+ req_id: ReqId
24
+ transfer_spec: TransferSpec
25
+
26
+
27
@dataclass
28
class OffloadingConnectorMetadata(KVConnectorMetadata):
13
- reqs_to_load: dict[ReqId, TransferSpec]
14
- reqs_to_store: dict[ReqId, TransferSpec]
15
- reqs_to_flush: set[str] | None = None
29
+ # Keyed by scheduler-assigned job IDs.
30
+ load_jobs: dict[int, TransferJob]
31
+ store_jobs: dict[int, TransferJob]
32
+ jobs_to_flush: set[int] | None = None
33
+
34
+
35
+@dataclass
36
+class OffloadingWorkerMetadata(KVConnectorWorkerMetadata):
37
+ """Worker -> Scheduler metadata for completed transfer jobs.
38
+
39
+ Each worker reports {job_id: 1} for newly completed transfer jobs
40
+ (load or store). aggregate() sums counts across workers within a step.
41
+ The scheduler accumulates across steps and processes
42
+ a transfer completion only when count reaches num_workers.
43
+ """
44
+
45
+ completed_jobs: dict[int, int] = field(default_factory=dict)
46
+
47
+ def mark_completed(self, job_id: int) -> None:
48
+ """Record a transfer job completion from this worker."""
49
+ self.completed_jobs[job_id] = 1
50
+
51
+ def aggregate(
52
+ self, other: "KVConnectorWorkerMetadata"
53
+ ) -> "KVConnectorWorkerMetadata":
54
+ assert isinstance(other, OffloadingWorkerMetadata)
55
+
56
+ merged = dict(self.completed_jobs)
57
+ for job_id, v in other.completed_jobs.items():
58
+ merged[job_id] = merged.get(job_id, 0) + v
59
+
60
+ return OffloadingWorkerMetadata(completed_jobs=merged)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+156
-48
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections import defaultdict
3
from collections.abc import Iterable, Sequence
4
from dataclasses import dataclass, field
5
from itertools import islice
@@ -11,7 +10,9 @@ from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
10
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
11
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
12
OffloadingConnectorMetadata,
13
+ OffloadingWorkerMetadata,
14
ReqId,
15
+ TransferJob,
16
)
17
from vllm.logger import init_logger
18
from vllm.utils.math_utils import cdiv
@@ -26,13 +27,27 @@ from vllm.v1.kv_offload.abstract import (
27
)
28
from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
29
from vllm.v1.kv_offload.spec import OffloadingSpec
29
-from vllm.v1.kv_offload.worker.worker import TransferSpec
30
from vllm.v1.outputs import KVConnectorOutput
31
from vllm.v1.request import Request
32
33
logger = init_logger(__name__)
34
35
36
+@dataclass(slots=True)
37
+class TransferJobStatus:
38
+ """Tracks scheduler-side state for a single transfer job."""
39
+
40
+ req_id: ReqId
41
+ # Number of workers still pending. Starts at num_workers,
42
+ # decremented as each worker reports completion. Job is done at 0.
43
+ pending_count: int
44
+ # Offload keys this job covers; passed to manager.complete_*().
45
+ keys: set[OffloadKey]
46
+ is_store: bool
47
+ # GPU blocks the fence tracks. Store src blocks; None for loads.
48
+ gpu_block_ids: list[int] | None = None
49
+
50
+
51
class GroupOffloadConfig(NamedTuple):
52
group_idx: int
53
gpu_block_size: int
@@ -43,10 +58,12 @@ class GroupOffloadConfig(NamedTuple):
58
class SchedulerOffloadConfig(NamedTuple):
59
kv_group_configs: tuple[GroupOffloadConfig, ...]
60
block_size_factor: int
61
+ num_workers: int
62
63
@classmethod
64
def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
65
return cls(
66
+ num_workers=spec.vllm_config.parallel_config.world_size,
67
kv_group_configs=tuple(
68
GroupOffloadConfig(
69
group_idx=idx,
@@ -79,6 +96,9 @@ class RequestOffloadState:
96
req_context: ReqContext = field(init=False)
97
# number of hits in the GPU cache
98
num_locally_computed_tokens: int = 0
99
+ # In-flight job IDs. Per the connector's invariant, at any given time
100
+ # this contains either a single load job, or one or more store jobs.
101
+ transfer_jobs: set[int] = field(default_factory=set)
102
103
def __post_init__(self) -> None:
104
self.group_states = tuple(
@@ -135,17 +155,26 @@ class OffloadingConnectorScheduler:
155
self.lookup_groups = attention_groups
156
157
self._req_status: dict[ReqId, RequestOffloadState] = {}
138
- # requests to load for the current scheduler step
139
- self._reqs_to_load: dict[ReqId, TransferSpec] = {}
158
+ self._current_batch_load_jobs: dict[int, TransferJob] = {}
159
+ self._current_batch_jobs_to_flush: set[int] = set()
160
# if GPU prefix caching is enabled,
161
# track loaded blocks to avoid redundant loads
162
self._blocks_being_loaded: set[OffloadKey] | None = (
163
set() if spec.vllm_config.cache_config.enable_prefix_caching else None
164
)
165
146
- # request ID -> set(offload keys being stored/loaded)
147
- self._reqs_being_stored = defaultdict[ReqId, set[OffloadKey]](set)
148
- self._reqs_being_loaded = defaultdict[ReqId, set[OffloadKey]](set)
166
+ # Job ID counter shared by loads and stores.
167
+ self._job_counter: int = 0
168
+ self._jobs: dict[int, TransferJobStatus] = {}
169
+
170
+ # block_id -> pending store job_ids. Populated only for finished
171
+ # requests (running-request blocks are protected by their ref_cnt).
172
+ self._block_id_to_pending_jobs: dict[int, set[int]] = {}
173
+
174
+ def _generate_job_id(self) -> int:
175
+ job_id = self._job_counter
176
+ self._job_counter += 1
177
+ return job_id
178
179
def _maximal_prefix_lookup(
180
self, keys: Iterable[OffloadKey], req_context: ReqContext
@@ -369,23 +398,46 @@ class OffloadingConnectorScheduler:
398
# entire KV cache so a remote decode node can consume it.
399
group_state.next_stored_block_idx = num_blocks
400
401
+ # Fence dst blocks against finished-request pending stores.
402
+ if (
403
+ self._block_id_to_pending_jobs
404
+ and not self._block_id_to_pending_jobs.keys().isdisjoint(dst_block_ids)
405
+ ):
406
+ self._current_batch_jobs_to_flush.update(
407
+ jid
408
+ for bid in dst_block_ids
409
+ for jid in self._block_id_to_pending_jobs.get(bid, ())
410
+ )
411
+
412
src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
413
dst_spec = GPULoadStoreSpec(
414
dst_block_ids, group_sizes=group_sizes, block_indices=block_indices
415
)
416
377
- self._reqs_to_load[request.request_id] = (src_spec, dst_spec)
378
- req_blocks_being_loaded = self._reqs_being_loaded[request.request_id]
379
- req_blocks_being_loaded.update(keys_to_load)
417
+ load_job_id = self._generate_job_id()
418
+ self._current_batch_load_jobs[load_job_id] = TransferJob(
419
+ req_id=request.request_id,
420
+ transfer_spec=(src_spec, dst_spec),
421
+ )
422
+ # a load can only be issued when no other jobs are pending.
423
+ assert not req_status.transfer_jobs
424
+ req_status.transfer_jobs.add(load_job_id)
425
+ self._jobs[load_job_id] = TransferJobStatus(
426
+ req_id=request.request_id,
427
+ pending_count=self.config.num_workers,
428
+ keys=set(keys_to_load),
429
+ is_store=False,
430
+ )
431
432
if self._blocks_being_loaded is not None:
382
- self._blocks_being_loaded.update(req_blocks_being_loaded)
433
+ self._blocks_being_loaded.update(keys_to_load)
434
384
- def _get_reqs_to_store(
385
- self, scheduler_output: SchedulerOutput
386
- ) -> dict[ReqId, TransferSpec]:
435
+ def _build_store_jobs(
436
+ self,
437
+ scheduler_output: SchedulerOutput,
438
+ ) -> dict[int, TransferJob]:
439
block_size_factor = self.config.block_size_factor
388
- reqs_to_store: dict[ReqId, TransferSpec] = {}
440
+ store_jobs: dict[int, TransferJob] = {}
441
# iterate over both new and cached requests
442
for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
443
req_status = self._req_status[req_id]
@@ -398,6 +450,19 @@ class OffloadingConnectorScheduler:
450
451
if new_block_id_groups:
452
req_status.update_block_id_groups(new_block_id_groups)
453
+ # Fence new blocks against in-flight stores.
454
+ if self._block_id_to_pending_jobs:
455
+ new_blocks_flat = [
456
+ bid for new_blocks in new_block_id_groups for bid in new_blocks
457
+ ]
458
+ if not self._block_id_to_pending_jobs.keys().isdisjoint(
459
+ new_blocks_flat
460
+ ):
461
+ self._current_batch_jobs_to_flush.update(
462
+ jid
463
+ for bid in new_blocks_flat
464
+ for jid in self._block_id_to_pending_jobs.get(bid, ())
465
+ )
466
467
num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
468
num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
@@ -491,36 +556,52 @@ class OffloadingConnectorScheduler:
556
)
557
dst_spec = store_output.store_spec
558
494
- reqs_to_store[req_id] = (src_spec, dst_spec)
495
- self._reqs_being_stored[req_id] |= keys_to_store
559
+ job_id = self._generate_job_id()
560
+ # a store can only be issued when no load is pending.
561
+ if req_status.transfer_jobs:
562
+ any_jid = next(iter(req_status.transfer_jobs))
563
+ assert self._jobs[any_jid].is_store
564
+ req_status.transfer_jobs.add(job_id)
565
+ self._jobs[job_id] = TransferJobStatus(
566
+ req_id=req_id,
567
+ pending_count=self.config.num_workers,
568
+ keys=set(keys_to_store),
569
+ is_store=True,
570
+ gpu_block_ids=src_block_ids,
571
+ )
572
+
573
+ store_jobs[job_id] = TransferJob(
574
+ req_id=req_id, transfer_spec=(src_spec, dst_spec)
575
+ )
576
577
logger.debug(
498
- "Request %s offloading %s blocks upto %d tokens",
578
+ "Request %s offloading %s blocks upto %d tokens (job %d)",
579
req_id,
580
len(keys_to_store),
581
num_offloadable_tokens,
582
+ job_id,
583
)
584
504
- return reqs_to_store
585
+ return store_jobs
586
587
def build_connector_meta(
588
self, scheduler_output: SchedulerOutput
589
) -> KVConnectorMetadata:
509
- meta = OffloadingConnectorMetadata(
510
- reqs_to_load=self._reqs_to_load,
511
- reqs_to_store=self._get_reqs_to_store(scheduler_output),
512
- reqs_to_flush=scheduler_output.preempted_req_ids,
513
- )
514
- self._reqs_to_load = {}
515
-
516
- # NOTE (orozery): we should move this logic to update_connector_output
517
- # once KVConnectorOutput allows us to report completed transfers
590
for req_id in scheduler_output.preempted_req_ids or ():
519
- keys = self._reqs_being_stored.get(req_id)
520
- if keys:
521
- self.manager.complete_store(keys)
522
- keys.clear()
591
+ req_status = self._req_status.get(req_id)
592
+ if req_status is None or not req_status.transfer_jobs:
593
+ continue
594
+ any_jid = next(iter(req_status.transfer_jobs))
595
+ assert self._jobs[any_jid].is_store
596
+ self._current_batch_jobs_to_flush.update(req_status.transfer_jobs)
597
598
+ meta = OffloadingConnectorMetadata(
599
+ load_jobs=self._current_batch_load_jobs,
600
+ store_jobs=self._build_store_jobs(scheduler_output),
601
+ jobs_to_flush=self._current_batch_jobs_to_flush,
602
+ )
603
+ self._current_batch_load_jobs = {}
604
+ self._current_batch_jobs_to_flush = set()
605
return meta
606
607
def update_connector_output(self, connector_output: KVConnectorOutput):
@@ -531,17 +612,37 @@ class OffloadingConnectorScheduler:
612
connector_output (KVConnectorOutput): the worker-side
613
connectors output.
614
"""
534
- for req_id in connector_output.finished_sending or []:
535
- keys = self._reqs_being_stored.pop(req_id, None)
536
- if keys:
537
- self.manager.complete_store(keys)
538
-
539
- for req_id in connector_output.finished_recving or []:
540
- keys = self._reqs_being_loaded.pop(req_id, None)
541
- if keys:
615
+ meta = connector_output.kv_connector_worker_meta
616
+ if not isinstance(meta, OffloadingWorkerMetadata):
617
+ assert meta is None
618
+ meta = OffloadingWorkerMetadata()
619
+ for job_id, count in meta.completed_jobs.items():
620
+ assert count > 0
621
+ job_status = self._jobs[job_id]
622
+ job_status.pending_count -= count
623
+ if job_status.pending_count > 0:
624
+ continue
625
+ assert job_status.pending_count == 0
626
+
627
+ if job_status.is_store:
628
+ self.manager.complete_store(job_status.keys)
629
+ else:
630
+ self.manager.complete_load(job_status.keys)
631
if self._blocks_being_loaded:
543
- self._blocks_being_loaded.difference_update(keys)
544
- self.manager.complete_load(keys)
632
+ self._blocks_being_loaded.difference_update(job_status.keys)
633
+
634
+ req_status = self._req_status[job_status.req_id]
635
+ if self._block_id_to_pending_jobs and req_status.req.is_finished():
636
+ for bid in job_status.gpu_block_ids or ():
637
+ pending = self._block_id_to_pending_jobs[bid]
638
+ pending.remove(job_id)
639
+ if not pending:
640
+ del self._block_id_to_pending_jobs[bid]
641
+
642
+ del self._jobs[job_id]
643
+ req_status.transfer_jobs.remove(job_id)
644
+ if not req_status.transfer_jobs and req_status.req.is_finished():
645
+ del self._req_status[job_status.req_id]
646
647
def request_finished(
648
self,
@@ -558,14 +659,21 @@ class OffloadingConnectorScheduler:
659
Optional KVTransferParams to be included in the request outputs
660
returned by the engine.
661
"""
561
- req_id = request.request_id
562
-
662
# TODO(orozery): possibly kickoff offload for last block
663
# which may have been deferred due to async scheduling
565
- self._req_status.pop(req_id, None)
566
-
567
- request_being_stored = req_id in self._reqs_being_stored
568
- return request_being_stored, None
664
+ req_status = self._req_status.get(request.request_id)
665
+ if req_status is None:
666
+ return False, None
667
+ if not req_status.transfer_jobs:
668
+ del self._req_status[request.request_id]
669
+ return False, None
670
+ # Pending stores will outlive the request's block ownership.
671
+ # Register them so future block reuse triggers a flush.
672
+ for job_id in req_status.transfer_jobs:
673
+ job_status = self._jobs[job_id]
674
+ for bid in job_status.gpu_block_ids or ():
675
+ self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
676
+ return False, None
677
678
def take_events(self) -> Iterable[KVCacheEvent]:
679
"""Take the KV cache events from the connector.
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+35
-71
@@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
11
)
12
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
13
OffloadingConnectorMetadata,
14
+ OffloadingWorkerMetadata,
15
ReqId,
16
)
17
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
@@ -45,24 +46,11 @@ class OffloadingConnectorWorker:
46
self.spec = spec
47
self.worker = OffloadingWorker()
48
48
- self._job_counter = 0
49
-
49
self.kv_connector_stats = OffloadingConnectorStats()
51
- # req_id -> (job_id, store)
52
- self._jobs: dict[int, tuple[ReqId, bool]] = {}
53
- # req_id -> active job IDs
54
- self._load_job: dict[ReqId, int] = {}
55
- # req_id -> set(active job IDs)
56
- self._store_jobs = defaultdict[ReqId, set[int]](set)
57
- # list of store jobs pending submission (job_id, transfer_spec)
50
+ # job_id -> req_id for in-flight loads.
51
+ self._load_jobs: dict[int, ReqId] = {}
52
self._unsubmitted_store_jobs: list[tuple[int, TransferSpec]] = []
59
-
60
- self._finished_reqs_waiting_for_store: set[ReqId] = set()
61
-
62
- def _generate_job_id(self) -> int:
63
- job_id = self._job_counter
64
- self._job_counter = job_id + 1
65
- return job_id
53
+ self._connector_worker_meta = OffloadingWorkerMetadata()
54
55
def _register_handlers(self, kv_caches: CanonicalKVCaches):
56
for src_cls, dst_cls, handler in self.spec.get_handlers(kv_caches):
@@ -301,10 +289,8 @@ class OffloadingConnectorWorker:
289
assert success
290
self._unsubmitted_store_jobs.clear()
291
304
- for req_id in kv_connector_metadata.reqs_to_flush or ():
305
- job_ids = self._store_jobs.get(req_id)
306
- if job_ids:
307
- self.worker.wait(job_ids)
292
+ if kv_connector_metadata.jobs_to_flush:
293
+ self.worker.wait(kv_connector_metadata.jobs_to_flush)
294
295
def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):
296
for job_id, transfer_spec in self._unsubmitted_store_jobs:
@@ -312,41 +298,33 @@ class OffloadingConnectorWorker:
298
assert success
299
self._unsubmitted_store_jobs.clear()
300
315
- for req_id, transfer_spec in metadata.reqs_to_load.items():
316
- job_id = self._generate_job_id()
317
- self._jobs[job_id] = (req_id, False)
318
- assert req_id not in self._load_job
319
- self._load_job[req_id] = job_id
320
- success = self.worker.transfer_async(job_id, transfer_spec)
301
+ for job_id, entry in metadata.load_jobs.items():
302
+ self._load_jobs[job_id] = entry.req_id
303
+ success = self.worker.transfer_async(job_id, entry.transfer_spec)
304
assert success
305
306
def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
324
- for req_id, transfer_spec in metadata.reqs_to_store.items():
325
- job_id = self._generate_job_id()
326
- self._jobs[job_id] = (req_id, True)
327
- self._store_jobs[req_id].add(job_id)
328
- # NOTE(orozery): defer the store to the beginning of the next engine step,
329
- # so that offloading starts AFTER transfers related to token sampling,
330
- # thereby avoiding delays to token generation due to offloading.
331
- self._unsubmitted_store_jobs.append((job_id, transfer_spec))
307
+ for job_id, entry in metadata.store_jobs.items():
308
+ # NOTE(orozery): defer the store to the beginning of the next
309
+ # engine step, so that offloading starts AFTER transfers related
310
+ # to token sampling, thereby avoiding delays to token generation.
311
+ self._unsubmitted_store_jobs.append((job_id, entry.transfer_spec))
312
313
def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
314
"""
335
- Notifies worker-side connector ids of requests that have
336
- finished generating tokens.
337
- Returns a list of request IDs that finished loading or storing.
338
-
315
Returns:
340
- ids of requests that have finished asynchronous transfer
341
- tuple of (sending/saving ids, recving/loading ids).
316
+ tuple of (finished_sending, finished_recving). Stores never
317
+ emit finished_sending — the scheduler tracks store completion
318
+ via kv_connector_worker_meta.completed_jobs and fences any
319
+ block reuse via jobs_to_flush. Loads still emit
320
+ finished_recving so the base scheduler can resume requests
321
+ blocked on remote KV (and free aborted-during-load reqs).
322
"""
343
- finished_sending = set()
344
- finished_recving = set()
323
+ finished_recving: set[str] = set()
324
for transfer_result in self.worker.get_finished():
325
# we currently do not support job failures
326
job_id = transfer_result.job_id
327
assert transfer_result.success
349
- req_id, store = self._jobs.pop(job_id)
328
if (
329
transfer_result.transfer_time
330
and transfer_result.transfer_size is not None
@@ -357,31 +335,21 @@ class OffloadingConnectorWorker:
335
time=transfer_result.transfer_time,
336
transfer_type=transfer_result.transfer_type,
337
)
360
- if store:
361
- req_jobs = self._store_jobs[req_id]
362
- req_jobs.remove(job_id)
363
- if req_jobs:
364
- continue
365
-
366
- if req_id in self._finished_reqs_waiting_for_store:
367
- self._finished_reqs_waiting_for_store.remove(req_id)
368
- finished_sending.add(req_id)
369
- del self._store_jobs[req_id]
370
- else:
371
- req_job = self._load_job[req_id]
372
- assert job_id == req_job
373
- del self._load_job[req_id]
338
+
339
+ self._connector_worker_meta.mark_completed(job_id)
340
+ req_id = self._load_jobs.pop(job_id, None)
341
+ if req_id is not None:
342
finished_recving.add(req_id)
343
376
- for req_id in finished_req_ids:
377
- pending_req_jobs = self._store_jobs.get(req_id)
378
- if pending_req_jobs:
379
- self._finished_reqs_waiting_for_store.add(req_id)
380
- elif pending_req_jobs is not None:
381
- finished_sending.add(req_id)
382
- del self._store_jobs[req_id]
344
+ return set(), finished_recving
345
384
- return finished_sending, finished_recving
346
+ def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:
347
+ """Return completed transfer job IDs since the last call."""
348
+ if not self._connector_worker_meta.completed_jobs:
349
+ return None
350
+ meta = self._connector_worker_meta
351
+ self._connector_worker_meta = OffloadingWorkerMetadata()
352
+ return meta
353
354
def get_kv_connector_stats(self) -> KVConnectorStats | None:
355
"""
@@ -396,11 +364,7 @@ class OffloadingConnectorWorker:
364
return kv_connector_stats
365
366
def shutdown(self) -> None:
399
- # Drop deferred store jobs: there is no point in submitting
400
- # them during shutdown.
367
self._unsubmitted_store_jobs.clear()
402
- self._jobs.clear()
403
- self._load_job.clear()
404
- self._store_jobs.clear()
405
- self._finished_reqs_waiting_for_store.clear()
368
+ self._load_jobs.clear()
369
+ self._connector_worker_meta = OffloadingWorkerMetadata()
370
self.worker.shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+6
@@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
20
)
21
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
22
OffloadingConnectorMetadata,
23
+ OffloadingWorkerMetadata,
24
)
25
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
26
OffloadingConnectorStats,
@@ -111,6 +112,11 @@ class OffloadingConnector(KVConnectorBase_V1):
112
assert self.connector_worker is not None
113
return self.connector_worker.get_finished(finished_req_ids)
114
115
+ def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:
116
+ if self.connector_worker is not None:
117
+ return self.connector_worker.build_connector_worker_meta()
118
+ return None
119
+
120
def get_num_new_matched_tokens(
121
self, request: "Request", num_computed_tokens: int
122
) -> tuple[int | None, bool]: