fix(vllm): support synchronous GPUDirect loads

integration/vllm/README.md

@@ -83,7 +83,7 @@ LMCache connector access logs, so one setting covers every integration. Format:
 | `batch_concurrency` | `0`=auto | client fan-out for batch ops; the real throughput lever (depth is flat). Auto = `min(max(nodes, 8), 32)`: 8-way parallel on single-node, one-per-node on multi-node. Set >0 to pin a fixed value. |
 | `rail_affinity` | `False` | Bind each vLLM worker process to a primary rail selected by world-group local rank; requires an ordered multi-rail `DFKV_RDMA_DEV`. |
 | `rail_affinity_fallbacks` | `1` | Number of ordered neighboring fallback rails when affinity is enabled. `0` keeps strict one-rank/one-rail; values above the available rail count are bounded. |
-| `load_async` | `True` | async KV load: the scheduler returns `WAITING_FOR_REMOTE_KVS` and the load runs off the critical path. Keep `True`. |
+| `load_async` | `True` | `True` returns `WAITING_FOR_REMOTE_KVS` and overlaps GPUDirect loads with unrelated model work. `False` performs each requested load synchronously in `start_load_kv`, before the forward pass. Use `False` for hybrid state-cache models when the engine cannot guarantee that remote writes target blocks disjoint from concurrent compute. |
 | `transfer_queue_capacity` | `256` | Maximum queued requests in each direction (`1..65536`). All receive workers consume one shared receive queue of this capacity; capacity is not multiplied by `recv_workers`. Submission is non-blocking: a full queue rejects new saves as completed (releasing finish/free fences) and rejects new loads as load errors (forcing recompute), so overload cannot grow memory or pin blocks indefinitely. Invalid or out-of-range values abort connector construction. |
 | `recv_workers` | `1` | Receive/load worker count (`1..32`). Workers consume the shared bounded receive queue and may execute independent native GETs concurrently. Invalid, boolean, or out-of-range values abort connector construction. |
 | `enable_cross_layers_blocks` | `False` | opt-in for engines whose paged layout interleaves layers within a block. Leave `False` unless you know the layout needs it. |

integration/vllm/src/dfkv_vllm/scheduler.py

@@ -64,6 +64,8 @@ def __init__(
             vllm_config.kv_transfer_config.kv_connector_extra_config
         )
         self.load_async = extra_config.get("load_async", True)
+        if not isinstance(self.load_async, bool):
+            raise ValueError("dfkv connector: load_async must be a boolean")
         self.lookup_async = extra_config.get("lookup_async", False)
         self.client = LookupKeyClient(vllm_config)
         self._closed = False

integration/vllm/src/dfkv_vllm/worker.py

@@ -1091,6 +1091,11 @@ def get_and_clear_block_ids_with_load_errors(self) -> set[int]:
             self._invalid_block_ids.clear()
         return invalid_block_ids
 
+    def load_request_sync(self, request: ReqMeta) -> None:
+        """Load one request on the model thread before its forward pass."""
+        self._handle_request(request)
+
+
     def _handle_request(self, req_meta: ReqMeta):
         req_id = req_meta.req_id
         enqueued_at = getattr(req_meta, "_dfkv_receive_enqueued_at", None)
@@ -1352,14 +1357,12 @@ def __init__(
         ensure_deterministic_block_hashing(vllm_config.cache_config)
         self.kv_role = vllm_config.kv_transfer_config.kv_role
         self.load_async = extra.get("load_async", True)
-        if not self.load_async:
-            # Fail at construction with an actionable message: the synchronous
-            # load mode was never implemented, and the old hot-path assert in
-            # get_finished() crashed the worker mid-serving instead.
-            raise ValueError(
-                "dfkv connector: kv_connector_extra_config.load_async=false is "
-                "not supported (loads are handled by the async recv thread); "
-                "remove the key or set it to true")
+        if not isinstance(self.load_async, bool):
+            raise ValueError("dfkv connector: load_async must be a boolean")
+        logger.info(
+            "dfkv load mode: %s",
+            "async-overlap" if self.load_async else "synchronous-before-forward",
+        )
         self.cache_config = vllm_config.cache_config
         self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
             kv_cache_config, vllm_config
@@ -1865,40 +1868,39 @@ def start_load_kv(
         self,
         metadata: DfkvStoreConnectorMetadata,
     ):
-        """Fence every preempted request's active GPU transfer.
+        """Fence preemptions and perform synchronous loads before forward."""
+        if metadata.preempted_req_ids:
+            if self.kv_recv_thread is not None:
+                self.kv_recv_thread.cancel_requests(
+                    metadata.preempted_req_ids,
+                    wait=True,
+                    fail_closed=False,
+                )
+            send_thread = self.kv_send_thread
+            if send_thread is not None:
+                # Drop queued entries first, then join whichever entry is active.
+                for req_id in metadata.preempted_req_ids:
+                    send_thread.delete_finished_stored_request(req_id)
+                for req_id in metadata.preempted_req_ids:
+                    wait_start = time.perf_counter()
+                    while not send_thread.wait_for_inflight_put(req_id):
+                        logger.error(
+                            "preemption fence still waiting for in-flight save "
+                            "of request %s after %.1fs; GPU block reuse remains "
+                            "fenced until the native save exits",
+                            req_id,
+                            time.perf_counter() - wait_start,
+                        )
 
-        A queued receive is cancelled without a native call.  An active receive
-        cannot be interrupted safely, so this waits for only that request's
-        GPUDirect write before its blocks may be reused.  The send-side fence
-        likewise waits for its request's active RDMA read.
-        """
-        if not metadata.preempted_req_ids:
-            return
-        if self.kv_recv_thread is not None:
-            self.kv_recv_thread.cancel_requests(
-                metadata.preempted_req_ids,
-                wait=True,
-                fail_closed=False,
-            )
-        send_thread = self.kv_send_thread
-        if send_thread is None:
+        if self.load_async:
             return
-        # Drop queued entries first (the dequeue gate re-checks per entry),
-        # then join whichever entry is already executing. The timeout is only
-        # a diagnostic interval: returning from this hook while native PUT can
-        # still read the blocks would let the scheduler reuse them.
-        for req_id in metadata.preempted_req_ids:
-            send_thread.delete_finished_stored_request(req_id)
-        for req_id in metadata.preempted_req_ids:
-            wait_start = time.perf_counter()
-            while not send_thread.wait_for_inflight_put(req_id):
-                logger.error(
-                    "preemption fence still waiting for in-flight save "
-                    "of request %s after %.1fs; GPU block reuse remains fenced "
-                    "until the native save exits",
-                    req_id,
-                    time.perf_counter() - wait_start,
-                )
+        assert self.kv_recv_thread is not None
+        for request in metadata.requests:
+            load_spec = request.load_spec
+            if load_spec is None or not load_spec.can_load:
+                continue
+            load_spec.token_len = load_spec.kvpool_cached_tokens
+            self.kv_recv_thread.load_request_sync(request)
 
     def wait_for_save(
         self,
@@ -1926,16 +1928,16 @@ def get_finished(
                 wait=True,
                 fail_closed=False,
             )
-        # Issue async loads
-        for request in meta.requests:
-            load_spec = request.load_spec
-            if load_spec is None or not load_spec.can_load:
-                continue
-
-            load_spec.token_len = load_spec.kvpool_cached_tokens
-
-            assert self.kv_recv_thread is not None
-            self.kv_recv_thread.add_request(request)
+        # Async mode overlaps loads with unrelated model work. Synchronous mode
+        # already completed them in start_load_kv, before this forward pass.
+        if self.load_async:
+            for request in meta.requests:
+                load_spec = request.load_spec
+                if load_spec is None or not load_spec.can_load:
+                    continue
+                load_spec.token_len = load_spec.kvpool_cached_tokens
+                assert self.kv_recv_thread is not None
+                self.kv_recv_thread.add_request(request)
 
         # Issue stores with CUDA event synchronization
         if self.kv_role in ["kv_producer", "kv_both"]:

integration/vllm/tests/test_worker_lifecycle.py

@@ -201,6 +201,64 @@ def test_repeated_close_joins_resources_and_closes_native_once(self):
         with self.assertRaisesRegex(RuntimeError, "worker is closed"):
             worker.register_kv_caches({"layer": object()})
 
+@unittest.skipUnless(HAVE_VLLM, "requires vllm (dfkv_vllm.worker imports it)")
+class SynchronousLoadTest(unittest.TestCase):
+    class FakeRecv:
+        def __init__(self):
+            self.loaded = []
+
+        def load_request_sync(self, request):
+            self.loaded.append(request.req_id)
+
+        def cancel_requests(self, *args, **kwargs):
+            raise AssertionError("no preemptions expected")
+
+    @staticmethod
+    def _request(req_id: str, can_load: bool):
+        return types.SimpleNamespace(
+            req_id=req_id,
+            load_spec=types.SimpleNamespace(
+                can_load=can_load,
+                kvpool_cached_tokens=128,
+                token_len=0,
+            ),
+        )
+
+    def test_sync_mode_loads_before_forward(self):
+        worker = DfkvStoreWorker.__new__(DfkvStoreWorker)
+        worker.load_async = False
+        worker.kv_recv_thread = self.FakeRecv()
+        worker.kv_send_thread = None
+        load = self._request("load", True)
+        skip = self._request("skip", False)
+        metadata = types.SimpleNamespace(
+            preempted_req_ids=set(),
+            requests=[load, skip],
+        )
+
+        worker.start_load_kv(metadata)
+
+        self.assertEqual(worker.kv_recv_thread.loaded, ["load"])
+        self.assertEqual(load.load_spec.token_len, 128)
+        self.assertEqual(skip.load_spec.token_len, 0)
+
+    def test_async_mode_defers_load_to_get_finished(self):
+        worker = DfkvStoreWorker.__new__(DfkvStoreWorker)
+        worker.load_async = True
+        worker.kv_recv_thread = self.FakeRecv()
+        worker.kv_send_thread = None
+        load = self._request("load", True)
+        metadata = types.SimpleNamespace(
+            preempted_req_ids=set(),
+            requests=[load],
+        )
+
+        worker.start_load_kv(metadata)
+
+        self.assertEqual(worker.kv_recv_thread.loaded, [])
+        self.assertEqual(load.load_spec.token_len, 0)
+
+
 
 if __name__ == "__main__":
     unittest.main()