[HiCache] Fix sidecar pool life-time issue

python/sglang/srt/managers/cache_controller.py

@@ -28,6 +28,7 @@
     HiCacheStorageExtraInfo,
     PoolName,
     PoolTransfer,
+    count_pool_hits,
 )
 
 if TYPE_CHECKING:
@@ -957,7 +958,7 @@ def append_host_mem_release(self, host_indices: torch.Tensor):
 
     def _page_get_zero_copy(
         self, operation, hash_values, host_indices, extra_info=None
-    ):
+    ) -> int:
         results = self.storage_backend.batch_get_v1(
             hash_values, host_indices, extra_info
         )
@@ -968,35 +969,43 @@ def _page_get_zero_copy(
                     f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
                 )
                 break
-            inc += self.page_size
-        operation.increment(inc)
+            inc += 1
+        return inc
 
     # todo: deprecate
-    def _generic_page_get(self, operation, hash_values, host_indices, extra_info=None):
+    def _generic_page_get(
+        self, operation, hash_values, host_indices, extra_info=None
+    ) -> int:
         dummy_page_dst = [
             self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values
         ]
         page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst)
         if page_data is None:
-            return
+            return 0
+        count = 0
         for i in range(len(hash_values)):
             if page_data[i] is None:
                 logger.warning(
                     f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
                 )
                 break
-            # Must set the data before increasing the completed tokens.
-            # Otherwise this page may be read before being set.
+            if operation.is_terminated():
+                break
             self.mem_pool_host.set_from_flat_data_page(
                 host_indices[i * self.page_size],
                 page_data[i],
             )
-            if not operation.increment(self.page_size):
-                break  # Operation terminated by controller
+            count += 1
+        return count
 
     def _page_transfer(self, operation):
         # Transfer batch by batch
         prefix_keys = operation.prefix_keys
+        kv_derived_transfers = [
+            transfer
+            for transfer in getattr(operation, "pool_transfers", None) or []
+            if transfer.indices_from_pool == PoolName.KV
+        ]
         for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
             batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE]
             batch_host_indices = operation.host_indices[
@@ -1009,21 +1018,66 @@ def _page_transfer(self, operation):
             if self.has_draft:
                 self._draft_page_get(batch_hashes, batch_host_indices)
 
-            prev_completed_tokens = operation.completed_tokens
             # Get one batch token, and update the completed_tokens if succeed
             extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
-            self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info)
+
+            hit_pages = self._page_transfer_kv_batch(
+                operation,
+                batch_hashes,
+                batch_host_indices,
+                extra_info,
+                kv_derived_transfers,
+            )
             # Check termination
-            if (
-                operation.completed_tokens
-                != prev_completed_tokens + len(batch_hashes) * self.page_size
-            ):
+            if not operation.increment(hit_pages * self.page_size):
+                # The scheduler thread has terminate this prefetch.
+                break
+            if hit_pages != len(batch_hashes):
                 operation.mark_terminate()
                 break  # Some operations fail or operation terminated by controller
 
             if prefix_keys and len(prefix_keys) > 0:
                 prefix_keys += batch_hashes
 
+    def _page_transfer_kv_batch(
+        self,
+        operation: PrefetchOperation,
+        batch_hashes: List[str],
+        batch_host_indices: torch.Tensor,
+        extra_info: HiCacheStorageExtraInfo,
+        kv_derived_transfers: List[PoolTransfer],
+    ) -> int:
+        """Read a single batch from KV and KV-derived pools (e.g. indexer pool).
+
+        Return the number of hit pages.  If the hits from KV and KV-derived pools differ,
+        clamp to the minimal number of hits.
+
+        Here, "batch" means a single unit of L3 read, not a "batch" in model forward.
+        """
+        # Read from KV pool.
+        kv_hits = self.page_get_func(
+            operation, batch_hashes, batch_host_indices, extra_info
+        )
+
+        # Read from KV-derived sidecar pools, if any.
+        sidecar_hits: dict[str, int] = {}
+        if len(kv_derived_transfers) > 0:
+            current_kv_derived_transfers = [
+                PoolTransfer(
+                    name=transfer.name,
+                    host_indices=batch_host_indices,
+                    keys=batch_hashes,
+                )
+                for transfer in kv_derived_transfers
+            ]
+            sidecar_results = self.storage_backend.batch_get_v2(
+                current_kv_derived_transfers
+            )
+            sidecar_hits = count_pool_hits(sidecar_results)
+
+        # Clamp to minimal number of hits.
+        return min([kv_hits, *sidecar_hits.values()])
+
     def prefetch_io_aux_func(self):
         """
         Auxiliary function conducting IO operations for prefetching.

python/sglang/srt/mem_cache/hicache_storage.py

@@ -139,12 +139,14 @@ def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None:
         Every extra pool contributes a prefix that must be contiguous from the
         start, so count the leading run of successes
         """
-        self.extra_pool_hit_pages.update(
-            {
-                name: (rs.index(False) if False in rs else len(rs))
-                for name, rs in results.items()
-            }
-        )
+        self.extra_pool_hit_pages.update(count_pool_hits(results))
+
+
+def count_pool_hits(results: dict[str, List[bool]]) -> dict[str, int]:
+    return {
+        name: (rs.index(False) if False in rs else len(rs))
+        for name, rs in results.items()
+    }
 
 
 class HiCacheStorage(ABC):

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -1583,9 +1583,13 @@ def can_terminate_prefetch(self, operation: PrefetchOperation):
         if len(operation.hash_value) == 0:
             completed = False
         else:
+            # kv pool
             completed = (
                 operation.completed_tokens == len(operation.hash_value) * self.page_size
             )
+            # sidecar pool (only present when using HybridCacheController)
+            if completed and getattr(operation, "pool_transfers", None):
+                completed = getattr(operation, "pool_transfers_done", False)
 
         if self.prefetch_stop_policy == "wait_complete":
             can_terminate = completed
@@ -1595,13 +1599,6 @@ def can_terminate_prefetch(self, operation: PrefetchOperation):
             # unknown prefetch stop policy, just return True
             return True
 
-        if (
-            completed
-            and getattr(operation, "pool_transfers", None)
-            and not getattr(operation, "pool_transfers_done", True)
-        ):
-            can_terminate = False
-
         operation_terminated = operation.is_terminated()
         states = torch.tensor(
             [1 - int(can_terminate), int(operation_terminated)],
@@ -1648,9 +1645,12 @@ def check_prefetch_progress(self, req_id: str) -> bool:
         )
         logger.debug(f"Prefetch {req_id} completed with {completed_tokens} tokens")
 
-        min_completed_tokens = self._sync_and_clamp_prefetch_result(
-            operation, completed_tokens
+        # Synchronize workers before mutating host cache tree state.
+        completed_tokens_tensor = torch.tensor(completed_tokens, dtype=torch.int)
+        self._all_reduce_attn_groups(
+            completed_tokens_tensor, torch.distributed.ReduceOp.MIN
         )
+        min_completed_tokens = completed_tokens_tensor.item()
 
         fetched_key = prefetch_key[:min_completed_tokens]
         written_indices = operation.host_indices[:min_completed_tokens]
@@ -1680,39 +1680,6 @@ def check_prefetch_progress(self, req_id: str) -> bool:
 
         return True
 
-    def _sync_and_clamp_prefetch_result(
-        self,
-        operation: PrefetchOperation,
-        completed_tokens: int,
-    ) -> int:
-        """Sync prefetch results across ATTN groups and decide the usable prefix.
-
-        HiRadixCache only wires DSA-style stacks (Full attention + a KV-derived
-        ALL_PAGES sidecar such as the DSA / MiniMax indexer); For the DSA case we *clamp*
-        to the minimum fetched prefix shared by the Full KV pool and every
-        sidecar rather than discarding everything. With no sidecar (FULL-only)
-        this is just the synced Full KV completion.
-        """
-        # Sync completed tokens and per-pool hit pages across ATTN groups, taking
-        # the minimum so every rank agrees on the same usable prefix length.
-        pool_transfers = getattr(operation, "pool_transfers", None) or []
-        hit_pages = (
-            operation.pool_storage_result.extra_pool_hit_pages if pool_transfers else {}
-        )
-        pool_hit_pages = [hit_pages.get(t.name, 0) for t in pool_transfers]
-        packed = torch.tensor([completed_tokens, *pool_hit_pages], dtype=torch.int)
-        self._all_reduce_attn_groups(packed, torch.distributed.ReduceOp.MIN)
-        min_completed_tokens = int(packed[0].item())
-        pool_hit_pages = list(map(int, packed[1:].tolist()))
-
-        # Clamp to the shared minimum prefix of the Full KV completion and each
-        # KV-derived ALL_PAGES sidecar (e.g. the DSA indexer). FULL-only has no
-        # sidecar, so the usable prefix is just the Full KV completion.
-        usable_pages = min_completed_tokens // self.page_size
-        if pool_transfers:
-            usable_pages = min(usable_pages, *pool_hit_pages)
-        return usable_pages * self.page_size
-
     def terminate_prefetch(self, req_id: str):
         if req_id not in self.ongoing_prefetch:
             return

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

@@ -149,6 +149,15 @@ def increment(self, num_tokens: int):
             self.completed_tokens += num_tokens
             return True
 
+    def complete_pool_transfers(self, result: dict[str, list[bool]]) -> bool:
+        with self._lock:
+            if self._terminated_flag:
+                return False
+            assert not self.pool_transfers_done
+            self.pool_transfers_done = True
+            self.pool_storage_result.update_extra_pool_hit_pages(result)
+            return True
+
     def mark_terminate(self):
         with self._lock:
             self._terminated_flag = True
@@ -738,26 +747,46 @@ def move_hybrid_indices(
                 )
         return host_indices, device_indices, resolved_pool_transfers
 
-    def _page_transfer(self, operation):
-        # KV pools first — determines actual completed page count
+    def _page_transfer(self, operation: PrefetchOperation):
+        # KV pools and KV-derived pools first — determines actual completed page count
         super()._page_transfer(operation)
 
+        # Read non-KV derived sidecar pool, e.g. SWA, Mamba.
+        self._page_transfer_sidecar(operation)
+
+    def _page_transfer_sidecar(self, operation: PrefetchOperation):
+        if operation.pool_transfers is None:
+            return
+
         # Extra pools only after KV fully completes. If KV terminated early
         # (IO failure, timeout, TP mismatch), skip extra IO entirely to avoid
         # data misalignment.
         kv_completed_pages = operation.completed_tokens // self.page_size
-        if (
-            operation.pool_transfers
-            and not operation.is_terminated()
-            and kv_completed_pages == len(operation.hash_value)
+        sidecar_completed_pages: dict[str, List[bool]] = {}
+        if not operation.is_terminated() and kv_completed_pages == len(
+            operation.hash_value
         ):
+            # KV-derived sidecar pools are handled in CacheController._page_transfer_kv_batch.
+            # Only handle non-KV-derived sidecar pools here.
+            transfers_nonkv = [
+                transfer
+                for transfer in operation.pool_transfers
+                if transfer.indices_from_pool != PoolName.KV
+            ]
             self._sync_trailing_keys(
-                operation.pool_transfers, operation.hash_value, kv_completed_pages
+                transfers_nonkv, operation.hash_value, kv_completed_pages
             )
-            self._resolve_sidecar_derived_pool_transfers(operation)
-            results = self.storage_backend.batch_get_v2(operation.pool_transfers)
-            operation.pool_storage_result.update_extra_pool_hit_pages(results)
-        operation.pool_transfers_done = True
+            self._resolve_sidecar_nonkv_derived_pool_transfers(operation)
+            sidecar_completed_pages = self.storage_backend.batch_get_v2(transfers_nonkv)
+
+        # It is tricky to determine which thread should release memory of extra pools.
+        # There are two cases:
+        # 1) If complete_pool_transfers() runs BEFORE mark_terminate(), then the scheduler
+        #    thread is responsible for releasing the extra pool.
+        # 2) If complete_pool_transfer() runs AFTER mark_terminate(), then the prefetch IO
+        #    thread (current thread) should release the extra pool (in below code).
+        if not operation.complete_pool_transfers(sidecar_completed_pages):
+            self.append_host_mem_release(extra_pools=operation.pool_transfers)
 
     def _page_backup(self, operation):
         # MLA KV is replicated across TP ranks and should still be written only
@@ -769,7 +798,8 @@ def _page_backup(self, operation):
         ]
 
         if backup_transfers:
-            self._resolve_sidecar_derived_pool_transfers(operation)
+            self._resolve_sidecar_kv_derived_pool_transfers(operation)
+            self._resolve_sidecar_nonkv_derived_pool_transfers(operation)
             results = self.storage_backend.batch_set_v2(backup_transfers)
             operation.pool_storage_result.update_extra_pool_hit_pages(results)
 
@@ -835,7 +865,14 @@ def backup_thread_func(self):
             except Empty:
                 continue
 
-    def _resolve_sidecar_derived_pool_transfers(self, operation):
+    def _resolve_sidecar_kv_derived_pool_transfers(self, operation):
+        for transfer in operation.pool_transfers:
+            if transfer.indices_from_pool == PoolName.KV:
+                transfer.host_indices = operation.host_indices
+                if transfer.keys is None:
+                    transfer.keys = operation.hash_value
+
+    def _resolve_sidecar_nonkv_derived_pool_transfers(self, operation):
         for transfer in operation.pool_transfers:
             if transfer.indices_from_pool is None:
                 continue
@@ -858,9 +895,7 @@ def _resolve_sidecar_derived_pool_transfers(self, operation):
                 if transfer.keys is None:
                     transfer.keys = source.keys
             else:
-                transfer.host_indices = operation.host_indices
-                if transfer.keys is None:
-                    transfer.keys = operation.hash_value
+                pass
 
     def _sync_trailing_keys(
         self,