[KVCache][Feature] implement storage prefetch ZMQ pipeline in Scheduler and Worker

Storage prefetch (Storage → CPU) previously had no runtime execution path in
Scheduler/Worker: the Scheduler prepared host-block metadata but the actual
data transfer was never triggered.  Workers also had no mechanism to receive
prefetch commands or report completion, leaving LOADING_FROM_STORAGE blocks
permanently stuck and never promoted to HOST.

- Add `_prefetch_node_map: Dict[int, BlockNode]` to track in-flight blocks by
  host_block_id for O(1) status lookup.
- `prepare_prefetch_metadata`: register returned nodes into `_prefetch_node_map`.
- New `update_storage_blocks_to_host(host_block_ids)`: transition
  LOADING_FROM_STORAGE → HOST after all TP workers confirm transfer done.
- New `abort_prefetch_blocks(host_block_ids)`: remove nodes from RadixTree and
  release host pool blocks on transfer failure.

- Add per-worker ZMQ PUSH/PULL servers (`_prefetch_cmd_servers`,
  `_prefetch_done_servers`), one pair per TP worker, keyed by local_rank.
- `_init_prefetch_zmq_servers()`: initialize servers at startup when storage
  backend is configured.
- `_prefetch_storage_cache()`: after inserting host blocks, serialize
  `StorageMetadata` and broadcast to all TP workers via ZMQ PUSH; then poll
  PULL done sockets until all workers reply, call
  `update_storage_blocks_to_host` on success or `abort_prefetch_blocks` on
  failure.

- Add `receive_pyobj_once(block=False)`: non-blocking (or blocking) receive
  helper returning `(error, data)` tuple; used by Scheduler to poll done
  messages and by Worker in the prefetch loop.

- Add `init_prefetch_zmq_clients()`: connect ZMQ PULL/PUSH clients to
  Scheduler servers for this worker's local_rank; start daemon
  `_prefetch_loop` thread.
- `_prefetch_loop()`: background thread receiving `StorageMetadata` commands,
  calling `cache_controller.prefetch_from_storage`, waiting for
  `AsyncTaskHandler.wait`, and replying with ok/error status.

- Add `TestUpdateStorageBlocksToHost` with 6 test cases covering:
  status transition, multi-block, unknown id, empty list, wrong status,
  and initial-empty-map assertions.

No additional build steps. Enable storage prefetch via existing config:
```bash
python -m fastdeploy.entrypoints.openai.api_server \
  --kvcache-storage-backend <backend> \
  --enable-prefix-caching \
  ...
```

Author: kevincheng2

fastdeploy/cache_manager/v1/cache_manager.py

@@ -107,6 +107,10 @@ def __init__(
         self._pending_backup: List[Tuple[List[BlockNode], List[int]]] = []
         self._pending_block_ids: List[int] = []
 
+        # Mapping from host_block_id -> BlockNode for LOADING_FROM_STORAGE blocks,
+        # used to quickly update status to HOST once prefetch completes.
+        self._prefetch_node_map: Dict[int, BlockNode] = {}
+
         # Storage scheduler (create using factory method if backend is configured)
         self._storage_scheduler = create_storage_scheduler(self.cache_config)
 
@@ -1004,11 +1008,79 @@ def prepare_prefetch_metadata(
                 if wasted_block_ids:
                     self._host_pool.release(wasted_block_ids)
 
+                # Register nodes in prefetch_node_map for fast status update on done
+                for node in prefetch_nodes:
+                    self._prefetch_node_map[node.block_id] = node
+
                 return prefetch_nodes
         except Exception as e:
             logger.error(f"prepare_prefetch_metadata error: {e}, {str(traceback.format_exc())}")
             return []
 
+    def update_storage_blocks_to_host(self, host_block_ids: List[int]) -> None:
+        """
+        Mark storage-prefetched blocks as HOST after data transfer completes.
+
+        Called by Scheduler when all TP workers report prefetch done for a batch
+        of blocks. Transitions block status LOADING_FROM_STORAGE → HOST so that
+        these blocks become eligible for swap-in scheduling.
+
+        Args:
+            host_block_ids: List of host block IDs that finished loading.
+        """
+        if not host_block_ids:
+            return
+        try:
+            with self._lock:
+                updated = 0
+                for block_id in host_block_ids:
+                    node = self._prefetch_node_map.pop(block_id, None)
+                    if node is None:
+                        logger.warning(
+                            f"[StoragePrefetch] update_storage_blocks_to_host: "
+                            f"block_id={block_id} not found in prefetch_node_map"
+                        )
+                        continue
+                    if node.cache_status == CacheStatus.LOADING_FROM_STORAGE:
+                        node.cache_status = CacheStatus.HOST
+                        updated += 1
+                    else:
+                        logger.warning(
+                            f"[StoragePrefetch] update_storage_blocks_to_host: "
+                            f"block_id={block_id} unexpected status={node.cache_status}"
+                        )
+                logger.info(
+                    f"[StoragePrefetch] update_storage_blocks_to_host: "
+                    f"requested={len(host_block_ids)}, updated={updated}"
+                )
+        except Exception as e:
+            logger.error(f"update_storage_blocks_to_host error: {e}, {str(traceback.format_exc())}")
+
+    def abort_prefetch_blocks(self, host_block_ids: List[int]) -> None:
+        """
+        Abort in-flight prefetch blocks on failure.
+
+        Removes nodes from the prefetch_node_map, deletes them from the RadixTree,
+        and releases their host pool blocks. Called when the storage→CPU transfer
+        fails so that LOADING_FROM_STORAGE blocks do not leak.
+
+        Args:
+            host_block_ids: List of host block IDs whose prefetch should be aborted.
+        """
+        if not host_block_ids:
+            return
+        try:
+            with self._lock:
+                for block_id in host_block_ids:
+                    node = self._prefetch_node_map.pop(block_id, None)
+                    if node is None:
+                        continue
+                    self._radix_tree._remove_node_from_tree(node)
+                self._host_pool.release(host_block_ids)
+            logger.warning(f"[StoragePrefetch] abort_prefetch_blocks: released {len(host_block_ids)} host blocks")
+        except Exception as e:
+            logger.error(f"abort_prefetch_blocks error: {e}, {str(traceback.format_exc())}")
+
     # ============ Reset Methods ============
 
     def reset_cache(self) -> bool:

fastdeploy/engine/sched/resource_manager_v1.py

@@ -22,17 +22,18 @@
 from collections.abc import Iterable
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
-from typing import List, Union
+from typing import Dict, List, Set, Union
 
 import numpy as np
 import paddle
+import zmq
 
 from fastdeploy import envs
 from fastdeploy.cache_manager.multimodal_cache_manager import (
     EncoderCacheManager,
     ProcessorCacheManager,
 )
-from fastdeploy.cache_manager.v1.metadata import CacheSwapMetadata
+from fastdeploy.cache_manager.v1.metadata import CacheSwapMetadata, StorageMetadata
 from fastdeploy.engine.request import (
     BatchRequest,
     ImagePosition,
@@ -44,6 +45,7 @@
 from fastdeploy.engine.resource_manager import ResourceManager
 from fastdeploy.input.utils import IDS_TYPE_FLAG
 from fastdeploy.inter_communicator import IPCSignal
+from fastdeploy.inter_communicator.zmq_server import ZmqIpcServer
 from fastdeploy.metrics.metrics import main_process_metrics
 from fastdeploy.multimodal.hasher import MultimodalHasher
 from fastdeploy.platforms import current_platform
@@ -252,6 +254,16 @@ def __init__(self, max_num_seqs, config, tensor_parallel_size, splitwise_role, l
         # Scheduler-side requests that have not been moved into resource manager waiting queue yet.
         self.scheduler_unhandled_request_num = 0
 
+        # ---- Storage Prefetch ZMQ channels (Scheduler side) ----
+        # Initialized only when storage backend is configured.
+        # One PUSH cmd socket + one PULL done socket per worker local_rank.
+        # local_rank = dp_rank * tp_size + tp_rank
+        self._prefetch_cmd_servers: Dict[int, ZmqIpcServer] = {}
+        self._prefetch_done_servers: Dict[int, ZmqIpcServer] = {}
+
+        if self.config.cache_config.kvcache_storage_backend and self.enable_cache_manager_v1:
+            self._init_prefetch_zmq_servers()
+
     def allocated_slots(self, request: Request):
         return len(request.block_tables) * self.config.cache_config.block_size
 
@@ -1264,6 +1276,29 @@ def apply_async_preprocess(self, request: Request) -> None:
                 self.async_preprocess_pool.submit(self._prefetch_storage_cache, request)
             )
 
+    def _init_prefetch_zmq_servers(self) -> None:
+        """
+        Initialize per-worker-rank ZMQ PUSH/PULL sockets for storage prefetch.
+
+        Called once during __init__ when storage backend is enabled.
+        Creates:
+          - prefetch_cmd_server[local_rank]:  PUSH → Worker (send StorageMetadata)
+          - prefetch_done_server[local_rank]: PULL ← Worker (receive done notification)
+
+        local_rank = dp_rank * tp_size + tp_rank, covers all workers in this DP group.
+        """
+        tp_size = self.config.parallel_config.tensor_parallel_size
+        dp_rank = self.config.parallel_config.local_data_parallel_id
+        port = self.config.parallel_config.local_engine_worker_queue_port
+
+        for tp_rank in range(tp_size):
+            local_rank = dp_rank * tp_size + tp_rank
+            cmd_name = f"prefetch_cmd_rank{local_rank}_{port}"
+            done_name = f"prefetch_done_rank{local_rank}_{port}"
+            self._prefetch_cmd_servers[local_rank] = ZmqIpcServer(cmd_name, zmq.PUSH)
+            self._prefetch_done_servers[local_rank] = ZmqIpcServer(done_name, zmq.PULL)
+            llm_logger.info(f"[StoragePrefetch] init ZMQ servers: cmd={cmd_name}, done={done_name}")
+
     def _prefetch_storage_cache(self, request: Request) -> None:
         """
         Asynchronously prefetch KV cache blocks from storage to host memory.
@@ -1274,29 +1309,118 @@ def _prefetch_storage_cache(self, request: Request) -> None:
           2. Allocate host blocks for them.
           3. Insert those blocks into the RadixTree with LOADING_FROM_STORAGE status.
 
-        The actual data transfer (storage → host memory) is handled by the Worker
-        via cache_controller.prefetch_from_storage once the batch is dispatched.
+        Then immediately sends a StorageMetadata message to all TP Workers via ZMQ,
+        so Workers can start the actual storage→CPU transfer independently of forward.
 
         Args:
             request: The request to prefetch cache for.
         """
+        host_block_ids: List[int] = []
         try:
             if not self.cache_manager.enable_prefix_caching:
                 return
             llm_logger.debug(f"[StoragePrefetch] start async prefetch for request_id={request.request_id}")
             self.cache_manager.match_prefix(request, skip_storage=False)
             match_result = request.match_result
-            if match_result is not None:
-                request.match_result = None
+            request.match_result = None
+            if match_result is None or match_result.matched_storage_nums == 0:
+                return
 
-                llm_logger.info(
-                    f"[StoragePrefetch] request_id={request.request_id} "
-                    f"storage_matched={match_result.matched_storage_nums} blocks"
+            # Collect host_block_ids and hash_values from matched storage nodes
+            storage_nodes = match_result.storage_nodes
+            host_block_ids = [node.block_id for node in storage_nodes]
+            hash_values = [node.hash_value for node in storage_nodes]
+
+            llm_logger.info(
+                f"[StoragePrefetch] request_id={request.request_id} "
+                f"storage_matched={match_result.matched_storage_nums} blocks, "
+                f"host_block_ids={host_block_ids}"
+            )
+
+            if not self._prefetch_cmd_servers:
+                return
+
+            metadata = StorageMetadata(
+                hash_values=hash_values,
+                block_ids=host_block_ids,
+                direction="load",
+            )
+
+            # Build the payload with request_id for done matching
+            payload = {
+                "request_id": request.request_id,
+                "metadata": metadata,
+            }
+
+            # Send to all TP workers in this DP group
+            for local_rank, cmd_server in self._prefetch_cmd_servers.items():
+                try:
+                    cmd_server.send_pyobj(payload)
+                except Exception as e:
+                    llm_logger.error(f"[StoragePrefetch] failed to send cmd to rank={local_rank}: {e}")
+
+            # Block in this thread until all TP workers report done.
+            # This mirrors _download_features: the future is considered complete only
+            # when the actual storage→CPU transfer has finished on every worker.
+            expected_count = len(self._prefetch_cmd_servers)
+            done_ranks: Set[int] = set()
+            failed_ranks: Set[int] = set()
+            poll_interval = 0.001  # 1ms
+
+            while len(done_ranks) + len(failed_ranks) < expected_count:
+                for local_rank, done_server in self._prefetch_done_servers.items():
+                    if local_rank in done_ranks or local_rank in failed_ranks:
+                        continue
+                    err, msg = done_server.receive_pyobj_once(block=False)
+                    if err is not None:
+                        llm_logger.warning(
+                            f"[StoragePrefetch] done_server rank={local_rank} socket error: {err}, "
+                            f"request_id={request.request_id}"
+                        )
+                        failed_ranks.add(local_rank)
+                        continue
+                    if msg is None:
+                        continue
+                    recv_req_id = msg.get("request_id", "")
+                    if recv_req_id != request.request_id:
+                        # Message for a different request; skip and let that request's
+                        # thread poll its own done message. This should not normally happen
+                        # since each worker sends done to the same socket, but guard anyway.
+                        llm_logger.warning(
+                            f"[StoragePrefetch] rank={local_rank} received done for unexpected "
+                            f"request_id={recv_req_id}, expected={request.request_id}, skipping"
+                        )
+                        continue
+                    if msg.get("status") != "ok":
+                        llm_logger.warning(
+                            f"[StoragePrefetch] rank={local_rank} worker reported prefetch failure for "
+                            f"request_id={request.request_id}: {msg.get('error')}"
+                        )
+                        failed_ranks.add(local_rank)
+                        continue
+                    done_ranks.add(local_rank)
+
+                if len(done_ranks) + len(failed_ranks) < expected_count:
+                    time.sleep(poll_interval)
+
+            if failed_ranks:
+                llm_logger.warning(
+                    f"[StoragePrefetch] request_id={request.request_id} prefetch failed on "
+                    f"ranks={failed_ranks}, aborting {len(host_block_ids)} host blocks"
                 )
-            # TODO: check if any of the block is still LOADING_FROM_STORAGE, if so, request.async_process_futures.append(self._prefetch_storage_cache)
+                self.cache_manager.abort_prefetch_blocks(host_block_ids)
+                return
+
+            # All workers done successfully: update CacheManager block status to HOST
+            self.cache_manager.update_storage_blocks_to_host(host_block_ids)
+            llm_logger.info(
+                f"[StoragePrefetch] request_id={request.request_id} all {expected_count} TP workers done, "
+                f"updated {len(host_block_ids)} blocks to HOST"
+            )
 
         except Exception as e:
             llm_logger.error(f"[StoragePrefetch] request_id={request.request_id} error: {e}")
+            self.cache_manager.abort_prefetch_blocks(host_block_ids)
 
     def _has_features_info(self, task):
         inputs = task.multimodal_inputs

fastdeploy/inter_communicator/zmq_client.py

@@ -143,6 +143,28 @@ def recv_pyobj(self, flags: int = 0):
                 return envelope["data"]
         return envelope
 
+    def receive_pyobj_once(self, block=False):
+        """
+        Receive a single Pickle-serializable message from the socket.
+
+        Args:
+            block: If True, block until a message arrives. If False, return immediately.
+
+        Returns:
+            Tuple of (error, data). error is None on success, data is None if no message.
+        """
+        self._ensure_socket()
+        if self.socket is None or self.socket.closed:
+            return "zmq socket has closed", None
+        try:
+            flags = 0 if block else zmq.NOBLOCK
+            return None, self.recv_pyobj(flags=flags)
+        except zmq.Again:
+            return None, None
+        except Exception as e:
+            llm_logger.warning(f"[ZmqClient] receive_pyobj_once error: {e}")
+            return str(e), None
+
     @abstractmethod
     def close(self):
         pass