From a7df84bd3b15f3b1d790a8f1dc6e0278a84b7582 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sat, 23 May 2026 21:02:13 +0800 Subject: [PATCH] Direct RDMA read: D reads cached KV from C's GPU without C's scheduler Complete implementation of direct RDMA read for KV cache migration: vLLM Mooncake connector (mooncake_connector.py): - PullReqMeta: add direct_read flag + block_hashes - MooncakeConnectorMetadata: add hash_table_updates/removals for scheduler->worker block hash sync - MooncakeConnectorScheduler: set_block_pool() to access BlockPool, build_connector_meta() computes hash table deltas each step, update_state_after_alloc() captures request block hashes for direct_read - MooncakeConnectorWorker: _start_direct_read() + _direct_read_single() implements D-side RDMA read via batch_transfer_sync_read, with HTTP query/unpin to C's bootstrap server Bootstrap server (mooncake_utils.py): - POST /query_blocks: look up block hashes, return block_ids + GPU layout - POST /unpin_blocks: release pin tracking - set_worker_kv_info(): register GPU addresses at init - update_hash_table(): receive scheduler deltas each step Scheduler (scheduler.py): - One-line hookup: pass block_pool to connector after KVCacheManager init Proxy (cache_aware_proxy.py): - _handle_direct_read_offload: sends request ONLY to D with direct_read=True + remote_bootstrap_addr. No request to C at all. - C's scheduler is completely uninvolved (0 GPU time on C) Co-Authored-By: Claude Opus 4.6 (1M context) --- scripts/cache_aware_proxy.py | 131 +++---------- .../v1/mooncake/mooncake_connector.py | 184 ++++++++++++++++-- .../v1/mooncake/mooncake_utils.py | 77 +++++++- .../vllm/vllm/v1/core/sched/scheduler.py | 2 + 4 files changed, 271 insertions(+), 123 deletions(-) diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index 7708bb3..e5b8ff4 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -351,33 +351,28 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h offload_reason = "colocated_cheaper_%.1fvs%.1f" % (colocated_cost, offload_cost) if use_offload: - p_inst = p_candidate + # Direct RDMA read: D reads cached KV from C_s's GPU, no request to C_s + c_inst = best_inst # has cache (not doing any work) d_inst = d_candidate d_idx = combined_instances.index(d_inst) - # Accounting: reserve both P and D immediately so router sees the load - p_new = max(0, input_length - p_inst.estimate_cache_hit(token_ids)) if token_ids else input_length - p_inst.ongoing_tokens += input_length - p_inst.pending_prefill_tokens += p_new - p_inst.num_requests += 1 - p_inst.active_p_offloads += 1 - breakdown["p_new_tokens"] = p_new - d_inst.ongoing_tokens += input_length + d_inst.pending_prefill_tokens += estimated_new d_inst.num_requests += 1 + c_inst.active_p_offloads += 1 breakdown["route_class"] = "HEAVY_OFFLOAD" breakdown["offload_reason"] = offload_reason - breakdown["p_inst"] = p_inst.url + breakdown["c_inst"] = c_inst.url breakdown["d_inst"] = d_inst.url - breakdown["p_load"] = p_inst.ongoing_tokens - breakdown["d_load"] = d_inst.ongoing_tokens + breakdown["cache_hit_tokens"] = cache_hit if session_id: session_affinity[session_id] = d_idx - return await _handle_heavy_offload(api, req_data, headers, token_ids, - input_length, p_inst, d_inst, breakdown) + return await _handle_direct_read_offload( + api, req_data, headers, token_ids, input_length, + c_inst, d_inst, cache_hit, estimated_new, breakdown) else: if estimated_new >= HEAVY_THRESHOLD: breakdown["route_class"] = "HEAVY_COLO" @@ -423,112 +418,31 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h PREFILL_TIMEOUT_S = 120 # max seconds to wait for P-instance prefill -async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length, - p_inst, d_inst, breakdown): - """HEAVY request with cache-aware KV migration. - - C_s (p_inst, has cache) exports cached KV blocks via Mooncake. - D (d_inst, idle) pulls cached blocks + does local prefill for new tokens + decodes. - C_s's blocks are auto-freed by Mooncake after D confirms receipt. - - On export failure, falls back to co-located prefill+decode on d_inst. +async def _handle_direct_read_offload(api, req_data, headers, token_ids, + input_length, c_inst, d_inst, + cache_hit, estimated_new, breakdown): + """HEAVY request: D direct-RDMA-reads cached KV from C_s, then does + local prefill for new tokens + decode. C_s's scheduler is NOT involved. """ request_id = headers.get("X-Request-Id", "") - estimated_new = breakdown.get("estimated_new_tokens", 0) - cache_hit = breakdown.get("cache_hit", 0) - p_prefill_release = breakdown.get("p_new_tokens", estimated_new) - # Step 1: C_s exports cached KV blocks - # Send TRUNCATED prompt (only cached portion) so C_s does 0 compute - # (full prefix cache hit), then pushes cached blocks to Mooncake. - breakdown["t_export_sent"] = _time.monotonic() - export_ok = False - - # Truncate prompt to cached portion (aligned to BLOCK_SIZE) + # Align cache_hit to block boundary for remote_num_tokens cached_tokens = (cache_hit // BLOCK_SIZE) * BLOCK_SIZE - if cached_tokens > 0 and token_ids: - export_prompt = token_ids[:cached_tokens] - else: - export_prompt = token_ids + breakdown["t_offload_sent"] = _time.monotonic() - try: - export_data = { - "model": req_data.get("model", ""), - "prompt": export_prompt, - "max_tokens": 1, - "temperature": 0, - "stream": False, - "kv_transfer_params": { - "do_remote_decode": True, - "do_remote_prefill": False, - "transfer_id": "xfer-" + request_id, - }, - } - - p_headers = {**headers, "X-data-parallel-rank": "0"} - resp = await asyncio.wait_for( - p_inst.client.post(api, json=export_data, headers=p_headers), - timeout=PREFILL_TIMEOUT_S, - ) - resp.raise_for_status() - await resp.aclose() - breakdown["t_export_done"] = _time.monotonic() - breakdown["exported_tokens"] = cached_tokens if cached_tokens > 0 else len(export_prompt) - export_ok = True - except Exception as e: - breakdown["t_export_done"] = _time.monotonic() - breakdown["export_error"] = str(e) - finally: - p_inst.ongoing_tokens -= input_length - p_inst.pending_prefill_tokens -= p_prefill_release - p_inst.num_requests -= 1 - p_inst.active_p_offloads = max(0, p_inst.active_p_offloads - 1) - - if not export_ok: - breakdown["route_class"] = "HEAVY_COLO_FALLBACK" - d_inst.pending_prefill_tokens += estimated_new - - async def generate_fallback(): - prefill_done = False - try: - async with d_inst.client.stream("POST", api, json=req_data, headers=headers) as resp: - resp.raise_for_status() - async for chunk in resp.aiter_bytes(): - if not prefill_done: - d_inst.pending_prefill_tokens -= estimated_new - d_inst.ongoing_decode_tokens += input_length - breakdown["t_first_token"] = _time.monotonic() - prefill_done = True - yield chunk - d_inst.record_prefix(token_ids) - finally: - if not prefill_done: - d_inst.pending_prefill_tokens -= estimated_new - else: - d_inst.ongoing_decode_tokens -= input_length - d_inst.ongoing_tokens -= input_length - d_inst.num_requests -= 1 - breakdown["t_done"] = _time.monotonic() - _breakdown_log.append(breakdown) - - return StreamingResponse(generate_fallback(), media_type="text/event-stream") - - # Step 2: D pulls cached blocks + does local prefill for new tokens + decodes - exported_tokens = breakdown.get("exported_tokens", 0) - d_inst.pending_prefill_tokens += estimated_new - breakdown["t_decode_sent"] = _time.monotonic() - - parsed = urllib.parse.urlparse(str(p_inst.client.base_url)) - bootstrap_addr = "http://%s:%s" % (parsed.hostname, p_inst.bootstrap_port) + parsed = urllib.parse.urlparse(str(c_inst.client.base_url)) + bootstrap_addr = "http://%s:%s" % (parsed.hostname, c_inst.bootstrap_port) + # Send full prompt to D with direct_read flag decode_data = req_data.copy() decode_data["kv_transfer_params"] = { "do_remote_decode": False, "do_remote_prefill": True, + "direct_read": True, "remote_bootstrap_addr": bootstrap_addr, - "remote_engine_id": p_inst.engine_id.get(0, ""), + "remote_engine_id": c_inst.engine_id.get(0, ""), "transfer_id": "xfer-" + request_id, - "remote_num_tokens": exported_tokens, + "remote_num_tokens": cached_tokens, } async def generate(): @@ -551,6 +465,7 @@ async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length, d_inst.ongoing_decode_tokens -= input_length d_inst.ongoing_tokens -= input_length d_inst.num_requests -= 1 + c_inst.active_p_offloads = max(0, c_inst.active_p_offloads - 1) breakdown["t_done"] = _time.monotonic() _breakdown_log.append(breakdown) diff --git a/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 2e62d26..b2bec04 100644 --- a/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -108,6 +108,9 @@ class PullReqMeta: expire_time: float = float("inf") # Designed for one D pairing to multiple P pull_tasks_count: int = 0 + # Direct RDMA read: D reads from C's GPU memory without C's scheduler + direct_read: bool = False + block_hashes: list[bytes] = field(default_factory=list) @dataclass @@ -124,11 +127,12 @@ class SendBlockMeta: class MooncakeConnectorMetadata(KVConnectorMetadata): def __init__(self): - # Use (engine_id, dp_rank) to group reqs with same dp. - # See comments in MooncakeBootstrapServer. self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict) self.reqs_to_send: dict[ReqId, tuple[TransferId, list[int]]] = {} self.reqs_not_processed: set[TransferId] = set() + # Hash table sync: scheduler → worker (for direct RDMA read) + self.hash_table_updates: dict[str, int] = {} # hex hash → block_id + self.hash_table_removals: set[str] = set() def add_new_req( self, @@ -136,6 +140,7 @@ class MooncakeConnectorMetadata(KVConnectorMetadata): local_block_ids: list[int], kv_transfer_params: dict[str, Any], load_remote_cache: bool = True, + block_hashes: list[bytes] | None = None, ): transfer_id = kv_transfer_params["transfer_id"] if load_remote_cache: @@ -146,6 +151,8 @@ class MooncakeConnectorMetadata(KVConnectorMetadata): remote_engine_id=remote_engine_id, remote_bootstrap_addr=kv_transfer_params["remote_bootstrap_addr"], transfer_id=transfer_id, + direct_read=bool(kv_transfer_params.get("direct_read")), + block_hashes=block_hashes or [], ) else: self.reqs_to_send[request_id] = (transfer_id, local_block_ids) @@ -173,6 +180,10 @@ class MooncakeConnector(KVConnectorBase_V1): self.connector_scheduler = None self.connector_worker = MooncakeConnectorWorker(vllm_config, self.engine_id) + def set_block_pool(self, block_pool): + if self.connector_scheduler is not None: + self.connector_scheduler.set_block_pool(block_pool) + ############################################################ # Scheduler Side Methods ############################################################ @@ -250,6 +261,8 @@ class MooncakeConnectorScheduler: def __init__(self, vllm_config: VllmConfig, engine_id: str): self.vllm_config = vllm_config + self._block_pool = None + self._known_hash_keys: set = set() assert vllm_config.kv_transfer_config self.is_kv_producer: bool = ( @@ -260,14 +273,14 @@ class MooncakeConnectorScheduler: ) logger.info("Initializing Mooncake Transfer Engine Scheduler %s", engine_id) - # Requests that need to start recv/send. - # New requests are added by update_state_after_alloc in - # the scheduler. Used to make metadata passed to Worker. self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {} self._reqs_need_send: dict[ReqId, tuple[Request, list[int]]] = {} - # Reqs to remove from processed set because they're not to send after - # remote prefill or aborted. self._reqs_not_processed: set[TransferId] = set() + # Per-request block hashes for direct_read passthrough + self._req_block_hashes: dict[ReqId, list[bytes]] = {} + + def set_block_pool(self, block_pool): + self._block_pool = block_pool def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int @@ -348,13 +361,21 @@ class MooncakeConnectorScheduler: else: local_block_ids = [] self._reqs_need_recv[request.request_id] = (request, local_block_ids) + # Capture block hashes for direct_read + if params.get("direct_read") and hasattr(request, "block_hashes"): + block_size = self.vllm_config.cache_config.block_size + num_remote_blocks = ( + (num_external_tokens + block_size - 1) // block_size + ) + self._req_block_hashes[request.request_id] = [ + bytes(h) for h in request.block_hashes[:num_remote_blocks] + ] else: logger.warning( "Got invalid KVTransferParams: %s. This " "request will not utilize KVTransfer", params, ) - # Only trigger 1 KV transfer per request. params["do_remote_prefill"] = False if params.get("do_remote_decode"): @@ -371,7 +392,6 @@ class MooncakeConnectorScheduler: ) -> KVConnectorMetadata: meta = MooncakeConnectorMetadata() - # Loop through scheduled reqs and convert to PullReqMeta. if not self.is_kv_producer: for req_id, (req, block_ids) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None @@ -379,9 +399,30 @@ class MooncakeConnectorScheduler: request_id=req_id, local_block_ids=block_ids, kv_transfer_params=req.kv_transfer_params, + block_hashes=self._req_block_hashes.pop(req_id, None), ) self._reqs_need_recv.clear() + # Sync hash table to worker for direct RDMA read block lookups + if self._block_pool is not None: + cache = self._block_pool.cached_block_hash_to_block._cache + current_keys = set(cache.keys()) + new_keys = current_keys - self._known_hash_keys + removed_keys = self._known_hash_keys - current_keys + if new_keys or removed_keys: + from vllm.v1.core.kv_cache_utils import get_block_hash + for k in new_keys: + block = cache[k] + if isinstance(block, dict): + bid = next(iter(block.values())).block_id + else: + bid = block.block_id + meta.hash_table_updates[get_block_hash(k).hex()] = bid + meta.hash_table_removals = { + get_block_hash(k).hex() for k in removed_keys + } + self._known_hash_keys = current_keys.copy() + if not self.is_kv_consumer: for req_id, (req, block_ids) in self._reqs_need_send.items(): assert req.kv_transfer_params is not None @@ -971,7 +1012,6 @@ class MooncakeConnectorWorker: "registered num_blocks=%d block_len=%d", self.num_blocks, self.block_len ) - # No need to launch server for D node. if self.is_kv_consumer: return @@ -979,7 +1019,13 @@ class MooncakeConnectorWorker: asyncio.run_coroutine_threadsafe( self._mooncake_sender_listener(ready_event), self.sender_loop ) - ready_event.wait() # Wait for listener ZMQ socket to be ready. + ready_event.wait() + + if self.bootstrap_server is not None: + self.bootstrap_server.set_worker_kv_info( + self.kv_caches_base_addr, self.block_len, + self.hostname, self.rpc_port, + ) async def fetch_finished_recving_reqs(self) -> set[ReqId]: finished_recving_reqs = self.finished_recving_reqs @@ -1197,6 +1243,93 @@ class MooncakeConnectorWorker: self.receive_kv(remote_engine_id, pull_metas) + async def _start_direct_read( + self, reqs_by_engine: dict[EngineId, dict[ReqId, PullReqMeta]] + ): + """Direct RDMA read: D reads cached KV blocks from C's GPU memory + without involving C's scheduler. + """ + for _engine_id, pull_metas in reqs_by_engine.items(): + for req_id, pm in pull_metas.items(): + asyncio.create_task( + self._direct_read_single(req_id, pm) + ) + + async def _direct_read_single(self, req_id: ReqId, pm: PullReqMeta): + pin_token = f"dr-{req_id}-{self.tp_rank}" + bootstrap_url = pm.remote_bootstrap_addr + + try: + # 1. Query C's bootstrap for block mapping + async with httpx.AsyncClient(timeout=30) as client: + resp = await client.post( + f"{bootstrap_url}/query_blocks", + json={ + "block_hashes": [h.hex() for h in pm.block_hashes], + "pin_token": pin_token, + }, + ) + resp.raise_for_status() + mapping = resp.json() + + remote_block_ids = [b for b in mapping["block_ids"] if b is not None] + num_matched = len(remote_block_ids) + + if num_matched == 0: + logger.debug("direct_read %s: no cache hit on remote", req_id) + self.finished_recving_reqs.add(req_id) + return + + # 2. Compute RDMA addresses + local_block_ids = pm.local_block_ids[:num_matched] + remote_base_addrs = mapping["kv_caches_base_addr"] + remote_block_len = mapping["block_len"] + remote_session = f"{mapping['hostname']}:{mapping['rpc_port']}" + + src_ptrs: list[int] = [] + dst_ptrs: list[int] = [] + lengths: list[int] = [] + + for local_layer_addr, remote_layer_addr in zip( + self.kv_caches_base_addr, remote_base_addrs + ): + grp_local, grp_remote = group_concurrent_contiguous( + local_block_ids, remote_block_ids + ) + for gl, gr in zip(grp_local, grp_remote): + src_ptrs.append(remote_layer_addr + gr[0] * remote_block_len) + dst_ptrs.append(local_layer_addr + gl[0] * self.block_len) + lengths.append(remote_block_len * len(gr)) + + # 3. RDMA READ (D pulls from C's GPU memory) + logger.debug( + "direct_read %s: reading %d blocks from %s", + req_id, num_matched, remote_session, + ) + ret = self.engine.batch_transfer_sync_read( + remote_session, src_ptrs, dst_ptrs, lengths, + ) + + if ret != 0: + logger.error("direct_read %s: RDMA read failed ret=%d", req_id, ret) + return + + logger.debug("direct_read %s: success (%d blocks)", req_id, num_matched) + self.finished_recving_reqs.add(req_id) + + except Exception as e: + logger.error("direct_read %s failed: %s", req_id, e) + finally: + # 4. Unpin blocks on C + try: + async with httpx.AsyncClient(timeout=10) as client: + await client.post( + f"{bootstrap_url}/unpin_blocks", + json={"pin_token": pin_token}, + ) + except Exception: + pass + async def _start_load_kv( self, reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] ): @@ -1237,11 +1370,34 @@ class MooncakeConnectorWorker: assert not send_meta.ready.is_set() def start_load_kv(self, metadata: MooncakeConnectorMetadata): - if not self.is_kv_producer and metadata.reqs_to_recv: - asyncio.run_coroutine_threadsafe( - self._start_load_kv(metadata.reqs_to_recv), self.receiver_loop + # Sync hash table to bootstrap server (for direct RDMA read queries) + if self.bootstrap_server is not None and ( + metadata.hash_table_updates or metadata.hash_table_removals + ): + self.bootstrap_server.update_hash_table( + metadata.hash_table_updates, metadata.hash_table_removals ) + if not self.is_kv_producer and metadata.reqs_to_recv: + # Split direct_read vs normal pull requests + direct_reqs: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict) + normal_reqs: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict) + for engine_id, pull_metas in metadata.reqs_to_recv.items(): + for req_id, pm in pull_metas.items(): + if pm.direct_read: + direct_reqs[engine_id][req_id] = pm + else: + normal_reqs[engine_id][req_id] = pm + + if normal_reqs: + asyncio.run_coroutine_threadsafe( + self._start_load_kv(normal_reqs), self.receiver_loop + ) + if direct_reqs: + asyncio.run_coroutine_threadsafe( + self._start_direct_read(direct_reqs), self.receiver_loop + ) + if not self.is_kv_consumer and ( metadata.reqs_to_send or metadata.reqs_not_processed ): diff --git a/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py b/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py index d1a9946..231cd3f 100644 --- a/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py +++ b/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py @@ -32,10 +32,28 @@ class EngineEntry: worker_addr: dict[int, dict[int, WorkerAddr]] +class QueryBlocksRequest(BaseModel): + block_hashes: list[str] # hex-encoded BlockHash values + pin_token: str + + +class QueryBlocksResponse(BaseModel): + block_ids: list[int | None] # None = cache miss (prefix match stops) + kv_caches_base_addr: list[int] + block_len: int + hostname: str + rpc_port: int + + +class UnpinBlocksRequest(BaseModel): + pin_token: str + + class MooncakeBootstrapServer: """ A centralized server running on the global rank 0 prefiller worker. Prefiller workers register their connection info (IP, port, ranks) here. + Also serves block mapping queries for direct RDMA read. """ def __init__(self, vllm_config: VllmConfig, host: str, port: int): @@ -48,13 +66,19 @@ class MooncakeBootstrapServer: self.server_thread: threading.Thread | None = None self.server: uvicorn.Server | None = None + # Direct RDMA read support: block hash → block_id mapping + self._hash_table: dict[str, int] = {} # hex BlockHash → block_id + self._kv_info: dict | None = None # set by worker at register_kv_caches + self._pinned: dict[str, list[int]] = {} # pin_token → block_ids + def __del__(self): self.shutdown() def _register_routes(self): - # All methods are async. No need to use lock to protect data. self.app.post("/register")(self.register_worker) self.app.get("/query", response_model=dict[int, EngineEntry])(self.query) + self.app.post("/query_blocks")(self.query_blocks) + self.app.post("/unpin_blocks")(self.unpin_blocks) def start(self): if self.server_thread: @@ -125,3 +149,54 @@ class MooncakeBootstrapServer: async def query(self) -> dict[int, EngineEntry]: return self.workers + + def set_worker_kv_info( + self, + kv_caches_base_addr: list[int], + block_len: int, + hostname: str, + rpc_port: int, + ): + self._kv_info = { + "kv_caches_base_addr": kv_caches_base_addr, + "block_len": block_len, + "hostname": hostname, + "rpc_port": rpc_port, + } + + def update_hash_table( + self, + updates: dict[str, int], + removals: set[str], + ): + for k in removals: + self._hash_table.pop(k, None) + self._hash_table.update(updates) + + async def query_blocks(self, req: QueryBlocksRequest): + if self._kv_info is None: + raise HTTPException(503, "Worker KV info not registered yet") + + block_ids: list[int | None] = [] + pinned_ids: list[int] = [] + for h in req.block_hashes: + bid = self._hash_table.get(h) + if bid is not None: + block_ids.append(bid) + pinned_ids.append(bid) + else: + block_ids.append(None) + break # prefix match: stop at first miss + + self._pinned[req.pin_token] = pinned_ids + return QueryBlocksResponse( + block_ids=block_ids, + kv_caches_base_addr=self._kv_info["kv_caches_base_addr"], + block_len=self._kv_info["block_len"], + hostname=self._kv_info["hostname"], + rpc_port=self._kv_info["rpc_port"], + ) + + async def unpin_blocks(self, req: UnpinBlocksRequest): + self._pinned.pop(req.pin_token, None) + return {"status": "ok"} diff --git a/third_party/vllm/vllm/v1/core/sched/scheduler.py b/third_party/vllm/vllm/v1/core/sched/scheduler.py index 5b2b40d..5350c90 100644 --- a/third_party/vllm/vllm/v1/core/sched/scheduler.py +++ b/third_party/vllm/vllm/v1/core/sched/scheduler.py @@ -234,6 +234,8 @@ class Scheduler(SchedulerInterface): hash_block_size=self.block_size, metrics_collector=self.kv_metrics_collector, ) + if self.connector is not None and hasattr(self.connector, "set_block_pool"): + self.connector.set_block_pool(self.kv_cache_manager.block_pool) self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER