[HiCache] Fix PP inconsistency with HiCache L3 (#22607) (#27010)

Co-authored-by: ybyang <ybyang7@iflytek.com>
Co-authored-by: hzh0425 <hzh0425@apache.org>
Co-authored-by: shangmingc <csmthu@gmail.com>
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>

python/sglang/srt/managers/cache_controller.py

@@ -13,9 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+
 import logging
 import threading
 import time
+from dataclasses import dataclass
 from queue import Empty, Queue
 from typing import TYPE_CHECKING, Callable, List, NamedTuple, Optional
 
@@ -27,6 +29,7 @@ from sglang.srt.mem_cache.hicache_storage import (
     HiCacheStorageExtraInfo,
     PoolName,
     PoolTransfer,
+    count_pool_hits,
 )
 
 if TYPE_CHECKING:
@@ -186,6 +189,32 @@ class HiCacheAck(NamedTuple):
     num_bytes: int = 0
 
 
+@dataclass
+class PrefetchAck:
+    """ACK for prefetch operation.
+
+    A sequence of PrefetchAck is sent to the scheduler thread via ack_prefetch_queue,
+    indicating progress or completion of the prefetch operation.
+
+    For example, a prefetch operation may results into the following sequence of PrefetchAck:
+
+    1. PrefetchAck(completed_tokens = 128)
+    2. PrefetchAck(completed_tokens = 256)
+    3. PrefetchAck(pool_hits={INDEXER: 256})
+    4. PrefetchAck(completed_req = True)
+
+    The last PrefetchAck always specifies completed_req = True.
+    """
+
+    rid: str
+    operation: PrefetchOperation
+    # Number of hits in KV pool.
+    completed_tokens: Optional[int] = None
+    # Number of hits in extra pools.
+    pool_hits: Optional[dict[str, int]] = None
+    completed_req: Optional[bool] = None
+
+
 class StorageOperation:
     counter = 0
 
@@ -244,19 +273,13 @@ class PrefetchOperation(StorageOperation):
 
         super().__init__(None, token_ids, last_hash, prefix_keys=prefix_keys)
 
-    def increment(self, num_tokens: int):
-        with self._lock:
-            if self._terminated_flag:
-                return False
-            self.completed_tokens += num_tokens
-            return True
-
     def mark_terminate(self):
         with self._lock:
             self._terminated_flag = True
 
     def is_terminated(self) -> bool:
-        return self._terminated_flag
+        with self._lock:
+            return self._terminated_flag
 
 
 class HiCacheController:
@@ -285,7 +308,8 @@ class HiCacheController:
         self.attn_cp_group = attn_cp_group
         self.attn_tp_group = attn_tp_group
         self.pp_group = pp_group
-        self.prefetch_sync_groups: List[torch.distributed.ProcessGroup] = []
+        self.prefetch_hits_sync_groups: List[torch.distributed.ProcessGroup] = []
+        self.prefetch_completion_sync_groups: List[torch.distributed.ProcessGroup] = []
         self.mem_pool_device_allocator = token_to_kv_pool_allocator
         mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
         from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
@@ -322,7 +346,10 @@ class HiCacheController:
         self.storage_stop_event = threading.Event()
 
         # Storage control queues, (re)created whenever the storage threads start.
+        self.prefetch_buffer: Optional[Queue[PrefetchOperation]] = None
+        self.prefetch_sync_queue: Optional[Queue[PrefetchAck]] = None
         self.prefetch_hit_queue: Optional[Queue[StorageOperation]] = None
+        self.ack_prefetch_queue = Queue[PrefetchAck]()
         self.ack_backup_queue: Optional[Queue[StorageOperation]] = None
         self.host_mem_release_queue: Optional[Queue[torch.Tensor]] = None
 
@@ -369,16 +396,18 @@ class HiCacheController:
             )
         return 0, 1
 
-    def _create_prefetch_sync_groups(self) -> None:
+    def _create_sync_groups(self) -> List[torch.distributed.ProcessGroup]:
         from sglang.srt.distributed.parallel_state import create_custom_parallel_group
 
-        self.prefetch_sync_groups = []
+        groups: List[torch.distributed.ProcessGroup] = []
         seen_rank_sets = set()
 
         if self.attn_cp_group is not None or self.attn_tp_group is not None:
             base_groups = [self.attn_cp_group, self.attn_tp_group]
         else:
             base_groups = [self.tp_group]
+        if self.pp_group is not None:
+            base_groups.append(self.pp_group)
 
         for group in base_groups:
             if group is None or torch.distributed.get_world_size(group=group) == 1:
@@ -387,22 +416,29 @@ class HiCacheController:
             if group_ranks in seen_rank_sets:
                 continue
             seen_rank_sets.add(group_ranks)
-            self.prefetch_sync_groups.append(
+            groups.append(
                 create_custom_parallel_group(
                     group_ranks=list(group_ranks), backend="gloo"
                 )
             )
+        return groups
 
-    def _destroy_prefetch_sync_groups(self) -> None:
-        for group in self.prefetch_sync_groups:
+    def _destroy_sync_groups(
+        self, groups: List[torch.distributed.ProcessGroup]
+    ) -> None:
+        for group in groups:
             try:
                 torch.distributed.destroy_process_group(group)
             except Exception:
                 pass
-        self.prefetch_sync_groups = []
 
-    def _all_reduce_prefetch_groups(self, tensor: torch.Tensor, op) -> None:
-        for group in self.prefetch_sync_groups:
+    def _all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op,
+        groups: List[torch.distributed.ProcessGroup],
+    ) -> None:
+        for group in groups:
             torch.distributed.all_reduce(tensor, op=op, group=group)
 
     def _start_storage_threads(self):
@@ -416,17 +452,27 @@ class HiCacheController:
         self.prefetch_thread = threading.Thread(
             target=self.prefetch_thread_func, daemon=True
         )
+        self.prefetch_io_aux_thread = threading.Thread(
+            target=self.prefetch_io_aux_func, daemon=True
+        )
+        self.prefetch_sync_thread = threading.Thread(
+            target=self.prefetch_sync_thread_func, daemon=True
+        )
         self.backup_thread = threading.Thread(
             target=self.backup_thread_func, daemon=True
         )
         self.prefetch_queue = Queue()
         self.backup_queue = Queue()
-
+        self.prefetch_buffer = Queue()
+        self.prefetch_sync_queue = Queue()
         self.prefetch_hit_queue = Queue()
+        self.ack_prefetch_queue = Queue()
         self.ack_backup_queue = Queue()
         self.host_mem_release_queue = Queue()
 
         self.prefetch_thread.start()
+        self.prefetch_io_aux_thread.start()
+        self.prefetch_sync_thread.start()
         self.backup_thread.start()
 
     def _stop_storage_threads(self):
@@ -449,6 +495,8 @@ class HiCacheController:
                 self.backup_queue.put_nowait(None)
             if hasattr(self, "prefetch_buffer"):
                 self.prefetch_buffer.put_nowait(None)
+            if hasattr(self, "prefetch_sync_queue"):
+                self.prefetch_sync_queue.put_nowait(None)
         except Exception:
             pass
 
@@ -460,6 +508,8 @@ class HiCacheController:
             threads.append(self.backup_thread)
         if hasattr(self, "prefetch_io_aux_thread"):
             threads.append(self.prefetch_io_aux_thread)
+        if hasattr(self, "prefetch_sync_thread"):
+            threads.append(self.prefetch_sync_thread)
 
         for t in threads:
             try:
@@ -542,7 +592,8 @@ class HiCacheController:
 
             # Use dedicated gloo groups so storage prefetch sync is isolated
             # from other collectives and consistent across CPxTP participants.
-            self._create_prefetch_sync_groups()
+            self.prefetch_hits_sync_groups = self._create_sync_groups()
+            self.prefetch_completion_sync_groups = self._create_sync_groups()
 
             # Select the get and set functions
             self.page_get_func = self._generic_page_get
@@ -569,7 +620,10 @@ class HiCacheController:
                 self._stop_storage_threads()
             except Exception:
                 pass
-            self._destroy_prefetch_sync_groups()
+            self._destroy_sync_groups(self.prefetch_hits_sync_groups)
+            self._destroy_sync_groups(self.prefetch_completion_sync_groups)
+            self.prefetch_hits_sync_groups = []
+            self.prefetch_completion_sync_groups = []
             try:
                 if (
                     hasattr(self, "storage_backend")
@@ -609,7 +663,11 @@ class HiCacheController:
             raise RuntimeError("Stop storage threads failed; detach aborted.") from e
 
         # Best-effort destroy process groups created for storage ops.
-        self._destroy_prefetch_sync_groups()
+        self._destroy_sync_groups(
+            self.prefetch_hits_sync_groups + self.prefetch_completion_sync_groups
+        )
+        self.prefetch_hits_sync_groups = []
+        self.prefetch_completion_sync_groups = []
 
         # Best-effort close (some backends rely on GC/destructor).
         try:
@@ -704,10 +762,15 @@ class HiCacheController:
         self.ack_load_queue.clear()
         if self.enable_storage:
             self.prefetch_thread.join()
+            self.prefetch_io_aux_thread.join()
+            self.prefetch_sync_thread.join()
             self.backup_thread.join()
             self.prefetch_queue.queue.clear()
             self.backup_queue.queue.clear()
+            self.prefetch_buffer.queue.clear()
+            self.prefetch_sync_queue.queue.clear()
             self.prefetch_hit_queue.queue.clear()
+            self.ack_prefetch_queue.queue.clear()
             self.ack_backup_queue.queue.clear()
             self.host_mem_release_queue.queue.clear()
             self.prefetch_tokens_occupied = 0
@@ -718,10 +781,18 @@ class HiCacheController:
             self.prefetch_thread = threading.Thread(
                 target=self.prefetch_thread_func, daemon=True
             )
+            self.prefetch_io_aux_thread = threading.Thread(
+                target=self.prefetch_io_aux_func, daemon=True
+            )
+            self.prefetch_sync_thread = threading.Thread(
+                target=self.prefetch_sync_thread_func, daemon=True
+            )
             self.backup_thread = threading.Thread(
                 target=self.backup_thread_func, daemon=True
             )
             self.prefetch_thread.start()
+            self.prefetch_io_aux_thread.start()
+            self.prefetch_sync_thread.start()
             self.backup_thread.start()
 
     def write(
@@ -984,6 +1055,14 @@ class HiCacheController:
         return operation
 
     def terminate_prefetch(self, operation):
+        """
+        Request to terminate a prefetch operation.
+
+        Must be called in the scheduler thread.
+
+        Asynchronous prefetch tasks may be running in background threads.  When all prefetch
+        tasks are terminated, a PrefetchAck with completed_req=True will be sent to ack_prefetch_queue.
+        """
         operation.mark_terminate()
         return operation.completed_tokens, operation.hash_value
 
@@ -996,7 +1075,7 @@ class HiCacheController:
 
     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
         )
@@ -1007,61 +1086,125 @@ class HiCacheController:
                     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):
+    def _page_transfer(self, operation: PrefetchOperation) -> int:
         # 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
+        ]
+        all_success = True
+        completed_pages = 0
         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[
-                i * self.page_size : (i + len(batch_hashes)) * self.page_size
+            # When an error is occurred, we should keep looping and produce the same number of
+            # PrefetchAck as other ranks do, because prefetch_sync_thread (i.e. consumer of
+            # prefetch_sync_queue) perform reduce on the results.  This is so tricky.
+            if all_success and operation.is_terminated():
+                all_success = False
+            if all_success:
+                batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE]
+                batch_host_indices = operation.host_indices[
+                    i * self.page_size : (i + len(batch_hashes)) * self.page_size
+                ]
+
+                # Best-effort draft L3 read before publishing target completion.
+                # Otherwise wait_complete can race and load back target KV before
+                # draft KV reaches host memory.
+                if self.has_draft:
+                    self._draft_page_get(batch_hashes, batch_host_indices)
+
+                # Get one batch token, and update the completed_tokens if succeed
+                extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
+
+                hit_pages = self._page_transfer_kv_batch(
+                    operation,
+                    batch_hashes,
+                    batch_host_indices,
+                    extra_info,
+                    kv_derived_transfers,
+                )
+                # Check termination
+                if hit_pages != len(batch_hashes):
+                    all_success = False
+                if prefix_keys and len(prefix_keys) > 0:
+                    prefix_keys += batch_hashes
+                completed_pages += hit_pages
+            ack = PrefetchAck(
+                rid=operation.request_id,
+                completed_tokens=completed_pages * self.page_size,
+                operation=operation,
+            )
+            self.prefetch_sync_queue.put(ack)
+        return completed_pages
+
+    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)
 
-            # Best-effort draft L3 read before publishing target completion.
-            # Otherwise wait_complete can race and load back target KV before
-            # draft KV reaches host memory.
-            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)
-            # Check termination
-            if (
-                operation.completed_tokens
-                != prev_completed_tokens + len(batch_hashes) * self.page_size
-            ):
-                operation.mark_terminate()
-                break  # Some operations fail or operation terminated by controller
-
-            if prefix_keys and len(prefix_keys) > 0:
-                prefix_keys += batch_hashes
+        # Clamp to minimal number of hits.
+        return min([kv_hits, *sidecar_hits.values()])
 
     def prefetch_io_aux_func(self):
         """
@@ -1073,9 +1216,13 @@ class HiCacheController:
                 if operation is None:
                     continue
                 self._page_transfer(operation)
-                # operation terminated by controller, release pre-allocated memory
-                self.append_host_mem_release(
-                    operation.host_indices[operation.completed_tokens :]
+
+                self.prefetch_sync_queue.put(
+                    PrefetchAck(
+                        rid=operation.request_id,
+                        completed_req=True,
+                        operation=operation,
+                    )
                 )
             except Empty:
                 continue
@@ -1129,11 +1276,6 @@ class HiCacheController:
         """
         Manage prefetching operations from storage backend to host memory.
         """
-        self.prefetch_buffer = Queue()
-        self.prefetch_io_aux_thread = threading.Thread(
-            target=self.prefetch_io_aux_func, daemon=True
-        )
-        self.prefetch_io_aux_thread.start()
         while (not self.storage_stop_event.is_set()) or not self.prefetch_queue.empty():
             try:
                 operation = self.prefetch_queue.get(block=True, timeout=1)
@@ -1146,8 +1288,10 @@ class HiCacheController:
                 storage_hit_count_tensor = torch.tensor(
                     storage_hit_count, dtype=torch.int
                 )
-                self._all_reduce_prefetch_groups(
-                    storage_hit_count_tensor, torch.distributed.ReduceOp.MIN
+                self._all_reduce(
+                    storage_hit_count_tensor,
+                    torch.distributed.ReduceOp.MIN,
+                    self.prefetch_hits_sync_groups,
                 )
                 storage_hit_count = storage_hit_count_tensor.item()
 
@@ -1300,3 +1444,29 @@ class HiCacheController:
 
             except Empty:
                 continue
+
+    def prefetch_sync_thread_func(self):
+        """Synchronize prefetch results across all PP and TP ranks."""
+        while not self.storage_stop_event.is_set():
+            try:
+                ack = self.prefetch_sync_queue.get(block=True, timeout=1)
+                if ack is None:
+                    continue
+                self._reduce_prefetch_ack(ack)
+                self.ack_prefetch_queue.put(ack)
+            except Empty:
+                continue
+
+    def _reduce_prefetch_ack(self, ack: PrefetchAck) -> None:
+        """Synchronize all ranks to agree on a PrefetchAck."""
+        if ack.completed_tokens is not None:
+            # Determine the minimal successful prefix of tokens.
+            completed_tokens_tensor = torch.tensor(
+                ack.completed_tokens, dtype=torch.int
+            )
+            self._all_reduce(
+                completed_tokens_tensor,
+                torch.distributed.ReduceOp.MIN,
+                self.prefetch_completion_sync_groups,
+            )
+            ack.completed_tokens = completed_tokens_tensor.item()

python/sglang/srt/managers/scheduler.py

@@ -2875,8 +2875,6 @@ class Scheduler(
                 if self.enable_hicache_storage:
                     # Release prefetch events associated with the request
                     self.tree_cache.release_aborted_request(candidate_req.rid)
-                elif self.enable_hierarchical_cache:
-                    self.tree_cache.terminate_prefetch(candidate_req.rid)
                 self.waiting_queue.pop(idx)
                 req_to_abort = candidate_req
                 message = "The request is aborted by a higher priority request."

python/sglang/srt/mem_cache/hicache_storage.py

@@ -133,18 +133,20 @@ class PoolTransferResult:
         """Accumulate kv_hit_pages across batches (max = last successful batch)."""
         self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages)
 
-    def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None:
+    def update_extra_pool_hit_pages(self, results: dict[str, int]) -> None:
         """Record actual load/write success counts per extra pool.
 
         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(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

@@ -5,9 +5,9 @@ import heapq
 import json
 import logging
 import os
+import queue
 import threading
 import time
-from queue import Empty
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
 import torch
@@ -228,15 +228,6 @@ class HiRadixCache(RadixCache):
         if not reduced and self.tp_world_size > 1:
             torch.distributed.all_reduce(tensor, op=op, group=self.tp_group)
 
-    def _barrier_attn_groups(self):
-        waited = False
-        for group in (self.attn_cp_group, self.attn_tp_group):
-            if group is not None and torch.distributed.get_world_size(group=group) > 1:
-                torch.distributed.barrier(group=group)
-                waited = True
-        if not waited and self.tp_world_size > 1:
-            torch.distributed.barrier(group=self.tp_group)
-
     def _drain_async_work(self):
         """
         Block until all outstanding async sends are consumed, then clear.
@@ -585,6 +576,7 @@ class HiRadixCache(RadixCache):
         """
         self._drain_storage_control_queues_impl(
             n_storage_hit=0,
+            n_ack_prefetch=None,
             n_backup=None,
             n_release=None,
             log_metrics=False,
@@ -593,21 +585,27 @@ class HiRadixCache(RadixCache):
     def _drain_storage_control_queues_impl(
         self,
         n_storage_hit: Optional[int],
+        n_ack_prefetch: Optional[int],
         n_backup: Optional[int],
         n_release: Optional[int],
         log_metrics: bool,
     ):
         cc = self.cache_controller
 
-        def _drain_queue(q, limit: Optional[int]):
-            drained = 0
-            while limit is None or drained < limit:
-                try:
-                    item = q.get_nowait()
-                except Empty:
-                    break
-                drained += 1
-                yield item
+        def _drain_queue(q: queue.Queue, n: Optional[int]):
+            """If n is None, consume all items from the queue.
+            Otherwise, consume n items from the queue.
+            """
+            if n is None:
+                while not q.empty():
+                    item = q.get()
+                    yield item
+            else:
+                for _ in range(n):
+                    # Block when there is no enough elements.
+                    # All TP/PP ranks must consume the same number of elements.
+                    item = q.get()
+                    yield item
 
         def _drain_and_alloc_storage_hit():
             # The L3 hit count is now known, so reserve exactly that much host
@@ -663,6 +661,25 @@ class HiRadixCache(RadixCache):
                 operation.host_indices = host_indices
                 cc.prefetch_buffer.put(operation)
 
+        def _drain_ack_prefetch():
+            for ack in _drain_queue(cc.ack_prefetch_queue, n_ack_prefetch):
+                operation = ack.operation
+                if ack.completed_tokens is not None:
+                    if operation.request_id in self.ongoing_prefetch:
+                        assert operation.completed_tokens <= ack.completed_tokens
+                        operation.completed_tokens = ack.completed_tokens
+                if ack.pool_hits is not None:
+                    if operation.request_id in self.ongoing_prefetch:
+                        operation.pool_storage_result.update_extra_pool_hit_pages(
+                            ack.pool_hits
+                        )
+                        operation.pool_transfers_done = True
+                if ack.completed_req:
+                    if operation.request_id in self.ongoing_prefetch:
+                        self._handle_prefetch_result(operation)
+                    tail = operation.host_indices[operation.completed_tokens :]
+                    self.cache_controller.mem_pool_host.free(tail)
+
         def _drain_backup():
             for operation in _drain_queue(cc.ack_backup_queue, n_backup):
                 ack_id = operation.id
@@ -683,6 +700,7 @@ class HiRadixCache(RadixCache):
                 cc.mem_pool_host.free(host_indices)
 
         _drain_and_alloc_storage_hit()
+        _drain_ack_prefetch()
         _drain_backup()
         _drain_release()
 
@@ -1008,6 +1026,7 @@ class HiRadixCache(RadixCache):
         storage_queue_sizes = (
             (
                 cache_controller.prefetch_hit_queue.qsize(),
+                cache_controller.ack_prefetch_queue.qsize(),
                 cache_controller.ack_backup_queue.qsize(),
                 cache_controller.host_mem_release_queue.qsize(),
             )
@@ -1541,9 +1560,12 @@ class HiRadixCache(RadixCache):
             self.loading_check(finish_count=load_finish_count)
 
             if self.enable_storage and storage_queue_sizes:
-                n_storage_hit, n_backup, n_release = storage_queue_sizes[:3]
+                n_storage_hit, n_ack_prefetch, n_backup, n_release = (
+                    storage_queue_sizes[:4]
+                )
                 self._drain_storage_control_queues_impl(
                     n_storage_hit=n_storage_hit,
+                    n_ack_prefetch=n_ack_prefetch,
                     n_backup=n_backup,
                     n_release=n_release,
                     log_metrics=True,
@@ -1563,16 +1585,18 @@ class HiRadixCache(RadixCache):
         qsizes = torch.tensor(
             [
                 cc.prefetch_hit_queue.qsize(),
+                cc.ack_prefetch_queue.qsize(),
                 cc.ack_backup_queue.qsize(),
                 cc.host_mem_release_queue.qsize(),
             ],
             dtype=torch.int,
         )
-        self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN)
 
-        n_storage_hit, n_backup, n_release = map(int, qsizes.tolist())
+        self._all_reduce(qsizes, torch.distributed.ReduceOp.MIN)
+        n_storage_hit, n_ack_prefetch, n_backup, n_release = map(int, qsizes.tolist())
         self._drain_storage_control_queues_impl(
             n_storage_hit=n_storage_hit,
+            n_ack_prefetch=n_ack_prefetch,
             n_backup=n_backup,
             n_release=n_release,
             log_metrics=True,
@@ -1585,47 +1609,17 @@ class HiRadixCache(RadixCache):
         timeout = min(cfg.max, cfg.base + cfg.per_ki_token * num_tokens / 1024)
         return time.monotonic() - operation.start_time > timeout
 
-    def can_terminate_prefetch(self, operation: PrefetchOperation):
-        can_terminate = True
-
+    def can_terminate_prefetch(self, operation: PrefetchOperation) -> bool:
         if self.prefetch_stop_policy == "best_effort":
-            return can_terminate
-
-        if len(operation.hash_value) == 0:
-            completed = False
-        else:
-            completed = (
-                operation.completed_tokens == len(operation.hash_value) * self.page_size
-            )
-
+            return True
         if self.prefetch_stop_policy == "wait_complete":
-            can_terminate = completed
+            return False
         elif self.prefetch_stop_policy == "timeout":
-            can_terminate = completed or self.is_prefetch_timeout(operation)
+            return self.is_prefetch_timeout(operation)
         else:
             # 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)],
-            dtype=torch.int,
-        )
-        self._all_reduce_attn_groups(states, torch.distributed.ReduceOp.MAX)
-        can_terminate = states[0].item() == 0
-        operation_terminated = states[1].item() == 1
-        # the operation should be terminated if it is already terminated on any TP worker
-        # or it meets the termination condition on all TP workers
-        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:
@@ -1642,44 +1636,56 @@ class HiRadixCache(RadixCache):
             # there is no ongoing prefetch for this request or it has been revoked
             return True
 
-        last_host_node, prefetch_key, operation = self.ongoing_prefetch[req_id]
+        _, _, operation = self.ongoing_prefetch[req_id]
 
-        if not self.can_terminate_prefetch(operation):
+        # Determine whether or not we should terminate this prefetch request.  Make all
+        # ranks agree on the decision.  When running with PP, PPn will follow PP0's decision.
+        should_terminate = False
+        if self.pp_rank == 0:
+            should_terminate = operation.is_terminated() or self.can_terminate_prefetch(
+                operation
+            )
+        should_terminate_tensor = torch.tensor(
+            int(should_terminate), dtype=torch.int, device="cpu"
+        )
+        self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX)
+        should_terminate = should_terminate_tensor.item() == 1
+
+        if not should_terminate:
             return False
 
+        # Terminate in-flight prefetch.
+        self.cache_controller.terminate_prefetch(operation)
         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
+        else:
+            # Handle partial or full completion.
+            self._handle_prefetch_result(operation)
+        return True
 
-        completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
-            operation
+    def _handle_prefetch_result(self, operation: PrefetchOperation) -> None:
+        req_id = operation.request_id
+        # All PP/TP ranks will get the same `min_completed_tokens`, because `completed_tokens`
+        # and `pool_hits` in their operations are same.  No need to sync cross-rank here.
+        min_completed_tokens = self._clamp_prefetch_result(operation)
+        logger.debug(
+            f"Prefetch {req_id} completed with {operation.completed_tokens} tokens"
         )
-        logger.debug(f"Prefetch {req_id} completed with {completed_tokens} tokens")
-
-        min_completed_tokens = self._sync_and_clamp_prefetch_result(
-            operation, completed_tokens
-        )
-
+        last_host_node, prefetch_key, operation = self.ongoing_prefetch.pop(req_id)
+        host_indices = operation.host_indices
         fetched_key = prefetch_key[:min_completed_tokens]
-        written_indices = operation.host_indices[:min_completed_tokens]
+        written_indices = host_indices[:min_completed_tokens]
         matched_length = self._insert_helper_host(
             last_host_node,
             fetched_key,
             written_indices,
-            hash_value[: min_completed_tokens // self.page_size],
+            operation.hash_value[: min_completed_tokens // self.page_size],
         )
 
-        self.cache_controller.mem_pool_host.free(
-            operation.host_indices[:matched_length]
-        )
-        self.cache_controller.append_host_mem_release(
-            operation.host_indices[min_completed_tokens:completed_tokens]
-        )
+        self.cache_controller.mem_pool_host.free(host_indices[:matched_length])
         last_host_node.release_host()
-        del self.ongoing_prefetch[req_id]
         self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
 
         # Track tokens actually loaded from storage for this request (L3 hits)
@@ -1688,21 +1694,19 @@ class HiRadixCache(RadixCache):
 
         if self.enable_storage_metrics:
             self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage)
+        return
 
-        return True
-
-    def _sync_and_clamp_prefetch_result(
+    def _clamp_prefetch_result(
         self,
         operation: PrefetchOperation,
-        completed_tokens: int,
     ) -> int:
-        """Sync prefetch results across ATTN groups and decide the usable prefix.
+        """Determine the minimal number of tokens from full KV hits and sidecar hits.
 
         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.
+        this is just 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.
@@ -1710,27 +1714,17 @@ class HiRadixCache(RadixCache):
         hit_pages = (
             operation.pool_storage_result.extra_pool_hit_pages if pool_transfers else {}
         )
+        completed_tokens = operation.completed_tokens
         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
+        usable_pages = 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
-
-        _, _, operation = self.ongoing_prefetch[req_id]
-        operation.mark_terminate()
-
     def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
         """
         Pop and return the number of tokens loaded from storage for a request.
@@ -2017,7 +2011,6 @@ class HiRadixCache(RadixCache):
             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(

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

@@ -19,6 +19,7 @@ from sglang.srt.managers.cache_controller import (
 )
 from sglang.srt.managers.cache_controller import (
     LayerDoneCounter,
+    PrefetchAck,
 )
 from sglang.srt.managers.cache_controller import (
     StorageOperation as BaseStorageOperation,
@@ -29,6 +30,7 @@ from sglang.srt.mem_cache.hicache_storage import (
     PoolName,
     PoolTransfer,
     PoolTransferResult,
+    count_pool_hits,
 )
 from sglang.srt.mem_cache.l2_transfer import L2Transfer
 from sglang.srt.mem_cache.memory_pool_host import HostPoolGroup, PoolEntry
@@ -78,19 +80,13 @@ class PrefetchOperation(StorageOperation):
         )
         self.pool_transfers_done = not bool(pool_transfers)
 
-    def increment(self, num_tokens: int):
-        with self._lock:
-            if self._terminated_flag:
-                return False
-            self.completed_tokens += num_tokens
-            return True
-
     def mark_terminate(self):
         with self._lock:
             self._terminated_flag = True
 
     def is_terminated(self) -> bool:
-        return self._terminated_flag
+        with self._lock:
+            return self._terminated_flag
 
 
 class HybridCacheController(BaseHiCacheController):
@@ -640,26 +636,50 @@ class HybridCacheController(BaseHiCacheController):
                 )
         return host_indices, device_indices, resolved_pool_transfers
 
-    def _page_transfer(self, operation):
-        # KV pools first — determines actual completed page count
-        super()._page_transfer(operation)
+    def _page_transfer(self, operation: PrefetchOperation) -> bool:
+        # KV pools and KV-derived pools first — determines actual completed page count
+        kv_completed_pages = super()._page_transfer(operation)
+
+        # Read non-KV derived sidecar pool, e.g. SWA, Mamba.
+        self._page_transfer_sidecar(operation, kv_completed_pages)
+
+    def _page_transfer_sidecar(
+        self, operation: PrefetchOperation, kv_completed_pages: int
+    ) -> None:
+        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)
+        pool_hits: dict[str, int] = {}
+        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)
+            results = self.storage_backend.batch_get_v2(transfers_nonkv)
+            pool_hits = count_pool_hits(results)
+        # Emit PrefetchAck to prefetch_sync_queue, even the operation has been canceled by the
+        # scheduler thread.  The prefetch sync thread expects the same number of PrefetchAck objects
+        # to perform all_reduce.
+        self.prefetch_sync_queue.put(
+            PrefetchAck(
+                rid=operation.request_id,
+                operation=operation,
+                pool_hits=pool_hits,
+            )
+        )
+        return
 
     def _page_backup(self, operation):
         # MLA KV is replicated across TP ranks and should still be written only
@@ -671,9 +691,11 @@ class HybridCacheController(BaseHiCacheController):
         ]
 
         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)
+            pool_hits = count_pool_hits(results)
+            operation.pool_storage_result.update_extra_pool_hit_pages(pool_hits)
 
         if not self.backup_skip:
             super()._page_backup(operation)
@@ -737,7 +759,14 @@ class HybridCacheController(BaseHiCacheController):
             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
@@ -760,9 +789,7 @@ class HybridCacheController(BaseHiCacheController):
                 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,
@@ -888,3 +915,23 @@ class HybridCacheController(BaseHiCacheController):
             pool.host_indices = source.host_indices
             pool.device_indices = source.device_indices
         return extra_pools
+
+    def _reduce_prefetch_ack(self, ack: PrefetchAck) -> None:
+        # Handle KV-derived pool.
+        super()._reduce_prefetch_ack(ack)
+
+        # Handle other sidecar pools, e.g. SWA, Mamba.
+        if ack.pool_hits is not None:
+            # "for ... in PoolName" ensures the same order across all ranks.
+            # On prefetch failure, pool_hits may be empty dict.
+            packed = torch.tensor(
+                [ack.pool_hits.get(pool.value, 0) for pool in PoolName],
+                dtype=torch.int,
+            )
+            self._all_reduce(
+                packed,
+                torch.distributed.ReduceOp.MIN,
+                self.prefetch_completion_sync_groups,
+            )
+            for i, pool in enumerate(PoolName):
+                ack.pool_hits[pool.value] = packed[i].item()

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -410,10 +410,14 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
                     "Mooncake package does not support ReplicateConfig.group_ids. "
                     "Falling back to the existing batch_put_from path."
                 )
-            tp_scale_factor = 1 if storage_config is None else storage_config.tp_size
+            rank_scale_factor = (
+                1
+                if storage_config is None
+                else (storage_config.tp_size * storage_config.pp_size)
+            )
 
-            per_tp_global_segment_size = (
-                self.config.global_segment_size // tp_scale_factor
+            per_rank_global_segment_size = (
+                self.config.global_segment_size // rank_scale_factor
             )
 
             # Use the backend tag and model name as a prefix to isolate tenants
@@ -510,7 +514,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
                         ret_code = self.store.setup(
                             client_hostname,
                             self.config.metadata_server,
-                            per_tp_global_segment_size,
+                            per_rank_global_segment_size,
                             DEFAULT_LOCAL_BUFFER_SIZE,  # Zero copy interface does not need local buffer
                             self.config.protocol,
                             device_name,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -5,7 +5,7 @@ import logging
 import threading
 import time
 from dataclasses import replace
-from queue import Empty, Queue
+from queue import Queue
 from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, Sequence, TypeVar
 
 import torch
@@ -35,7 +35,6 @@ from sglang.srt.mem_cache.buffer_mode.storage_existence_cache import (
 )
 from sglang.srt.mem_cache.common import RetractionBackup
 from sglang.srt.mem_cache.hicache_storage import (
-    PoolHitPolicy,
     PoolName,
     PoolTransfer,
     SidecarPoolSpec,
@@ -1717,64 +1716,67 @@ class UnifiedRadixCache(BasePrefixCache):
     def can_terminate_prefetch(self, operation: PrefetchOperation) -> bool:
         if self.prefetch_stop_policy == "best_effort":
             return True
-
-        if len(operation.hash_value) == 0:
-            completed = False
-        else:
-            completed = (
-                operation.completed_tokens == len(operation.hash_value) * self.page_size
-            )
-
         if self.prefetch_stop_policy == "wait_complete":
-            can_terminate = completed
+            return False
         elif self.prefetch_stop_policy == "timeout":
-            can_terminate = completed or self._prefetch_timeout_check_linear_func(
-                operation
-            )
+            return self._prefetch_timeout_check_linear_func(operation)
         else:
             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)],
-            dtype=torch.int,
-        )
-        self._all_reduce_attn_groups(states, torch.distributed.ReduceOp.MAX)
-        can_terminate = states[0].item() == 0
-        operation_terminated = states[1].item() == 1
-        return can_terminate or operation_terminated
 
     @rank_consensus(same_params=True, same_results=True)
     def check_prefetch_progress(self, req_id: str) -> bool:
         if req_id not in self.ongoing_prefetch:
             return True
 
+        _, _, _, operation, _, _ = self.ongoing_prefetch[req_id]
+
+        # Determine whether or not we should terminate this prefetch request.  Make all
+        # ranks agree on the decision.  When running with PP, PPn will follow PP0's decision.
+        should_terminate = False
+        if self.pp_rank == 0:
+            should_terminate = operation.is_terminated() or self.can_terminate_prefetch(
+                operation
+            )
+        should_terminate_tensor = torch.tensor(
+            int(should_terminate), dtype=torch.int, device="cpu"
+        )
+        self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX)
+        should_terminate = should_terminate_tensor.item() == 1
+
+        if not should_terminate:
+            return False
+
+        self.cache_controller.terminate_prefetch(operation)
+        if operation.host_indices is None:
+            self.revoke_pending_prefetch(req_id)
+        else:
+            self._handle_prefetch_result(operation)
+        return True
+
+    def _handle_prefetch_result(self, operation: PrefetchOperation) -> None:
+        # This function **owns**:
+        # - host_indices[0 : completed_tokens]
+        # - sidecar pool hits if operation.pool_transfers_done is true
+        #
+        # That is, when this function returns the host memory referenced must be inserted
+        # into the radix tree or released to pool.
+
+        req_id = operation.request_id
+        completed_tokens = operation.completed_tokens
+        hash_value = operation.hash_value
+
         (
             last_host_node_id,
             prefetch_key,
             host_indices,
-            operation,
+            _,
             anchor_lock_params,
             comp_xfers,
         ) = self.ongoing_prefetch[req_id]
-        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
-        )
-
-        min_completed_tokens = self._sync_and_check_hybrid_prefetch_result(
+        # All PP/TP ranks will get the same `min_completed_tokens`, because `completed_tokens`
+        # and `pool_hits` in their operations are same.  No need to sync cross-rank here.
+        if not self._check_hybrid_prefetch_result(
             req_id,
             operation,
             completed_tokens,
@@ -1783,27 +1785,23 @@ class UnifiedRadixCache(BasePrefixCache):
             last_host_node_id,
             anchor_lock_params,
             prefetch_key,
-        )
-        if min_completed_tokens is None:
+        ):
             # Hybrid all-or-nothing check failed; result already discarded.
-            return True
+            return
 
         if self.buffer_pipeline is not None:
             # No graft: release the rank-local tail beyond the synced usable
             # length, then park the bounce for admission-time consumption.
-            self.cache_controller.append_host_mem_release(
-                host_indices[min_completed_tokens:completed_tokens]
-            )
             return self.buffer_pipeline.stage_completed_prefetch(
-                req_id, min_completed_tokens, hash_value
+                req_id, completed_tokens, hash_value
             )
 
-        fetched_key = prefetch_key[:min_completed_tokens]
+        fetched_key = prefetch_key[:completed_tokens]
         insert_result = self.tree_core.insert_host(
             last_host_node_id,
             fetched_key,
-            host_indices[:min_completed_tokens],
-            hash_value[: min_completed_tokens // self.page_size],
+            host_indices[:completed_tokens],
+            hash_value[: completed_tokens // self.page_size],
         )
 
         # Apply the host-insert walk's actions before the transfer commit.
@@ -1815,7 +1813,6 @@ class UnifiedRadixCache(BasePrefixCache):
                 extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
             )
             loaded_from_storage = 0
-            released_tokens = completed_tokens
         else:
             commit_actions: list[CacheAction | ComponentAction] = []
             self.tree_core.commit_hicache_transfers(
@@ -1833,11 +1830,7 @@ class UnifiedRadixCache(BasePrefixCache):
             self.cache_controller.mem_pool_host.free(
                 host_indices[: insert_result.prefix_len]
             )
-            self.cache_controller.append_host_mem_release(
-                host_indices[min_completed_tokens:completed_tokens]
-            )
-            loaded_from_storage = min_completed_tokens - insert_result.prefix_len
-            released_tokens = completed_tokens - min_completed_tokens
+            loaded_from_storage = completed_tokens - insert_result.prefix_len
 
         self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
         del self.ongoing_prefetch[req_id]
@@ -1845,21 +1838,19 @@ class UnifiedRadixCache(BasePrefixCache):
 
         self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage
         logger.info(
-            "HiCache prefetch %s req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d released=%d occupied=%d",
+            "HiCache prefetch %s req=%s completed=%d matched=%d loaded=%d occupied=%d",
             "dropped" if insert_result.host_insert_dropped else "success",
             req_id,
             completed_tokens,
-            min_completed_tokens,
             insert_result.prefix_len,
             loaded_from_storage,
-            released_tokens,
             self.cache_controller.prefetch_tokens_occupied,
         )
         if self.enable_storage_metrics and self.storage_metrics_collector is not None:
             self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage)
-        return True
+        return
 
-    def _sync_and_check_hybrid_prefetch_result(
+    def _check_hybrid_prefetch_result(
         self,
         req_id: str,
         operation: PrefetchOperation,
@@ -1869,8 +1860,8 @@ class UnifiedRadixCache(BasePrefixCache):
         last_host_node_id: NodeId,
         anchor_lock_params: DecLockRefParams,
         prefetch_key: RadixKey,
-    ) -> Optional[int]:
-        """Sync prefetch results across ATTN groups and decide the usable prefix.
+    ) -> bool:
+        """Decide the length of usable prefix.
 
         Two strategies depending on the hybrid layout:
 
@@ -1882,41 +1873,29 @@ class UnifiedRadixCache(BasePrefixCache):
           *all-or-nothing*. Their pools only cover a window / tail and cannot be
           truncated page by page, so any shortfall discards the whole prefetch.
 
-        Returns the synced usable token count (possibly clamped, possibly 0), or
-        ``None`` when an all-or-nothing prefetch was discarded (the caller should
-        then treat the prefetch as finished).
+        Returns true if prefetch success, or false when an all-or-nothing prefetch
+        was discarded (the caller should then treat the prefetch as finished).
         """
         # 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 = operation.pool_transfers or []
+        #
+        # Skip KV-derived pools, which do not report hits in operation.pool_storage_result.
+        # Their hit lengths are stored in completed_tokens.
+        pool_transfers = [
+            transfer
+            for transfer in operation.pool_transfers or []
+            if transfer.indices_from_pool != PoolName.KV
+        ]
         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()))
-        for transfer, count in zip(pool_transfers, pool_hit_pages):
-            hit_pages[transfer.name] = count
-
-        # DSA-style clamp: every sidecar is KV-derived and required for the whole
-        # prefix (ALL_PAGES), so the usable length is simply the shared minimum of
-        # the Full KV completion and each sidecar hit.
-        clampable = bool(pool_transfers) and all(
-            t.hit_policy == PoolHitPolicy.ALL_PAGES
-            and t.indices_from_pool == PoolName.KV
-            for t in pool_transfers
-        )
-        if clampable:
-            usable_pages = min(min_completed_tokens // self.page_size, *pool_hit_pages)
-            return usable_pages * self.page_size
-
+        completed_tokens = operation.completed_tokens
         # Hybrid cache state is all-or-nothing: every extra pool (SWA / Mamba / ...)
         # must cover the same fetched prefix. If any pool falls short the whole
         # prefetch result is unusable, so discard it and release everything.
         expected_tokens = len(hash_value) * self.page_size
-        all_succeeded = min_completed_tokens == expected_tokens and all(
+        all_succeeded = completed_tokens == expected_tokens and all(
             transfer.keys is not None and count == len(transfer.keys)
             for transfer, count in zip(pool_transfers, pool_hit_pages)
         )
@@ -1925,14 +1904,14 @@ class UnifiedRadixCache(BasePrefixCache):
             # tail (host_indices[completed_tokens:])
             self.cache_controller.append_host_mem_release(
                 host_indices=host_indices[:completed_tokens],
-                extra_pools=pool_transfers,
+                extra_pools=pool_transfers if operation.pool_transfers_done else None,
             )
             if anchor_lock_params is not None:
                 self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
-            del self.ongoing_prefetch[req_id]
             if self.buffer_pipeline is not None:
                 self.buffer_pipeline.pop_prefix_ctx(req_id)
                 self.buffer_pipeline.release_anchor_lock(req_id)
+            del self.ongoing_prefetch[req_id]
             self.cache_controller.prefetch_tokens_occupied -= (
                 self._prefetch_occupied_span(prefetch_key, host_indices)
             )
@@ -1943,14 +1922,8 @@ class UnifiedRadixCache(BasePrefixCache):
                 completed_tokens,
                 expected_tokens,
             )
-            return None
-        return min_completed_tokens
-
-    def terminate_prefetch(self, req_id: str) -> None:
-        if req_id not in self.ongoing_prefetch:
-            return
-        operation = self.ongoing_prefetch[req_id].operation
-        operation.mark_terminate()
+            return False
+        return True
 
     def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
         return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)
@@ -1994,16 +1967,16 @@ class UnifiedRadixCache(BasePrefixCache):
             return
 
         completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
-        self._barrier_attn_groups()
         if anchor_lock_params is not None:
             self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
         del self.ongoing_prefetch[rid]
         if self.buffer_pipeline is not None:
             self.buffer_pipeline.pop_prefix_ctx(rid)
             self.buffer_pipeline.release_anchor_lock(rid)
+        pool_transfers = [x for xfers in comp_xfers.values() for x in xfers]
         self.cache_controller.append_host_mem_release(
             host_indices=host_indices[:completed_tokens],
-            extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
+            extra_pools=pool_transfers if operation.pool_transfers_done else None,
         )
         # Buffer mode granted occupancy at hit-alloc, sized to the bounce;
         # cache mode reserved the requested span at enqueue.
@@ -2099,6 +2072,7 @@ class UnifiedRadixCache(BasePrefixCache):
     def _drain_storage_control_queues_impl(
         self,
         n_storage_hit: Optional[int],
+        n_ack_prefetch: Optional[int],
         n_backup: Optional[int],
         n_release: Optional[int],
         extra_release_counts: Optional[dict[PoolName, int]],
@@ -2106,15 +2080,25 @@ class UnifiedRadixCache(BasePrefixCache):
     ) -> None:
         cc = self.cache_controller
 
-        def _drain_queue(q: Queue[T], limit: Optional[int]) -> Iterator[T]:
-            drained = 0
-            while limit is None or drained < limit:
-                try:
-                    item = q.get_nowait()
-                except Empty:
-                    break
-                drained += 1
-                yield item
+        def _drain_queue(q: Queue[T], n: Optional[int]) -> Iterator[T]:
+            """If n is None, consume all items from the queue.
+            Otherwise, consume n items from the queue.  Blocking if there are no enough n items.
+
+            In TP, each rank consumes the a minimal number of items of all ranks.
+            In PP, each rank consumes the exact number of items of PP0.  Refer to _pp_sync for more details.
+
+            This prevents TP/PP divergence.
+            """
+            if n is None:
+                while not q.empty():
+                    item = q.get()
+                    yield item
+            else:
+                for _ in range(n):
+                    # Block when there are not enough elements.
+                    # All TP/PP ranks must consume the same number of elements.
+                    item = q.get()
+                    yield item
 
         buffer_mode = self.host_memory_mode == "buffer_only"
 
@@ -2202,6 +2186,33 @@ class UnifiedRadixCache(BasePrefixCache):
                     self._prefetch_outcome_stats["declined_rate_limited"] += 1
                     self.buffer_pipeline.pending_hit_allocs.append(operation)
 
+        def _drain_ack_prefetch():
+            for ack in _drain_queue(cc.ack_prefetch_queue, n_ack_prefetch):
+                operation = ack.operation
+                if ack.completed_tokens is not None:
+                    if operation.request_id in self.ongoing_prefetch:
+                        assert operation.completed_tokens <= ack.completed_tokens
+                        operation.completed_tokens = ack.completed_tokens
+                if ack.pool_hits is not None:
+                    if operation.request_id in self.ongoing_prefetch:
+                        operation.pool_storage_result.update_extra_pool_hit_pages(
+                            ack.pool_hits
+                        )
+                        operation.pool_transfers_done = True
+                if ack.completed_req:
+                    if operation.request_id in self.ongoing_prefetch:
+                        # check_prefetch_progress() is not called for this rid yet.
+                        # Let us insert the prefetch result into the radix tree.
+                        self._handle_prefetch_result(operation)
+                    cc.append_host_mem_release(
+                        operation.host_indices[operation.completed_tokens :],
+                        (
+                            operation.pool_transfers
+                            if not operation.pool_transfers_done
+                            else None
+                        ),
+                    )
+
         def _drain_backup():
             drained = 0
             for operation in _drain_queue(cc.ack_backup_queue, n_backup):
@@ -2255,6 +2266,7 @@ class UnifiedRadixCache(BasePrefixCache):
             return drained
 
         _drain_and_alloc_storage_hit()
+        _drain_ack_prefetch()
         _drain_backup()
         _drain_release()
         _drain_extra_release()
@@ -2265,6 +2277,7 @@ class UnifiedRadixCache(BasePrefixCache):
         extra_pool_names = list(extra_release_queues)
         local_qsize_list = [
             cc.prefetch_hit_queue.qsize(),
+            cc.ack_prefetch_queue.qsize(),
             cc.ack_backup_queue.qsize(),
             cc.host_mem_release_queue.qsize(),
             *[
@@ -2276,15 +2289,16 @@ class UnifiedRadixCache(BasePrefixCache):
             local_qsize_list,
             dtype=torch.int,
         )
-        self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN)
+        self._all_reduce(qsizes, torch.distributed.ReduceOp.MIN)
         qsize_list = list(map(int, qsizes.tolist()))
-        n_storage_hit, n_backup, n_release = qsize_list[:3]
+        n_storage_hit, n_ack_prefetch, 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_storage_hit=n_storage_hit,
+            n_ack_prefetch=n_ack_prefetch,
             n_backup=n_backup,
             n_release=n_release,
             extra_release_counts=extra_release_counts,
@@ -2306,6 +2320,7 @@ class UnifiedRadixCache(BasePrefixCache):
             return
         self._drain_storage_control_queues_impl(
             n_storage_hit=0,
+            n_ack_prefetch=0,
             n_backup=None,
             n_release=None,
             extra_release_counts={
@@ -2389,6 +2404,7 @@ class UnifiedRadixCache(BasePrefixCache):
             storage_queue_sizes = (
                 (
                     cc.prefetch_hit_queue.qsize(),
+                    cc.ack_prefetch_queue.qsize(),
                     cc.ack_backup_queue.qsize(),
                     cc.host_mem_release_queue.qsize(),
                     *(extra_release_queues[name].qsize() for name in extra_pool_names),
@@ -2612,16 +2628,19 @@ class UnifiedRadixCache(BasePrefixCache):
             self.loading_check(finish_count=load_finish_count)
 
             if self.enable_storage and storage_queue_sizes:
-                n_storage_hit, n_backup, n_release = storage_queue_sizes[:3]
+                n_storage_hit, n_ack_prefetch, n_backup, n_release = (
+                    storage_queue_sizes[:4]
+                )
                 extra_release_counts = {
                     pool_name: count
                     for pool_name, count in zip(
                         extra_pool_names,
-                        storage_queue_sizes[3:],
+                        storage_queue_sizes[4:],
                     )
                 }
                 self._drain_storage_control_queues_impl(
                     n_storage_hit=n_storage_hit,
+                    n_ack_prefetch=n_ack_prefetch,
                     n_backup=n_backup,
                     n_release=n_release,
                     extra_release_counts=extra_release_counts,