diff --git a/integration/vllm/README.md b/integration/vllm/README.md index 5b11ec4..9521a7b 100644 --- a/integration/vllm/README.md +++ b/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. | diff --git a/integration/vllm/src/dfkv_vllm/scheduler.py b/integration/vllm/src/dfkv_vllm/scheduler.py index 8f9390c..7f787be 100644 --- a/integration/vllm/src/dfkv_vllm/scheduler.py +++ b/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 diff --git a/integration/vllm/src/dfkv_vllm/worker.py b/integration/vllm/src/dfkv_vllm/worker.py index 101a1b5..cddd985 100644 --- a/integration/vllm/src/dfkv_vllm/worker.py +++ b/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"]: diff --git a/integration/vllm/tests/test_worker_lifecycle.py b/integration/vllm/tests/test_worker_lifecycle.py index 12b1560..a9bfb24 100644 --- a/integration/vllm/tests/test_worker_lifecycle.py +++ b/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()