[HiCache] Optimize L2 mem allocation when cache miss in L3 (#19320)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

benchmark/hf3fs/bench_zerocopy.py

@@ -114,7 +114,6 @@ print(f"{tok-tik:.6f} s")
 operations = [
     PrefetchOperation(
         f"{i}",
-        torch.tensor(list(range(i, i + op_size))),
         list(range(i, i + op_size)),
         f"{i}",
     )
@@ -122,6 +121,7 @@ operations = [
 ]
 
 for operation in operations:
+    operation.host_indices = torch.tensor(operation.token_ids)
     operation.hash_value = [
         f"{j}"
         for j in range(

python/sglang/srt/managers/cache_controller.py

@@ -170,7 +170,7 @@ class StorageOperation:
 
     def __init__(
         self,
-        host_indices: torch.Tensor,
+        host_indices: Optional[torch.Tensor],
         token_ids: List[int],
         last_hash: Optional[str] = None,
         hash_value: Optional[List[str]] = None,
@@ -194,7 +194,6 @@ class PrefetchOperation(StorageOperation):
     def __init__(
         self,
         request_id: str,
-        host_indices: torch.Tensor,
         token_ids: List[int],
         last_hash: Optional[str] = None,
         prefix_keys: Optional[List[str]] = None,
@@ -203,9 +202,10 @@ class PrefetchOperation(StorageOperation):
 
         self._lock = threading.Lock()
         self._terminated_flag = False
+        self.storage_hit_count = 0
         self.start_time = time.monotonic()
 
-        super().__init__(host_indices, token_ids, last_hash, prefix_keys=prefix_keys)
+        super().__init__(None, token_ids, last_hash, prefix_keys=prefix_keys)
 
     def increment(self, num_tokens: int):
         with self._lock:
@@ -374,6 +374,7 @@ class HiCacheController:
         self.prefetch_queue = Queue()
         self.backup_queue = Queue()
 
+        self.prefetch_hit_queue: Queue[StorageOperation] = Queue()
         self.prefetch_revoke_queue: Queue[str] = Queue()
         self.ack_backup_queue: Queue[StorageOperation] = Queue()
         self.host_mem_release_queue: Queue[torch.Tensor] = Queue()
@@ -652,6 +653,7 @@ class HiCacheController:
             self.prefetch_queue.queue.clear()
             self.backup_queue.queue.clear()
             self.prefetch_revoke_queue.queue.clear()
+            self.prefetch_hit_queue.queue.clear()
             self.ack_backup_queue.queue.clear()
             self.host_mem_release_queue.queue.clear()
             self.prefetch_tokens_occupied = 0
@@ -898,7 +900,6 @@ class HiCacheController:
     def prefetch(
         self,
         request_id: str,
-        host_indices: torch.Tensor,
         new_input_tokens: List[int],
         last_hash: Optional[str] = None,
         prefix_keys: Optional[List[str]] = None,
@@ -907,7 +908,7 @@ class HiCacheController:
         Prefetch KV caches from storage backend to host memory.
         """
         operation = PrefetchOperation(
-            request_id, host_indices, new_input_tokens, last_hash, prefix_keys
+            request_id, new_input_tokens, last_hash, prefix_keys
         )
         self.prefetch_queue.put(operation)
         return operation
@@ -1057,7 +1058,10 @@ class HiCacheController:
                 operation = self.prefetch_queue.get(block=True, timeout=1)
                 if operation is None:
                     continue
-                hash_value, storage_hit_count = self._storage_hit_query(operation)
+                if operation.is_terminated():
+                    hash_value, storage_hit_count = [], 0
+                else:
+                    hash_value, storage_hit_count = self._storage_hit_query(operation)
                 storage_hit_count_tensor = torch.tensor(
                     storage_hit_count, dtype=torch.int
                 )
@@ -1069,23 +1073,16 @@ class HiCacheController:
                 if storage_hit_count < self.prefetch_threshold:
                     # not to prefetch if not enough benefits
                     self.prefetch_revoke_queue.put(operation.request_id)
-                    self.append_host_mem_release(operation.host_indices)
                     logger.debug(
                         f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})."
                     )
                 else:
+                    # Record hit count, so the scheduler thread will know the exact memory to allocate
                     operation.hash_value = hash_value[
                         : (storage_hit_count // self.page_size)
                     ]
-                    # free the pre-allocated memory for pages that are not hit
-                    self.append_host_mem_release(
-                        operation.host_indices[storage_hit_count:]
-                    )
-                    operation.host_indices = operation.host_indices[:storage_hit_count]
-                    logger.debug(
-                        f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}."
-                    )
-                    self.prefetch_buffer.put(operation)
+                    operation.storage_hit_count = storage_hit_count
+                    self.prefetch_hit_queue.put(operation)
 
             except Empty:
                 continue

python/sglang/srt/mem_cache/hi_mamba_radix_cache.py

@@ -1476,14 +1476,28 @@ class HiMambaRadixCache(MambaRadixCache):
     def _drain_storage_control_queues_local(self):
         self._drain_storage_control_queues_impl(
             n_revoke=None,
+            n_storage_hit=0,
             n_backup=None,
             n_release=None,
             log_metrics=False,
         )
 
+    def _revoke_pending_prefetch(self, req_id: str):
+        info = self.ongoing_prefetch.pop(req_id, None)
+        if info is None:
+            return
+        last_host_node, token_ids, _host_indices, operation = info
+        self.prefetch_abort(operation.pool_transfers)
+        self._release_host_node(last_host_node)
+        cc = self.cache_controller
+        cc.prefetch_tokens_occupied = max(
+            0, cc.prefetch_tokens_occupied - len(token_ids)
+        )
+
     def _drain_storage_control_queues_impl(
         self,
         n_revoke: Optional[int],
+        n_storage_hit: Optional[int],
         n_backup: Optional[int],
         n_release: Optional[int],
         log_metrics: bool,
@@ -1502,14 +1516,57 @@ class HiMambaRadixCache(MambaRadixCache):
 
         def _drain_revoke():
             for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
-                info = self.ongoing_prefetch.pop(req_id, None)
-                if info is not None:
-                    last_host_node, token_ids, _, operation = info
-                    self.prefetch_abort(operation.pool_transfers)
-                    self._release_host_node(last_host_node)
-                    cc.prefetch_tokens_occupied -= len(token_ids)
-                    if cc.prefetch_tokens_occupied < 0:
-                        cc.prefetch_tokens_occupied = 0
+                self._revoke_pending_prefetch(req_id)
+
+        def _drain_and_alloc_storage_hit():
+            # The L3 hit count is now known, so reserve exactly that much host
+            # KV memory. NOTE: alloc/evict here is rank-local but deterministic
+            # across TP ranks without extra synchronization: host pool
+            # mutations only happen on the scheduler thread at lockstep points
+            # (releases are page-granular and drained by the TP-min count), so
+            # every rank reaches the same success / fallback / revoke decision.
+            for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit):
+                req_id = operation.request_id
+                info = self.ongoing_prefetch.get(req_id)
+                if info is None:
+                    # request already aborted/cleaned up, skip
+                    continue
+                if operation.is_terminated():
+                    # request was aborted while the storage query was in flight
+                    self._revoke_pending_prefetch(req_id)
+                    continue
+
+                alloc_len = operation.storage_hit_count
+                host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    self.evict_host(alloc_len)
+                    host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    # Memory-pressure fallback: a shorter page-aligned prefix.
+                    available_size = cc.mem_pool_host.available_size()
+                    alloc_len = min(
+                        operation.storage_hit_count,
+                        available_size - (available_size % self.page_size),
+                    )
+                    if alloc_len >= self.prefetch_threshold:
+                        host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    self._revoke_pending_prefetch(req_id)
+                    continue
+
+                operation.storage_hit_count = alloc_len
+                operation.hash_value = operation.hash_value[
+                    : alloc_len // self.page_size
+                ]
+                operation.host_indices = host_indices
+                last_host_node, token_ids, _, op = info
+                self.ongoing_prefetch[req_id] = (
+                    last_host_node,
+                    token_ids,
+                    host_indices,
+                    op,
+                )
+                cc.prefetch_buffer.put(operation)
 
         def _drain_backup():
             for operation in _drain_queue(cc.ack_backup_queue, n_backup):
@@ -1532,6 +1589,7 @@ class HiMambaRadixCache(MambaRadixCache):
                 cc.mem_pool_host.free(host_indices)
 
         _drain_revoke()
+        _drain_and_alloc_storage_hit()
         _drain_backup()
         _drain_release()
 
@@ -1642,6 +1700,7 @@ class HiMambaRadixCache(MambaRadixCache):
         qsizes = torch.tensor(
             [
                 cc.prefetch_revoke_queue.qsize(),
+                cc.prefetch_hit_queue.qsize(),
                 cc.ack_backup_queue.qsize(),
                 cc.host_mem_release_queue.qsize(),
             ],
@@ -1652,9 +1711,10 @@ class HiMambaRadixCache(MambaRadixCache):
                 qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
             )
 
-        n_revoke, n_backup, n_release = map(int, qsizes.tolist())
+        n_revoke, n_storage_hit, n_backup, n_release = map(int, qsizes.tolist())
         self._drain_storage_control_queues_impl(
             n_revoke=n_revoke,
+            n_storage_hit=n_storage_hit,
             n_backup=n_backup,
             n_release=n_release,
             log_metrics=True,
@@ -1710,8 +1770,6 @@ class HiMambaRadixCache(MambaRadixCache):
             return
 
         _, _, _, operation = self.ongoing_prefetch[req_id]
-        if operation.host_indices is None:
-            return
         operation.mark_terminate()
 
     def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
@@ -1756,30 +1814,9 @@ class HiMambaRadixCache(MambaRadixCache):
 
         self._protect_host_node(last_host_node, protect_mamba=False)
 
-        # Allocate host KV memory
-        host_indices = self._alloc_with_evict(
-            self.cache_controller.mem_pool_host,
-            prefetch_length,
-            self.evict_host,
-        )
-        if host_indices is None:
-            # truncate the prefetch length to the page-aligned available host size
-            available_size = self.cache_controller.mem_pool_host.available_size()
-            prefetch_length = available_size - (available_size % self.page_size)
-            if prefetch_length < self.prefetch_threshold:
-                self._release_host_node(last_host_node, release_mamba=False)
-                return
-            new_input_tokens = new_input_tokens[:prefetch_length]
-            host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
-
-        if host_indices is None:
-            self._release_host_node(last_host_node, release_mamba=False)
-            return
-
         # Allocate host mamba slot
         extra_pools = self.mamba_prefetch_alloc(new_input_tokens, last_hash)
         if extra_pools is None:
-            self.cache_controller.mem_pool_host.free(host_indices)
             self._release_host_node(last_host_node, release_mamba=False)
             return
 
@@ -1790,7 +1827,6 @@ class HiMambaRadixCache(MambaRadixCache):
 
         operation = self.cache_controller.prefetch(
             req_id,
-            host_indices,
             new_input_tokens,
             last_hash,
             prefix_keys,
@@ -1799,7 +1835,7 @@ class HiMambaRadixCache(MambaRadixCache):
         self.ongoing_prefetch[req_id] = (
             last_host_node,
             new_input_tokens,
-            host_indices,
+            None,
             operation,
         )
         self.cache_controller.prefetch_tokens_occupied += len(new_input_tokens)
@@ -1812,12 +1848,17 @@ class HiMambaRadixCache(MambaRadixCache):
             req_id
         ]
 
-        if operation.host_indices is None:
-            return True
-
         if not self.can_terminate_prefetch(operation):
             return False
 
+        if operation.host_indices is None:
+            # Stopping before host memory was committed (best_effort, timeout,
+            # or still mid-query): signal the worker to stop, then release the request.
+            self.cache_controller.terminate_prefetch(operation)
+            self._revoke_pending_prefetch(req_id)
+            return True
+
+        host_indices = operation.host_indices
         completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
             operation
         )
@@ -1957,6 +1998,8 @@ class HiMambaRadixCache(MambaRadixCache):
 
         last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[rid]
         if operation.host_indices is None:
+            self.cache_controller.terminate_prefetch(operation)
+            self._revoke_pending_prefetch(rid)
             return
 
         completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -524,15 +524,15 @@ class HiRadixCache(RadixCache):
         try:
             for req_id, info in list(self.ongoing_prefetch.items()):
                 try:
-                    last_host_node, token_ids, host_indices, _operation = info
+                    last_host_node, prefetch_key, _operation = info
                 except Exception:
                     # Unexpected shape; just drop it.
                     self.ongoing_prefetch.pop(req_id, None)
                     continue
 
                 try:
-                    if host_indices is not None:
-                        cc.mem_pool_host.free(host_indices)
+                    if _operation.host_indices is not None:
+                        cc.mem_pool_host.free(_operation.host_indices)
                 except Exception:
                     logger.exception(
                         "Failed to free host indices for prefetch %s", req_id
@@ -546,7 +546,7 @@ class HiRadixCache(RadixCache):
                     )
 
                 try:
-                    cc.prefetch_tokens_occupied -= len(token_ids)
+                    cc.prefetch_tokens_occupied -= len(prefetch_key)
                     if cc.prefetch_tokens_occupied < 0:
                         cc.prefetch_tokens_occupied = 0
                 except Exception:
@@ -577,6 +577,7 @@ class HiRadixCache(RadixCache):
         """
         self._drain_storage_control_queues_impl(
             n_revoke=None,
+            n_storage_hit=0,
             n_backup=None,
             n_release=None,
             log_metrics=False,
@@ -585,6 +586,7 @@ class HiRadixCache(RadixCache):
     def _drain_storage_control_queues_impl(
         self,
         n_revoke: Optional[int],
+        n_storage_hit: Optional[int],
         n_backup: Optional[int],
         n_release: Optional[int],
         log_metrics: bool,
@@ -603,13 +605,54 @@ class HiRadixCache(RadixCache):
 
         def _drain_revoke():
             for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
-                info = self.ongoing_prefetch.pop(req_id, None)
-                if info is not None:
-                    last_host_node, token_ids, _, _ = info
-                    last_host_node.release_host()
-                    cc.prefetch_tokens_occupied -= len(token_ids)
-                    if cc.prefetch_tokens_occupied < 0:
-                        cc.prefetch_tokens_occupied = 0
+                self._revoke_pending_prefetch(req_id)
+
+        def _drain_and_alloc_storage_hit():
+            # The L3 hit count is now known, so reserve exactly that much host
+            # memory (this is the whole point: no over-allocation up front).
+            # NOTE: alloc/evict here is rank-local but deterministic across TP
+            # ranks without extra synchronization: host pool mutations only
+            # happen on the scheduler thread at lockstep points (releases are
+            # page-granular and drained by the TP-min count), so every rank
+            # reaches the same success / fallback / revoke decision.
+            for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit):
+                req_id = operation.request_id
+                info = self.ongoing_prefetch.get(req_id)
+                if info is None:
+                    # request already aborted/cleaned up, skip
+                    continue
+                if operation.is_terminated():
+                    # request was aborted while the storage query was in flight
+                    self._revoke_pending_prefetch(req_id)
+                    continue
+
+                alloc_len = operation.storage_hit_count
+                host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    self.evict_host(alloc_len)
+                    host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    # Memory-pressure fallback: a shorter page-aligned prefix.
+                    available_size = cc.mem_pool_host.available_size()
+                    alloc_len = min(
+                        operation.storage_hit_count,
+                        available_size - (available_size % self.page_size),
+                    )
+                    if alloc_len >= self.prefetch_threshold:
+                        host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    self._revoke_pending_prefetch(req_id)
+                    logger.debug(
+                        f"Revoking prefetch for request {req_id} due to host memory allocation failure."
+                    )
+                    continue
+
+                operation.storage_hit_count = alloc_len
+                operation.hash_value = operation.hash_value[
+                    : alloc_len // self.page_size
+                ]
+                operation.host_indices = host_indices
+                cc.prefetch_buffer.put(operation)
 
         def _drain_backup():
             for operation in _drain_queue(cc.ack_backup_queue, n_backup):
@@ -631,6 +674,7 @@ class HiRadixCache(RadixCache):
                 cc.mem_pool_host.free(host_indices)
 
         _drain_revoke()
+        _drain_and_alloc_storage_hit()
         _drain_backup()
         _drain_release()
 
@@ -1370,7 +1414,6 @@ class HiRadixCache(RadixCache):
             extra_kwargs["pool_transfers"] = self._get_extra_pools().get("extra_pools")
         operation = prefetch_op_cls(
             "__storage_hit_query__",
-            self.cache_controller.mem_pool_host.get_dummy_flat_data_page()[:0],
             prefetch_key,
             last_hash,
             prefix_keys,
@@ -1419,6 +1462,7 @@ class HiRadixCache(RadixCache):
         qsizes = torch.tensor(
             [
                 cc.prefetch_revoke_queue.qsize(),
+                cc.prefetch_hit_queue.qsize(),
                 cc.ack_backup_queue.qsize(),
                 cc.host_mem_release_queue.qsize(),
             ],
@@ -1426,9 +1470,10 @@ class HiRadixCache(RadixCache):
         )
         self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN)
 
-        n_revoke, n_backup, n_release = map(int, qsizes.tolist())
+        n_revoke, n_storage_hit, n_backup, n_release = map(int, qsizes.tolist())
         self._drain_storage_control_queues_impl(
             n_revoke=n_revoke,
+            n_storage_hit=n_storage_hit,
             n_backup=n_backup,
             n_release=n_release,
             log_metrics=True,
@@ -1482,24 +1527,34 @@ class HiRadixCache(RadixCache):
         can_terminate = can_terminate or operation_terminated
         return can_terminate
 
+    def _revoke_pending_prefetch(self, req_id: str):
+        info = self.ongoing_prefetch.pop(req_id, None)
+        if info is None:
+            return
+        last_host_node, prefetch_key, _ = info
+        last_host_node.release_host()
+        cc = self.cache_controller
+        cc.prefetch_tokens_occupied = max(
+            0, cc.prefetch_tokens_occupied - len(prefetch_key)
+        )
+
     def check_prefetch_progress(self, req_id: str) -> bool:
         if req_id not in self.ongoing_prefetch:
             # there is no ongoing prefetch for this request or it has been revoked
             return True
 
-        # todo: more policies for prefetch progress such as timeout
-        # the current policy is to prefetch with best effort and terminate when queuing is over
-        last_host_node, prefetch_key, host_indices, operation = self.ongoing_prefetch[
-            req_id
-        ]
-
-        if operation.host_indices is None:
-            # prefetch has not been issued due to insufficient host memory
-            return True
+        last_host_node, prefetch_key, operation = self.ongoing_prefetch[req_id]
 
         if not self.can_terminate_prefetch(operation):
             return False
 
+        if operation.host_indices is None:
+            # Stopping before host memory was committed (best_effort, timeout, or
+            # still mid-query): signal the worker to stop, then release the request.
+            self.cache_controller.terminate_prefetch(operation)
+            self._revoke_pending_prefetch(req_id)
+            return True
+
         completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
             operation
         )
@@ -1513,7 +1568,7 @@ class HiRadixCache(RadixCache):
         )
         min_completed_tokens = completed_tokens_tensor.item()
         fetched_key = prefetch_key[:min_completed_tokens]
-        written_indices = host_indices[:min_completed_tokens]
+        written_indices = operation.host_indices[:min_completed_tokens]
         matched_length = self._insert_helper_host(
             last_host_node,
             fetched_key,
@@ -1521,9 +1576,11 @@ class HiRadixCache(RadixCache):
             hash_value[: min_completed_tokens // self.page_size],
         )
 
-        self.cache_controller.mem_pool_host.free(host_indices[:matched_length])
+        self.cache_controller.mem_pool_host.free(
+            operation.host_indices[:matched_length]
+        )
         self.cache_controller.append_host_mem_release(
-            host_indices[min_completed_tokens:completed_tokens]
+            operation.host_indices[min_completed_tokens:completed_tokens]
         )
         last_host_node.release_host()
         del self.ongoing_prefetch[req_id]
@@ -1542,9 +1599,7 @@ class HiRadixCache(RadixCache):
         if req_id not in self.ongoing_prefetch:
             return
 
-        _, _, _, operation = self.ongoing_prefetch[req_id]
-        if operation.host_indices is None:
-            return
+        _, _, operation = self.ongoing_prefetch[req_id]
         operation.mark_terminate()
 
     def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
@@ -1612,28 +1667,11 @@ class HiRadixCache(RadixCache):
             return
 
         last_host_node.protect_host()
-        host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
-        if host_indices is None:
-            self.evict_host(prefetch_length)
-            host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
-        if host_indices is None:
-            available_size = self.cache_controller.mem_pool_host.available_size()
-            prefetch_length = available_size - (available_size % self.page_size)
-            if prefetch_length >= self.prefetch_threshold:
-                prefetch_key = prefetch_key[:prefetch_length]
-                host_indices = self.cache_controller.mem_pool_host.alloc(
-                    prefetch_length
-                )
-                if host_indices is None:
-                    last_host_node.release_host()
-                    return
-            else:
-                last_host_node.release_host()
-                # no sufficient host memory for prefetch
-                return
+        # NOTE: host_indices is no longer pre-allocated here. It is allocated
+        # lazily in _drain_and_alloc_storage_hit() once the L3 storage hit count is known,
+        # so we only reserve host memory for pages that actually hit.
         operation = self.cache_controller.prefetch(
             req_id,
-            host_indices,
             prefetch_key,
             last_hash,
             prefix_keys,
@@ -1642,7 +1680,6 @@ class HiRadixCache(RadixCache):
         self.ongoing_prefetch[req_id] = (
             last_host_node,
             prefetch_key,
-            host_indices,
             operation,
         )
         self.cache_controller.prefetch_tokens_occupied += len(prefetch_key)
@@ -1838,15 +1875,17 @@ class HiRadixCache(RadixCache):
         if rid not in self.ongoing_prefetch:
             return
 
-        last_host_node, prefetch_key, host_indices, operation = self.ongoing_prefetch[
-            rid
-        ]
+        last_host_node, prefetch_key, operation = self.ongoing_prefetch[rid]
         if operation.host_indices is None:
+            self.cache_controller.terminate_prefetch(operation)
+            self._revoke_pending_prefetch(rid)
             return
 
         completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
         self._barrier_attn_groups()
         last_host_node.release_host()
         del self.ongoing_prefetch[rid]
-        self.cache_controller.append_host_mem_release(host_indices[:completed_tokens])
+        self.cache_controller.append_host_mem_release(
+            operation.host_indices[:completed_tokens]
+        )
         self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -122,7 +122,6 @@ class PrefetchOperation(StorageOperation):
     def __init__(
         self,
         request_id: str,
-        host_indices: torch.Tensor,
         token_ids: List[int],
         last_hash: Optional[str] = None,
         prefix_keys: Optional[List[str]] = None,
@@ -131,9 +130,10 @@ class PrefetchOperation(StorageOperation):
         self.request_id = request_id
         self._lock = threading.Lock()
         self._terminated_flag = False
+        self.storage_hit_count = 0
         self.start_time = time.monotonic()
         super().__init__(
-            host_indices,
+            None,
             token_ids,
             last_hash,
             prefix_keys=prefix_keys,
@@ -559,7 +559,6 @@ class HybridCacheController(BaseHiCacheController):
     def prefetch(
         self,
         request_id: str,
-        host_indices: torch.Tensor,
         new_input_tokens: List[int],
         last_hash: Optional[str] = None,
         prefix_keys: Optional[List[str]] = None,
@@ -567,7 +566,6 @@ class HybridCacheController(BaseHiCacheController):
     ) -> PrefetchOperation:
         operation = PrefetchOperation(
             request_id,
-            host_indices,
             new_input_tokens,
             last_hash,
             prefix_keys=prefix_keys,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1789,7 +1789,11 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 if phase == CacheTransferPhase.BACKUP_HOST
                 else indices_source.host_indices
             )
-            if indices is None or len(indices) == 0:
+            defer_kv_sidecar = (
+                phase == CacheTransferPhase.PREFETCH
+                and spec.indices_from_pool == PoolName.KV
+            )
+            if (indices is None or len(indices) == 0) and not defer_kv_sidecar:
                 continue
             transfers.append(
                 PoolTransfer(
@@ -1889,25 +1893,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             return
 
         anchor_lock_params = self.inc_host_lock_ref(last_host_node).to_dec_params()
-        host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
-        if host_indices is None:
-            self.evict_host(prefetch_length)
-            host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
-        if host_indices is None:
-            available_size = self.cache_controller.mem_pool_host.available_size()
-            prefetch_length = available_size - (available_size % self.page_size)
-            if prefetch_length >= self.prefetch_threshold:
-                prefetch_key = prefetch_key[:prefetch_length]
-                host_indices = self.cache_controller.mem_pool_host.alloc(
-                    prefetch_length
-                )
-            else:
-                self.dec_host_lock_ref(last_host_node, anchor_lock_params)
-                return
-        if host_indices is None:
-            self.dec_host_lock_ref(last_host_node, anchor_lock_params)
-            return
-
         comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
         alloc_failed = False
         for comp in self._components_tuple:
@@ -1925,13 +1910,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 break
             if transfers:
                 comp_xfers[comp.component_type] = transfers
-        kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices)
+        kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=None)
         sidecar_xfers = self._build_sidecar_transfers(
             CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers
         )
         if alloc_failed:
             self.cache_controller.append_host_mem_release(
-                host_indices=host_indices,
                 extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
             )
             self.dec_host_lock_ref(last_host_node, anchor_lock_params)
@@ -1941,7 +1925,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         aux_xfers.extend(sidecar_xfers)
         operation = self.cache_controller.prefetch(
             req_id,
-            host_indices,
             prefetch_key,
             last_hash,
             prefix_keys,
@@ -1950,7 +1933,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         self.ongoing_prefetch[req_id] = _OngoingPrefetch(
             last_host_node,
             prefetch_key,
-            host_indices,
+            None,
             operation,
             anchor_lock_params,
             comp_xfers,
@@ -2012,10 +1995,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             anchor_lock_params,
             comp_xfers,
         ) = self.ongoing_prefetch[req_id]
-        if operation.host_indices is None:
-            return True
         if not self.can_terminate_prefetch(operation):
             return False
+        if operation.host_indices is None:
+            self.cache_controller.terminate_prefetch(operation)
+            self._revoke_pending_prefetch(req_id)
+            return True
 
         completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
             operation
@@ -2082,8 +2067,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         if req_id not in self.ongoing_prefetch:
             return
         operation = self.ongoing_prefetch[req_id].operation
-        if operation.host_indices is None:
-            return
         operation.mark_terminate()
 
     def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
@@ -2103,6 +2086,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             comp_xfers,
         ) = self.ongoing_prefetch[rid]
         if operation.host_indices is None:
+            self.cache_controller.terminate_prefetch(operation)
+            self._revoke_pending_prefetch(rid)
             return
 
         completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
@@ -2115,9 +2100,31 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         )
         self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
 
+    def _revoke_pending_prefetch(self, req_id: str) -> None:
+        info = self.ongoing_prefetch.pop(req_id, None)
+        if info is None:
+            return
+        (
+            last_host_node,
+            prefetch_key,
+            _host_indices,
+            _operation,
+            anchor_lock_params,
+            comp_xfers,
+        ) = info
+        cc = self.cache_controller
+        cc.append_host_mem_release(
+            extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
+        )
+        self.dec_host_lock_ref(last_host_node, anchor_lock_params)
+        cc.prefetch_tokens_occupied = max(
+            0, cc.prefetch_tokens_occupied - len(prefetch_key)
+        )
+
     def _drain_storage_control_queues_impl(
         self,
         n_revoke: Optional[int],
+        n_storage_hit: Optional[int],
         n_backup: Optional[int],
         n_release: Optional[int],
         extra_release_counts: Optional[dict[PoolName, int]],
@@ -2136,28 +2143,46 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 yield item
 
         def _drain_revoke():
-            drained = 0
             for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
-                info = self.ongoing_prefetch.pop(req_id, None)
+                self._revoke_pending_prefetch(req_id)
+
+        def _drain_and_alloc_storage_hit():
+            for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit):
+                req_id = operation.request_id
+                info = self.ongoing_prefetch.get(req_id)
                 if info is None:
+                    # request already aborted/cleaned up, skip
                     continue
-                drained += 1
-                (
-                    last_host_node,
-                    prefetch_key,
-                    _host_indices,
-                    _operation,
-                    anchor_lock_params,
-                    comp_xfers,
-                ) = info
-                cc.append_host_mem_release(
-                    extra_pools=[x for xfers in comp_xfers.values() for x in xfers]
-                )
-                self.dec_host_lock_ref(last_host_node, anchor_lock_params)
-                cc.prefetch_tokens_occupied -= len(prefetch_key)
-                if cc.prefetch_tokens_occupied < 0:
-                    cc.prefetch_tokens_occupied = 0
-            return drained
+                if operation.is_terminated():
+                    # request was aborted while the storage query was in flight
+                    self._revoke_pending_prefetch(req_id)
+                    continue
+
+                alloc_len = operation.storage_hit_count
+                host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    self.evict_host(alloc_len)
+                    host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    # Memory-pressure fallback: a shorter page-aligned prefix.
+                    available_size = cc.mem_pool_host.available_size()
+                    alloc_len = min(
+                        operation.storage_hit_count,
+                        available_size - (available_size % self.page_size),
+                    )
+                    if alloc_len >= self.prefetch_threshold:
+                        host_indices = cc.mem_pool_host.alloc(alloc_len)
+                if host_indices is None:
+                    self._revoke_pending_prefetch(req_id)
+                    continue
+
+                operation.storage_hit_count = alloc_len
+                operation.hash_value = operation.hash_value[
+                    : alloc_len // self.page_size
+                ]
+                operation.host_indices = host_indices
+                self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices)
+                cc.prefetch_buffer.put(operation)
 
         def _drain_backup():
             drained = 0
@@ -2208,6 +2233,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             return drained
 
         _drain_revoke()
+        _drain_and_alloc_storage_hit()
         _drain_backup()
         _drain_release()
         _drain_extra_release()
@@ -2218,6 +2244,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         extra_pool_names = list(extra_release_queues)
         local_qsize_list = [
             cc.prefetch_revoke_queue.qsize(),
+            cc.prefetch_hit_queue.qsize(),
             cc.ack_backup_queue.qsize(),
             cc.host_mem_release_queue.qsize(),
             *[
@@ -2231,13 +2258,14 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         )
         self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN)
         qsize_list = list(map(int, qsizes.tolist()))
-        n_revoke, n_backup, n_release = qsize_list[:3]
+        n_revoke, n_storage_hit, n_backup, n_release = qsize_list[:4]
         extra_release_counts = {
             pool_name: count
-            for pool_name, count in zip(extra_pool_names, qsize_list[3:])
+            for pool_name, count in zip(extra_pool_names, qsize_list[4:])
         }
         self._drain_storage_control_queues_impl(
             n_revoke=n_revoke,
+            n_storage_hit=n_storage_hit,
             n_backup=n_backup,
             n_release=n_release,
             extra_release_counts=extra_release_counts,

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -380,7 +380,6 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
             def prefetch(
                 self,
                 request_id,
-                host_indices,
                 new_input_tokens,
                 last_hash=None,
                 prefix_keys=None,
@@ -388,7 +387,6 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
             ):
                 self.prefetch_args = (
                     request_id,
-                    host_indices,
                     new_input_tokens,
                     last_hash,
                     prefix_keys,
@@ -400,7 +398,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
         cache.cache_controller = controller
         cache.prefetch_from_storage("req", cache.root_node, tokens)
 
-        _, _, storage_key, _, _, _ = controller.prefetch_args
+        _, storage_key, _, _, _ = controller.prefetch_args
         self.assertIsInstance(storage_key, RadixKey)
         self.assertTrue(storage_key.is_bigram)
         self.assertEqual(len(storage_key), len(tokens) - 1)
@@ -2251,6 +2249,10 @@ class UnifiedRadixCacheSuite:
     def _run_prefetch_to_completion(self, cache, req_id, timeout: float = 10.0):
         deadline = time.time() + timeout
         while time.time() < deadline:
+            # Host memory is reserved (and IO started) by the scheduler-thread
+            # drain once the L3 hit count is known, so pump it like the real
+            # scheduler loop does (check_hicache_events before progress checks).
+            cache.drain_storage_control_queues()
             if cache.check_prefetch_progress(req_id):
                 return
             time.sleep(0.01)
@@ -2385,15 +2387,15 @@ class UnifiedRadixCacheSuite:
                 comp_xfers = info[-1]
                 names = [t.name for xfers in comp_xfers.values() for t in xfers]
                 if PoolName.SWA in names:
-                    return 1 + names.index(PoolName.SWA)
-            return None
+                    return 1 + names.index(PoolName.SWA), 1 + len(names)
+            return None, None
 
         def fake(tensor, op=None, group=None):
             if op == dist.ReduceOp.MIN:
                 min_sizes.append(tensor.numel())
                 if drop_swa:
-                    idx = swa_packed_index()
-                    if idx is not None and idx < tensor.numel():
+                    idx, packed_numel = swa_packed_index()
+                    if idx is not None and tensor.numel() == packed_numel:
                         tensor[idx] = 0
             return None