[KV Offload] Per-job store completion for CPU offloading connector

tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
+    OffloadingWorkerMetadata,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def test_aggregate_sums_counts():
+    meta1 = OffloadingWorkerMetadata(completed_store_jobs={42: 1})
+    meta2 = OffloadingWorkerMetadata(completed_store_jobs={42: 1})
+    result = meta1.aggregate(meta2)
+    assert result.completed_store_jobs == {42: 2}
+
+
+def test_aggregate_disjoint_jobs():
+    meta1 = OffloadingWorkerMetadata(completed_store_jobs={42: 1})
+    meta2 = OffloadingWorkerMetadata(completed_store_jobs={43: 1})
+    result = meta1.aggregate(meta2)
+    assert result.completed_store_jobs == {42: 1, 43: 1}
+
+
+def test_aggregate_multiple_workers():
+    meta1 = OffloadingWorkerMetadata(completed_store_jobs={42: 1, 43: 1})
+    meta2 = OffloadingWorkerMetadata(completed_store_jobs={42: 1})
+    meta3 = OffloadingWorkerMetadata(completed_store_jobs={42: 1, 43: 1})
+    result = meta1.aggregate(meta2).aggregate(meta3)
+    assert result.completed_store_jobs == {42: 3, 43: 2}

vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py

@@ -2,14 +2,44 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
 
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+    KVConnectorWorkerMetadata,
+)
 from vllm.v1.kv_offload.worker.worker import TransferSpec
 
 ReqId = str
 
+# A store entry bundles a scheduler-assigned job ID with the transfer spec.
+# The job ID allows per-job completion tracking: the worker reports it back
+# when the DMA finishes, and the scheduler calls complete_store immediately.
+StoreJobEntry = tuple[int, TransferSpec]
+
 
 @dataclass
 class OffloadingConnectorMetadata(KVConnectorMetadata):
     reqs_to_load: dict[ReqId, TransferSpec]
-    reqs_to_store: dict[ReqId, TransferSpec]
+    reqs_to_store: dict[ReqId, StoreJobEntry]
     reqs_to_flush: set[str] | None = None
+
+
+@dataclass
+class OffloadingWorkerMetadata(KVConnectorWorkerMetadata):
+    """Worker -> Scheduler metadata for completed store jobs.
+
+    Each worker reports {job_id: 1} for newly completed stores.
+    aggregate() sums counts across workers within a step.
+    The scheduler accumulates across steps and processes
+    a store completion only when count reaches world_size.
+    """
+
+    completed_store_jobs: dict[int, int]
+
+    def aggregate(
+        self, other: "KVConnectorWorkerMetadata"
+    ) -> "KVConnectorWorkerMetadata":
+        assert isinstance(other, OffloadingWorkerMetadata)
+        merged = dict(self.completed_store_jobs)
+        for k, v in other.completed_store_jobs.items():
+            merged[k] = merged.get(k, 0) + v
+        return OffloadingWorkerMetadata(completed_store_jobs=merged)

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -11,7 +11,9 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
     OffloadingConnectorMetadata,
+    OffloadingWorkerMetadata,
     ReqId,
+    StoreJobEntry,
 )
 from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -129,6 +131,16 @@ def __init__(self, spec: OffloadingSpec):
         self._reqs_being_stored = defaultdict[ReqId, set[OffloadKey]](set)
         self._reqs_being_loaded = defaultdict[ReqId, set[OffloadKey]](set)
 
+        # Per-job store completion tracking.
+        # Scheduler assigns job IDs so it can call complete_store per-job.
+        self._store_job_counter: int = 0
+        self._store_job_hashes: dict[int, set[OffloadKey]] = {}
+        self._store_job_to_req: dict[int, ReqId] = {}
+        self._req_to_jobs: dict[ReqId, set[int]] = defaultdict(set)
+        # Accumulate partial TP completions across steps.
+        self._expected_worker_count: int = spec.vllm_config.parallel_config.world_size
+        self._store_job_pending_counts: dict[int, int] = {}
+
     def get_num_new_matched_tokens(
         self, request: Request, num_computed_tokens: int
     ) -> tuple[int | None, bool]:
@@ -264,12 +276,14 @@ def update_state_after_alloc(
         if self._blocks_being_loaded is not None:
             self._blocks_being_loaded.update(req_blocks_being_loaded)
 
-    def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
+    def _get_reqs_to_store(
+        self, scheduler_output: SchedulerOutput
+    ) -> dict[ReqId, StoreJobEntry]:
         # Below assertion will be removed once this function supports HMA
         assert len(self.config.kv_group_configs) == 1
         group_config = self.config.kv_group_configs[0]
 
-        reqs_to_store: dict[ReqId, TransferSpec] = {}
+        reqs_to_store: dict[ReqId, StoreJobEntry] = {}
         # iterate over both new and cached requests
         for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
             req_status = self._req_status[req_id]
@@ -332,18 +346,33 @@ def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
                 src_block_ids, group_sizes=(len(src_block_ids),)
             )
 
-            reqs_to_store[req_id] = (src_spec, dst_spec)
+            # Assign a scheduler job ID for per-job completion tracking.
+            job_id = self._store_job_counter
+            self._store_job_counter += 1
+            self._store_job_hashes[job_id] = set(keys_to_store)
+            self._store_job_to_req[job_id] = req_id
+            self._req_to_jobs[req_id].add(job_id)
+
+            reqs_to_store[req_id] = (job_id, (src_spec, dst_spec))
             self._reqs_being_stored[req_id] |= keys_to_store
 
             logger.debug(
-                "Request %s offloading %s blocks starting from block #%d",
+                "Request %s offloading %s blocks starting from block #%d (job %d)",
                 req_id,
                 len(keys_to_store),
                 start_block_idx,
+                job_id,
             )
 
         return reqs_to_store
 
+    def _cleanup_store_jobs_for_req(self, req_id: ReqId) -> None:
+        """Remove per-job tracking state for a given request."""
+        for jid in self._req_to_jobs.pop(req_id, ()):
+            self._store_job_hashes.pop(jid, None)
+            self._store_job_to_req.pop(jid, None)
+            self._store_job_pending_counts.pop(jid, None)
+
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
@@ -361,6 +390,7 @@ def build_connector_meta(
             if keys:
                 self.manager.complete_store(keys)
                 keys.clear()
+            self._cleanup_store_jobs_for_req(req_id)
 
         return meta
 
@@ -372,10 +402,39 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
             connector_output (KVConnectorOutput): the worker-side
                 connectors output.
         """
+        # Process per-job store completions via worker metadata.
+        # Each worker reports {job_id: 1} when its DMA finishes.
+        # We accumulate across steps; once count >= world_size,
+        # the job is fully complete and we call complete_store.
+        meta = connector_output.kv_connector_worker_meta
+        if isinstance(meta, OffloadingWorkerMetadata):
+            for job_id, count in meta.completed_store_jobs.items():
+                total = self._store_job_pending_counts.get(job_id, 0) + count
+                if total >= self._expected_worker_count:
+                    self._store_job_pending_counts.pop(job_id, None)
+                    block_hashes = self._store_job_hashes.pop(job_id, None)
+                    if block_hashes:
+                        self.manager.complete_store(block_hashes)
+                        req_id = self._store_job_to_req.pop(job_id, None)
+                        if req_id is not None:
+                            self._req_to_jobs.get(req_id, set()).discard(job_id)
+                            remaining = self._reqs_being_stored.get(req_id)
+                            if remaining is not None:
+                                remaining -= block_hashes
+                                # Keep the empty set so request_finished()
+                                # still returns True, ensuring _free_blocks
+                                # waits for finished_sending.
+                else:
+                    self._store_job_pending_counts[job_id] = total
+
+        # Handle request-level completion (for _free_blocks in scheduler).
+        # If per-job tracking already cleared _reqs_being_stored, the pop
+        # returns None and complete_store is a no-op.
         for req_id in connector_output.finished_sending or []:
             keys = self._reqs_being_stored.pop(req_id, None)
             if keys:
                 self.manager.complete_store(keys)
+            self._cleanup_store_jobs_for_req(req_id)
 
         for req_id in connector_output.finished_recving or []:
             keys = self._reqs_being_loaded.pop(req_id, None)

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -11,6 +11,7 @@
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
     OffloadingConnectorMetadata,
+    OffloadingWorkerMetadata,
     ReqId,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
@@ -59,6 +60,12 @@ def __init__(self, spec: OffloadingSpec):
 
         self._finished_reqs_waiting_for_store: set[ReqId] = set()
 
+        # Maps internal worker job IDs to scheduler-assigned job IDs
+        # for per-job store completion reporting.
+        self._internal_to_sched_job: dict[int, int] = {}
+        # Completed store jobs to report via build_connector_worker_meta.
+        self._completed_store_jobs: dict[int, int] = {}
+
     def _generate_job_id(self) -> int:
         job_id = self._job_counter
         self._job_counter = job_id + 1
@@ -321,10 +328,12 @@ def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):
             assert success
 
     def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
-        for req_id, transfer_spec in metadata.reqs_to_store.items():
+        for req_id, (sched_job_id, transfer_spec) in metadata.reqs_to_store.items():
             job_id = self._generate_job_id()
             self._jobs[job_id] = (req_id, True)
             self._store_jobs[req_id].add(job_id)
+            # Map internal job ID to scheduler job ID for per-job reporting.
+            self._internal_to_sched_job[job_id] = sched_job_id
             # NOTE(orozery): defer the store to the beginning of the next engine step,
             # so that offloading starts AFTER transfers related to token sampling,
             # thereby avoiding delays to token generation due to offloading.
@@ -358,6 +367,11 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
                     transfer_type=transfer_result.transfer_type,
                 )
             if store:
+                # Report per-job completion via worker metadata.
+                sched_job_id = self._internal_to_sched_job.pop(job_id, None)
+                if sched_job_id is not None:
+                    self._completed_store_jobs[sched_job_id] = 1
+
                 req_jobs = self._store_jobs[req_id]
                 req_jobs.remove(job_id)
                 if req_jobs:
@@ -383,6 +397,16 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
 
         return finished_sending, finished_recving
 
+    def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None:
+        """Return completed store job IDs since the last call."""
+        if not self._completed_store_jobs:
+            return None
+        meta = OffloadingWorkerMetadata(
+            completed_store_jobs=self._completed_store_jobs,
+        )
+        self._completed_store_jobs = {}
+        return meta
+
     def get_kv_connector_stats(self) -> KVConnectorStats | None:
         """
         Get the KV transfer stats for the connector.

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -105,6 +105,11 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
         assert self.connector_worker is not None
         return self.connector_worker.get_finished(finished_req_ids)
 
+    def build_connector_worker_meta(self):
+        if self.connector_worker is not None:
+            return self.connector_worker.build_connector_worker_meta()
+        return None
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int | None, bool]: