fix(hicache): TP/PP write-through & load-back consensus (#28429)

python/sglang/srt/managers/cache_controller.py

@@ -20,6 +20,7 @@
 from queue import Empty, Queue
 from typing import TYPE_CHECKING, List, NamedTuple, Optional
 
+import msgspec
 import torch
 
 from sglang.srt.mem_cache.hicache_storage import (
@@ -116,6 +117,25 @@ def reset(self):
         self.consumer_index = -1
 
 
+class WriteReservation(msgspec.Struct, frozen=True):
+    """Host-side reservation for a two-phase write-through backup.
+
+    Produced by ``reserve_write`` (allocation only, no DMA/ack) and consumed by
+    either ``commit_write`` (schedule the backup DMA) or ``abort_write`` (free
+    the reservation). Splitting the host allocation from the DMA schedule lets
+    TP/PP ranks reach a collective consensus on the enqueue decision BEFORE any
+    rank mutates ``ongoing_write_through`` / ``ack_write_queue`` -- see
+    ``HiRadixCache.write_backup`` and sglang#28429.
+    """
+
+    host_indices: torch.Tensor
+    device_indices: torch.Tensor
+    node_id: int = -1
+    priority: Optional[int] = None
+    # Only used by HybridCacheController (indexer / aux pool transfers).
+    pool_transfers: Optional[list] = None
+
+
 class CacheOperation:
 
     counter = 0
@@ -670,23 +690,58 @@ def reset(self):
             self.prefetch_thread.start()
             self.backup_thread.start()
 
-    def write(
+    def reserve_write(
         self,
         device_indices: torch.Tensor,
         priority: Optional[int] = None,
         node_id: int = -1,
-    ) -> Optional[torch.Tensor]:
-        """
-        Back up KV caches from device memory to host memory.
+    ) -> Optional[WriteReservation]:
+        """Phase 1 of a two-phase write-through backup: reserve host slots only,
+        WITHOUT scheduling the DMA or enqueuing an ack. Lets TP/PP ranks agree
+        on whether every rank secured host memory before any rank commits.
+        Returns None if this rank could not allocate.
         """
         host_indices = self.mem_pool_host.alloc(len(device_indices))
         if host_indices is None:
             return None
+        return WriteReservation(
+            host_indices=host_indices,
+            device_indices=device_indices,
+            node_id=node_id,
+            priority=priority,
+        )
+
+    def commit_write(self, reservation: WriteReservation) -> torch.Tensor:
+        """Phase 2: schedule the reserved backup DMA and enqueue its ack."""
         self.write_queue.append(
-            CacheOperation(host_indices, device_indices, node_id, priority)
+            CacheOperation(
+                reservation.host_indices,
+                reservation.device_indices,
+                reservation.node_id,
+                reservation.priority,
+            )
         )
         self.start_writing()
-        return host_indices
+        return reservation.host_indices
+
+    def abort_write(self, reservation: WriteReservation) -> None:
+        """Undo a reservation the group rejected: free the host slots. No DMA or
+        ack was created for it, so nothing else needs unwinding."""
+        self.mem_pool_host.free(reservation.host_indices)
+
+    def write(
+        self,
+        device_indices: torch.Tensor,
+        priority: Optional[int] = None,
+        node_id: int = -1,
+    ) -> Optional[torch.Tensor]:
+        """
+        Back up KV caches from device memory to host memory.
+        """
+        reservation = self.reserve_write(device_indices, priority, node_id)
+        if reservation is None:
+            return None
+        return self.commit_write(reservation)
 
     def start_writing(self) -> None:
         if len(self.write_queue) == 0:

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -839,27 +839,48 @@ def write_backup(self, node: TreeNode, write_back=False) -> int:
         ):
             return 0
 
-        host_indices = self.cache_controller.write(
+        # Phase 1: reserve host slots only (no DMA / ack yet), so TP/PP ranks can
+        # agree on the enqueue decision before any rank mutates shared state.
+        reservation = self.cache_controller.reserve_write(
             device_indices=node.value,
             node_id=node.id,
             **self._get_extra_pools(),
         )
-        if host_indices is None:
+        if reservation is None:
             self.evict_host(len(node.value))
-            host_indices = self.cache_controller.write(
+            reservation = self.cache_controller.reserve_write(
                 device_indices=node.value,
                 node_id=node.id,
                 **self._get_extra_pools(),
             )
-        if host_indices is not None:
-            node.host_value = host_indices.clone()
-            assert len(node.host_value) > 0
-            self._track_write_through_node(node, len(node.key))
-            if not write_back:
-                self.inc_lock_ref(node)
-        else:
+
+        # Phase 2 (write-through, TP/PP > 1): reach a cross-rank consensus so
+        # every rank makes the SAME enqueue decision. Host-pool occupancy is
+        # per-rank physical state, so reserve_write can succeed on some ranks and
+        # fail on others; committing on only a subset diverges node.backuped /
+        # ongoing_write_through / ack_write_queue and desyncs the TP forward
+        # collectives (the write_through HiCache hang, sglang#28429). Commit only
+        # if EVERY rank secured its reservation; otherwise every rank aborts --
+        # skipping a write-through is harmless, the KV simply stays on device.
+        if not write_back and (self.tp_world_size > 1 or self.pp_size > 1):
+            reserved = torch.tensor(
+                1 if reservation is not None else 0, dtype=torch.int, device="cpu"
+            )
+            self._all_reduce(reserved, torch.distributed.ReduceOp.MIN)
+            if reserved.item() == 0:
+                if reservation is not None:
+                    self.cache_controller.abort_write(reservation)
+                return 0
+
+        if reservation is None:
             return 0
 
+        host_indices = self.cache_controller.commit_write(reservation)
+        node.host_value = host_indices.clone()
+        assert len(node.host_value) > 0
+        self._track_write_through_node(node, len(node.key))
+        if not write_back:
+            self.inc_lock_ref(node)
         return len(host_indices)
 
     def _track_write_through_node(self, node: TreeNode, backup_len: int) -> None:
@@ -1312,45 +1333,79 @@ def load_back(
 
         # load it all or not at all
         host_indices = torch.cat([n.host_value for n in nodes_to_load])
-        if len(host_indices) < self.load_back_threshold or (
-            len(host_indices) > mem_quota + delta if mem_quota is not None else False
-        ):
-            # skip loading back if the total size is too small or exceeding the memory quota
-            self.dec_lock_ref(ancester_node)
-            return None
+
+        # Local viability: too small to bother, or would exceed this rank's quota.
+        num_tokens = len(host_indices)
+        too_small = num_tokens < self.load_back_threshold
+        exceeds_quota = mem_quota is not None and num_tokens > mem_quota + delta
+        local_ok = not (too_small or exceeds_quota)
 
         # Protect the nodes being loaded from host eviction.
         for n in nodes_to_load:
             n.protect_host()
 
-        device_indices = self.cache_controller.load(
-            host_indices=host_indices,
-            node_id=last_hit_node.id,
-            **self._get_extra_pools(),
-        )
-        if device_indices is None:
-            self.evict(EvictParams(num_tokens=len(host_indices)))
+        multi_rank = self.tp_world_size > 1 or self.pp_size > 1
+        device_indices = None
+        if local_ok:
             device_indices = self.cache_controller.load(
                 host_indices=host_indices,
                 node_id=last_hit_node.id,
                 **self._get_extra_pools(),
             )
-        self.dec_lock_ref(ancester_node)
-        if device_indices is None:
-            # no sufficient GPU memory to load back KV caches
-            for n in nodes_to_load:
-                n.release_host()
-            logger.warning(
-                "load_back: FAILED to load %d tokens for node %d "
-                "even after eviction (evictable_size=%d)",
-                len(host_indices),
-                last_hit_node.id,
-                self.evictable_size_,
-            )
-            return None
+            if device_indices is None and not multi_rank:
+                # Single rank: safe to evict-and-retry. Under TP/PP > 1 we do NOT
+                # evict here -- a conditional, per-rank eviction would mutate the
+                # radix tree on only the ranks that failed the first alloc and
+                # re-introduce the very cross-rank divergence this fix removes.
+                # load-back is best-effort, so instead every rank skips this
+                # round together via the consensus below.
+                self.evict(EvictParams(num_tokens=len(host_indices)))
+                device_indices = self.cache_controller.load(
+                    host_indices=host_indices,
+                    node_id=last_hit_node.id,
+                    **self._get_extra_pools(),
+                )
+            local_ok = device_indices is not None
+
+        # Transactional consensus (sglang#28429, load side): load back on EVERY
+        # rank or on none. device-pool occupancy is per-rank physical state, so
+        # loading on a subset diverges prefix_indices and the next forward's
+        # collective shapes mismatch -> TP hang. If any rank cannot load, all
+        # ranks roll back to the pre-load state (radix tree left untouched).
+        if multi_rank:
+            loaded = torch.tensor(1 if local_ok else 0, dtype=torch.int, device="cpu")
+            self._all_reduce(loaded, torch.distributed.ReduceOp.MIN)
+            group_ok = loaded.item() == 1
+        else:
+            group_ok = local_ok
 
+        self.dec_lock_ref(ancester_node)
         for n in nodes_to_load:
             n.release_host()
+
+        if not group_ok:
+            # Roll back anything this rank did; leave the tree unchanged so it
+            # stays identical across ranks.
+            if device_indices is not None:
+                self.cache_controller.mem_pool_device_allocator.free(device_indices)
+            if multi_rank:
+                logger.debug(
+                    "load_back: group consensus rejected loading %d tokens for node %d",
+                    len(host_indices),
+                    last_hit_node.id,
+                )
+            else:
+                # no sufficient GPU memory to load back KV caches
+                logger.warning(
+                    "load_back: FAILED to load %d tokens for node %d "
+                    "even after eviction (evictable_size=%d)",
+                    len(host_indices),
+                    last_hit_node.id,
+                    self.evictable_size_,
+                )
+            return None
+
+        # All ranks succeeded -> finalize (device_indices is not None on all).
         self.ongoing_load_back[last_hit_node.id] = last_hit_node
         offset = 0
         for node in nodes_to_load:

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

@@ -14,6 +14,9 @@
 from sglang.srt.managers.cache_controller import (
     HiCacheAck,
 )
+from sglang.srt.managers.cache_controller import (
+    WriteReservation,
+)
 from sglang.srt.managers.cache_controller import (
     HiCacheController as BaseHiCacheController,
 )
@@ -361,13 +364,13 @@ def reset(self):
                 release_queue.queue.clear()
             self.prefetch_tokens_occupied = 0
 
-    def write(
+    def reserve_write(
         self,
         device_indices: torch.Tensor,
         priority: Optional[int] = None,
         node_id: int = -1,
         extra_pools: Optional[list[PoolTransfer]] = None,
-    ) -> Optional[torch.Tensor]:
+    ) -> Optional[WriteReservation]:
         host_indices = self.mem_pool_host.alloc(len(device_indices))
         if host_indices is None:
             return None
@@ -380,18 +383,55 @@ def write(
         if pool_transfers is None and extra_pools:
             self.mem_pool_host.free(host_indices)
             return None
+        return WriteReservation(
+            host_indices=host_indices,
+            device_indices=device_indices,
+            node_id=node_id,
+            priority=priority,
+            pool_transfers=pool_transfers,
+        )
 
+    def commit_write(self, reservation: WriteReservation) -> torch.Tensor:
         self.write_queue.append(
             CacheOperation(
-                host_indices,
-                device_indices,
-                node_id,
-                priority,
-                pool_transfers=pool_transfers or None,
+                reservation.host_indices,
+                reservation.device_indices,
+                reservation.node_id,
+                reservation.priority,
+                pool_transfers=reservation.pool_transfers or None,
             )
         )
         self.start_writing()
-        return host_indices
+        return reservation.host_indices
+
+    def abort_write(self, reservation: WriteReservation) -> None:
+        # Free the indexer / aux pool host slots we allocated in reserve_write
+        # (mirrors the alloc_host branch of _resolve_pool_transfers_allocation),
+        # then free the KV host slots. Derived transfers (indices_from_pool set)
+        # share the KV indices and are not separately allocated -> skip them.
+        pool_transfers = reservation.pool_transfers
+        if pool_transfers:
+            for pool in pool_transfers:
+                if pool.indices_from_pool is not None:
+                    continue
+                entry = self.mem_pool_host.entry_map.get(pool.name)
+                if entry is None or pool.host_indices is None:
+                    continue
+                entry.host_pool.free(pool.host_indices)
+                pool.host_indices = None
+        self.mem_pool_host.free(reservation.host_indices)
+
+    def write(
+        self,
+        device_indices: torch.Tensor,
+        priority: Optional[int] = None,
+        node_id: int = -1,
+        extra_pools: Optional[list[PoolTransfer]] = None,
+    ) -> Optional[torch.Tensor]:
+        reservation = self.reserve_write(device_indices, priority, node_id, extra_pools)
+        if reservation is None:
+            return None
+        return self.commit_write(reservation)
 
     def start_writing(self) -> None:
         if not self.write_queue: