Make CP HiCache backup admission deterministic

CP HiCache write-through under shared KV was still using rank-wide collectives to decide host reservation eviction, and per-layer backup registration could be bypassed before the final forward boundary. This moves backup registration to the final run_batch pre-forward boundary, forwards it through SessionAwareCache, exposes fallback paths as explicit warnings, and introduces deterministic owner-lane capacity planning for CP host reservation.

Constraint: CP shared-KV ranks must keep target and draft host reservations owner-lane consistent without adding hot-path collective synchronization

Constraint: Remote CUDA validation must run in the g0034 container, not locally

Rejected: Keep reserve_slots_max all_reduce as the default admission path | observed reserve collectives reaching double-digit and occasional 100ms+ latency

Rejected: Silent post-forward catch-up backup | hides when per-layer forward-overlap backup is not actually active

Confidence: medium

Scope-risk: broad

Directive: Do not reintroduce CP HiCache hot-path collectives without a measured mismatch case and explicit fallback warning

Tested: py_compile for modified Python modules and CP HiCache metadata test file in remote g0034 container

Tested: python3 -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q in remote g0034 container (75 passed, 5 warnings)

Tested: git diff --check HEAD~1..HEAD

Not-tested: Local pytest blocked by missing pybase64 in the local environment

Not-tested: Full CP HiCache + MTP E2E after the no-collective reservation change

Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/environ.py

@@ -215,6 +215,7 @@ class Envs:
     SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
     SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False)
     SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False)
+    SGLANG_CP_HICACHE_CAPACITY_DEBUG = EnvBool(False)
     SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)

python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py

@@ -38,8 +38,8 @@ def _prefetch_log(message: str, *args) -> None:
 
 
 def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None:
-    logger.info(
-        "CP shared KV index prefetch fallback (%s): " + message,
+    logger.warning(
+        "[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message,
         reason,
         *args,
     )

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -123,8 +123,8 @@ def _log_cp_shared_kv_index_prefetch_fallback(
     dedupe here or later real fallbacks become invisible during profiling.
     """
 
-    logger.info(
-        "CP shared KV index prefetch fallback (%s): " + message,
+    logger.warning(
+        "[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message,
         reason,
         *args,
     )

python/sglang/srt/layers/attention/nsa/utils.py

@@ -58,8 +58,8 @@ def log_cp_shared_kv_direct_write_fallback(
     Warmup can hit the same fallback reason as a later real request, so
     de-duplicating by reason hides correctness/performance issues after startup.
     """
-    logger.info(
-        "CP shared KV direct-write fallback (%s): " + message,
+    logger.warning(
+        "[CP_SHARED_KV_FALLBACK][direct_write] reason=%s " + message,
         reason,
         *args,
     )
@@ -821,8 +821,8 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor(
 
 
 def _log_tai_in_seq_rerange_fallback(reason: str, message: str, *args) -> None:
-    logger.info(
-        "CP NSA in-seq all-gather rerange tai-kernel fallback (%s): " + message,
+    logger.warning(
+        "[CP_SHARED_KV_FALLBACK][tai_in_seq_rerange] reason=%s " + message,
         reason,
         *args,
     )

python/sglang/srt/managers/cache_controller.py

@@ -1026,7 +1026,9 @@ class HiCacheController:
 
     def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None:
         logger.warning(
-            "[CacheCtrl-write] CP HiCache all-layer backup fallback: node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
+            "[CP_HICACHE_FALLBACK][submit_write_cp_all_layer] "
+            "CP HiCache all-layer backup fallback. "
+            "node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
             reservation.node_id,
             reservation.metadata.logical_len,
             reservation.metadata.owned_positions.numel(),
@@ -1250,6 +1252,16 @@ class HiCacheController:
             )
             return
 
+        logger.warning(
+            "[CP_HICACHE_FALLBACK][catch_up_all_layers] "
+            "CP HiCache write is submitting all layers after KV materialization; "
+            "forward-overlap per-layer backup was not active. "
+            "node_id=%d logical_len=%d layers=%d draft=%s",
+            reservation.node_id,
+            reservation.metadata.logical_len,
+            state.total_layers,
+            reservation.draft_host_indices is not None,
+        )
         total_layers = state.total_layers
         for layer_id in range(total_layers):
             self.submit_write_cp_layer(reservation, layer_id)

python/sglang/srt/managers/scheduler.py

@@ -2457,11 +2457,6 @@ class Scheduler(
             )
 
         new_batch.prepare_for_extend()
-        if self.enable_hierarchical_cache and hasattr(
-            self.tree_cache, "prepare_write_backup_for_req"
-        ):
-            for req in can_run_list:
-                self.tree_cache.prepare_write_backup_for_req(req)
 
         # Record prefill stats for logging after forward
         new_batch.prefill_stats = PrefillStats.from_adder(
@@ -2579,6 +2574,38 @@ class Scheduler(
         self.batch_record_ct = (self.batch_record_ct + 1) % 2
         self.batch_record_buf[self.batch_record_ct] = model_worker_batch
 
+    def _prepare_hicache_write_backups_before_forward(
+        self, batch: ScheduleBatch
+    ) -> None:
+        """Reserve/register CP HiCache write-through backups immediately before forward.
+
+        CP shared-KV HiCache per-layer D2H backup must be registered after
+        `prepare_for_extend()` has allocated output KV slots, but before the
+        model starts emitting per-layer KV.  Keeping this at the final
+        `run_batch` boundary prevents alternate prefill event loops from
+        bypassing the early-backup path and falling back to post-forward
+        catch-up copies.
+        """
+
+        if not self.enable_hierarchical_cache:
+            return
+        if batch.forward_mode not in (
+            ForwardMode.EXTEND,
+            ForwardMode.SPLIT_PREFILL,
+            ForwardMode.DLLM_EXTEND,
+        ):
+            return
+
+        prepare_fn = getattr(self.tree_cache, "prepare_write_backup_for_req", None)
+        if prepare_fn is None:
+            inner_cache = getattr(self.tree_cache, "inner", None)
+            prepare_fn = getattr(inner_cache, "prepare_write_backup_for_req", None)
+        if prepare_fn is None:
+            return
+
+        for req in batch.reqs:
+            prepare_fn(req)
+
     def run_batch(
         self,
         batch: ScheduleBatch,
@@ -2601,6 +2628,8 @@ class Scheduler(
         if batch.forward_mode.is_prebuilt():
             return self._run_batch_prebuilt(batch)
 
+        self._prepare_hicache_write_backups_before_forward(batch)
+
         # Run forward
         if self.is_generation:
             if self.spec_algorithm.is_none() or self.enable_overlap:

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -257,6 +257,37 @@ class PreparedCpHiCacheBackup:
     attached: bool = False
 
 
+@dataclass(frozen=True)
+class CpHiCacheCapacitySnapshot:
+    target_capacity: Tuple[int, ...]
+    draft_capacity: Optional[Tuple[int, ...]]
+    committed_target: Tuple[int, ...]
+    committed_draft: Tuple[int, ...]
+    pending_target: Tuple[int, ...]
+    pending_draft: Tuple[int, ...]
+
+    @property
+    def target_used(self) -> Tuple[int, ...]:
+        return tuple(
+            committed + pending
+            for committed, pending in zip(self.committed_target, self.pending_target)
+        )
+
+    @property
+    def draft_used(self) -> Tuple[int, ...]:
+        return tuple(
+            committed + pending
+            for committed, pending in zip(self.committed_draft, self.pending_draft)
+        )
+
+
+@dataclass(frozen=True)
+class CpHiCacheEvictionPlan:
+    victims: Tuple[TreeNode, ...]
+    planned_freed: Tuple[int, ...]
+    remaining_deficit: Tuple[int, ...]
+
+
 class HiCachePendingBackupSplit(Exception):
     def __init__(self, node: TreeNode):
         self.node = node
@@ -473,6 +504,14 @@ class HiRadixCache(RadixCache):
         # record the ongoing prefetch requests
         self.ongoing_prefetch = {}
         self.ongoing_backup = {}
+        # Temporary aggregate accounting for CP HiCache collectives.  These
+        # calls are on latency-sensitive scheduler paths; keep logging coarse
+        # enough to avoid per-request spam while still exposing hot collectives.
+        self._cp_hicache_collective_stats = {}
+        self._cp_hicache_capacity_debug = (
+            envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get()
+        )
+        self._cp_hicache_capacity_debug_stats = {}
         # track per-request tokens loaded from storage (L3 hits)
         # key: request_id, value: number of tokens actually loaded from storage
         self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
@@ -501,6 +540,441 @@ class HiRadixCache(RadixCache):
 
         super().__init__(params=params)
 
+    def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None:
+        start = time.perf_counter()
+        torch.distributed.all_reduce(tensor, op=op, group=group)
+        elapsed_ms = (time.perf_counter() - start) * 1000
+
+        if not hasattr(self, "_cp_hicache_collective_stats"):
+            self._cp_hicache_collective_stats = {}
+        stats = self._cp_hicache_collective_stats.setdefault(
+            tag,
+            {
+                "count": 0,
+                "total_ms": 0.0,
+                "max_ms": 0.0,
+            },
+        )
+        stats["count"] += 1
+        stats["total_ms"] += elapsed_ms
+        stats["max_ms"] = max(stats["max_ms"], elapsed_ms)
+
+        count = stats["count"]
+        if count == 1 or count % 128 == 0 or elapsed_ms >= 10.0:
+            logger.info(
+                "[HiCache-collective] tag=%s count=%d total_ms=%.3f avg_ms=%.3f max_ms=%.3f last_ms=%.3f",
+                tag,
+                count,
+                stats["total_ms"],
+                stats["total_ms"] / count,
+                stats["max_ms"],
+                elapsed_ms,
+            )
+
+    def _cp_hicache_cp_size(self, page_owners: Optional[torch.Tensor] = None) -> int:
+        layout = getattr(
+            getattr(self, "cache_controller", None), "cp_shared_kv_layout", None
+        )
+        cp_size = getattr(layout, "cp_size", None)
+        if cp_size is not None:
+            return int(cp_size)
+        tp_world_size = int(getattr(self, "tp_world_size", 1) or 1)
+        if tp_world_size > 1:
+            return tp_world_size
+        if page_owners is not None and page_owners.numel() > 0:
+            non_negative = page_owners.to(device="cpu", dtype=torch.int64)
+            non_negative = non_negative[non_negative >= 0]
+            if non_negative.numel() > 0:
+                return int(non_negative.max().item()) + 1
+        return 1
+
+    def _cp_hicache_cp_rank(self) -> int:
+        layout = getattr(
+            getattr(self, "cache_controller", None), "cp_shared_kv_layout", None
+        )
+        return int(getattr(layout, "cp_rank", getattr(self, "_tp_group_rank", 0)) or 0)
+
+    def _cp_zero_counts(self, cp_size: Optional[int] = None) -> Tuple[int, ...]:
+        cp_size = self._cp_hicache_cp_size() if cp_size is None else int(cp_size)
+        return tuple(0 for _ in range(cp_size))
+
+    @staticmethod
+    def _cp_add_counts(lhs: Tuple[int, ...], rhs: Tuple[int, ...]) -> Tuple[int, ...]:
+        if len(lhs) != len(rhs):
+            raise RuntimeError(
+                f"CP HiCache count length mismatch: {len(lhs)} vs {len(rhs)}"
+            )
+        return tuple(int(a) + int(b) for a, b in zip(lhs, rhs))
+
+    @staticmethod
+    def _cp_sub_counts_floor_zero(
+        lhs: Tuple[int, ...], rhs: Tuple[int, ...]
+    ) -> Tuple[int, ...]:
+        if len(lhs) != len(rhs):
+            raise RuntimeError(
+                f"CP HiCache count length mismatch: {len(lhs)} vs {len(rhs)}"
+            )
+        return tuple(max(0, int(a) - int(b)) for a, b in zip(lhs, rhs))
+
+    def _cp_owner_token_counts(
+        self,
+        page_owners: torch.Tensor,
+        *,
+        page_size: int,
+        cp_size: Optional[int] = None,
+    ) -> Tuple[int, ...]:
+        cp_size = (
+            self._cp_hicache_cp_size(page_owners)
+            if cp_size is None
+            else int(cp_size)
+        )
+        counts = [0 for _ in range(cp_size)]
+        owners = page_owners.to(device="cpu", dtype=torch.int64)
+        for owner in owners.tolist():
+            owner = int(owner)
+            if owner < 0:
+                continue
+            if owner >= cp_size:
+                raise RuntimeError(
+                    "CP HiCache metadata owner exceeds CP size: "
+                    f"owner={owner} cp_size={cp_size}"
+                )
+            counts[owner] += int(page_size)
+        return tuple(counts)
+
+    def _cp_metadata_token_counts(
+        self, metadata: CpHiCacheNodeMetadata
+    ) -> Tuple[int, ...]:
+        return self._cp_owner_token_counts(
+            metadata.page_owners,
+            page_size=metadata.page_size,
+        )
+
+    def _cp_assert_metadata_counts(
+        self, metadata: CpHiCacheNodeMetadata, *, context: str
+    ) -> Tuple[int, ...]:
+        counts = self._cp_metadata_token_counts(metadata)
+        cp_rank = self._cp_hicache_cp_rank()
+        if cp_rank >= len(counts):
+            raise RuntimeError(
+                "CP HiCache owner lane mismatch: "
+                f"context={context} cp_rank={cp_rank} cp_size={len(counts)}"
+            )
+        expected_local = counts[cp_rank]
+        actual_local = int(metadata.owned_positions.numel())
+        if actual_local != expected_local:
+            raise RuntimeError(
+                "CP HiCache owner lane mismatch: "
+                f"context={context} cp_rank={cp_rank} expected_local={expected_local} "
+                f"owned_positions={actual_local}"
+            )
+        if int(metadata.host_indices.numel()) != expected_local:
+            raise RuntimeError(
+                "CP HiCache host index owner lane mismatch: "
+                f"context={context} cp_rank={cp_rank} expected_local={expected_local} "
+                f"host_indices={metadata.host_indices.numel()}"
+            )
+        draft_host_indices = getattr(metadata, "draft_host_indices", None)
+        if draft_host_indices is not None and int(draft_host_indices.numel()) != expected_local:
+            raise RuntimeError(
+                "CP HiCache draft host index owner lane mismatch: "
+                f"context={context} cp_rank={cp_rank} expected_local={expected_local} "
+                f"draft_host_indices={draft_host_indices.numel()}"
+            )
+        return counts
+
+    def _cp_walk_radix_nodes(self) -> List[TreeNode]:
+        root = getattr(self, "root_node", None)
+        if root is None:
+            return []
+        nodes: List[TreeNode] = []
+        stack = list(getattr(root, "children", {}).values())
+        while stack:
+            node = stack.pop()
+            nodes.append(node)
+            stack.extend(getattr(node, "children", {}).values())
+        return nodes
+
+    def _cp_host_capacity_snapshot(self) -> CpHiCacheCapacitySnapshot:
+        cp_size = self._cp_hicache_cp_size()
+        target_capacity = tuple(
+            int(getattr(getattr(self, "token_to_kv_pool_host", None), "size", 0))
+            for _ in range(cp_size)
+        )
+        controller = getattr(self, "cache_controller", None)
+        draft_pool = getattr(controller, "draft_mem_pool_host", None)
+        has_draft = bool(getattr(controller, "has_draft_hicache", False))
+        draft_capacity = (
+            tuple(int(getattr(draft_pool, "size", 0)) for _ in range(cp_size))
+            if has_draft
+            else None
+        )
+        committed_target = self._cp_zero_counts(cp_size)
+        committed_draft = self._cp_zero_counts(cp_size)
+        pending_target = self._cp_zero_counts(cp_size)
+        pending_draft = self._cp_zero_counts(cp_size)
+
+        pending = getattr(self, "pending_host_backups", {})
+        pending_node_ids = set(pending.keys())
+        for node in self._cp_walk_radix_nodes():
+            metadata = getattr(node, "cp_hicache", None)
+            if metadata is None or getattr(node, "host_len", 0) <= 0:
+                continue
+            if getattr(node, "id", None) in pending_node_ids:
+                continue
+            counts = self._cp_assert_metadata_counts(
+                metadata, context=f"snapshot_committed node_id={getattr(node, 'id', '?')}"
+            )
+            committed_target = self._cp_add_counts(committed_target, counts)
+            if getattr(metadata, "draft_host_indices", None) is not None:
+                committed_draft = self._cp_add_counts(committed_draft, counts)
+
+        for node_id, pending_backup in pending.items():
+            metadata = pending_backup.metadata
+            counts = self._cp_assert_metadata_counts(
+                metadata, context=f"snapshot_pending node_id={node_id}"
+            )
+            pending_target = self._cp_add_counts(pending_target, counts)
+            if getattr(metadata, "draft_host_indices", None) is not None:
+                pending_draft = self._cp_add_counts(pending_draft, counts)
+
+        return CpHiCacheCapacitySnapshot(
+            target_capacity=target_capacity,
+            draft_capacity=draft_capacity,
+            committed_target=committed_target,
+            committed_draft=committed_draft,
+            pending_target=pending_target,
+            pending_draft=pending_draft,
+        )
+
+    def _cp_required_host_tokens_by_rank(
+        self, device_indices: torch.Tensor
+    ) -> Tuple[int, ...]:
+        layout = getattr(
+            getattr(self, "cache_controller", None), "cp_shared_kv_layout", None
+        )
+        if layout is None:
+            return self._cp_zero_counts()
+        logical_len = int(device_indices.numel())
+        if logical_len == 0:
+            return self._cp_zero_counts(layout.cp_size)
+        if logical_len % self.page_size != 0:
+            raise RuntimeError(
+                "CP HiCache capacity planning requires page-aligned device indices: "
+                f"logical_len={logical_len} page_size={self.page_size}"
+            )
+        logical_pages = torch.div(
+            device_indices[:: self.page_size],
+            self.page_size,
+            rounding_mode="floor",
+        )
+        page_owners = layout.owner_for_logical_pages(logical_pages).to(
+            dtype=torch.int8, device="cpu"
+        )
+        return self._cp_owner_token_counts(
+            page_owners,
+            page_size=self.page_size,
+            cp_size=layout.cp_size,
+        )
+
+    def _cp_host_available_tokens_by_rank(
+        self, snapshot: CpHiCacheCapacitySnapshot, *, draft: bool = False
+    ) -> Tuple[int, ...]:
+        if draft:
+            if snapshot.draft_capacity is None:
+                return tuple(2**62 for _ in snapshot.target_capacity)
+            return self._cp_sub_counts_floor_zero(
+                snapshot.draft_capacity, snapshot.draft_used
+            )
+        return self._cp_sub_counts_floor_zero(
+            snapshot.target_capacity, snapshot.target_used
+        )
+
+    def _cp_host_evict_key(self, node: TreeNode):
+        return (
+            int(getattr(node, "priority", 0)),
+            int(getattr(node, "id", 0)),
+        )
+
+    def _cp_host_leaf_is_plannable_victim(self, node: TreeNode) -> bool:
+        if node == getattr(self, "root_node", None):
+            return False
+        if not getattr(node, "evicted", False):
+            return False
+        parent = getattr(node, "parent", None)
+        if parent is None:
+            return False
+        try:
+            parent_key = self.get_child_key_fn(node.key)
+        except Exception:
+            return False
+        if getattr(parent, "children", {}).get(parent_key) is not node:
+            return False
+        pin_expiry = float(getattr(node, "pin_expiry", 0.0) or 0.0)
+        pin_active = pin_expiry > 0 and time.monotonic() <= pin_expiry
+        if pin_active:
+            return False
+        effective_host_refs = int(getattr(node, "host_ref_counter", 0) or 0)
+        if pin_expiry > 0:
+            # Current eviction clears an expired pin before checking host refs.
+            # Mirror that decision without mutating the tree in observer mode.
+            effective_host_refs = max(0, effective_host_refs - 1)
+        if effective_host_refs > 0:
+            return False
+        if self._node_host_write_pending(node):
+            return False
+        try:
+            return self._node_backuped(node)
+        except RuntimeError:
+            return False
+
+    def _plan_cp_host_evictions(
+        self, required_tokens_by_rank: Tuple[int, ...]
+    ) -> CpHiCacheEvictionPlan:
+        deficits = [max(0, int(v)) for v in required_tokens_by_rank]
+        cp_size = len(deficits)
+        planned_freed = [0 for _ in range(cp_size)]
+        victims: List[TreeNode] = []
+        if all(v <= 0 for v in deficits):
+            return CpHiCacheEvictionPlan((), tuple(planned_freed), tuple(deficits))
+
+        leaves = sorted(
+            list(getattr(self, "evictable_host_leaves", set())),
+            key=self._cp_host_evict_key,
+        )
+        for node in leaves:
+            if not self._cp_host_leaf_is_plannable_victim(node):
+                continue
+            metadata = getattr(node, "cp_hicache", None)
+            if metadata is None:
+                continue
+            counts = self._cp_assert_metadata_counts(
+                metadata, context=f"eviction_plan node_id={getattr(node, 'id', '?')}"
+            )
+            if len(counts) != cp_size:
+                raise RuntimeError(
+                    "CP HiCache eviction plan CP size mismatch: "
+                    f"required={cp_size} metadata={len(counts)}"
+                )
+            victims.append(node)
+            for rank, count in enumerate(counts):
+                planned_freed[rank] += int(count)
+                deficits[rank] = max(0, deficits[rank] - int(count))
+            if all(v <= 0 for v in deficits):
+                break
+
+        return CpHiCacheEvictionPlan(
+            victims=tuple(victims),
+            planned_freed=tuple(planned_freed),
+            remaining_deficit=tuple(deficits),
+        )
+
+    def _cp_capacity_debug_enabled(self) -> bool:
+        return bool(
+            getattr(self, "_cp_hicache_capacity_debug", False)
+            or envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get()
+        )
+
+    def _cp_build_write_admission_plan(
+        self, device_indices: torch.Tensor, *, node_id: int, phase: str
+    ) -> dict:
+        required = self._cp_required_host_tokens_by_rank(device_indices)
+        snapshot = self._cp_host_capacity_snapshot()
+        target_available = self._cp_host_available_tokens_by_rank(snapshot)
+        draft_available = self._cp_host_available_tokens_by_rank(snapshot, draft=True)
+        target_deficit = tuple(
+            max(0, req - avail) for req, avail in zip(required, target_available)
+        )
+        draft_deficit = tuple(
+            max(0, req - avail) for req, avail in zip(required, draft_available)
+        )
+        deficit = tuple(max(a, b) for a, b in zip(target_deficit, draft_deficit))
+        current_compatible_need = tuple(
+            req if missing > 0 else 0 for req, missing in zip(required, deficit)
+        )
+        eviction_plan = self._plan_cp_host_evictions(deficit)
+        return {
+            "node_id": node_id,
+            "phase": phase,
+            "required": required,
+            "target_available": target_available,
+            "draft_available": draft_available,
+            "deficit": deficit,
+            "current_compatible_need": current_compatible_need,
+            "eviction_plan": eviction_plan,
+        }
+
+    def _cp_debug_build_write_admission_plan(
+        self, device_indices: torch.Tensor, *, node_id: int, phase: str
+    ) -> Optional[dict]:
+        if not self._cp_capacity_debug_enabled():
+            return None
+        try:
+            return self._cp_build_write_admission_plan(
+                device_indices, node_id=node_id, phase=phase
+            )
+        except Exception as exc:
+            logger.warning(
+                "[HiCache-capacity-debug] failed to build write admission plan: node_id=%d phase=%s error=%s",
+                node_id,
+                phase,
+                exc,
+            )
+            return None
+
+    def _cp_debug_compare_write_admission(
+        self,
+        plan: Optional[dict],
+        *,
+        result,
+        planned_required_slots: int,
+    ) -> None:
+        if plan is None or not self._cp_capacity_debug_enabled():
+            return
+        stats = getattr(self, "_cp_hicache_capacity_debug_stats", None)
+        if stats is None:
+            stats = self._cp_hicache_capacity_debug_stats = {}
+        tag = "write_admission"
+        entry = stats.setdefault(tag, {"count": 0, "mismatch": 0})
+        entry["count"] += 1
+
+        planned_required_slots = max(plan["current_compatible_need"], default=0)
+        planned_needs_eviction = planned_required_slots > 0
+        reservation_failed = isinstance(result, HiCacheWriteFailure)
+        mismatch = (
+            reservation_failed
+            and not planned_needs_eviction
+            and all(v <= 0 for v in plan["eviction_plan"].remaining_deficit)
+        )
+        if mismatch:
+            entry["mismatch"] += 1
+        should_log = mismatch or entry["count"] == 1 or entry["count"] % 128 == 0
+        if not should_log:
+            return
+
+        eviction_plan = plan["eviction_plan"]
+        log_fn = logger.warning if mismatch else logger.info
+        log_fn(
+            "[HiCache-capacity-debug] write_admission_compare node_id=%d phase=%s "
+            "count=%d mismatch=%d planned_required=%d local_failed=%s "
+            "planned_current_required=%d required=%s target_avail=%s draft_avail=%s "
+            "deficit=%s planned_freed=%s remaining_deficit=%s victims=%s",
+            plan["node_id"],
+            plan["phase"],
+            entry["count"],
+            entry["mismatch"],
+            int(planned_required_slots),
+            isinstance(result, HiCacheWriteFailure),
+            int(planned_required_slots),
+            plan["required"],
+            plan["target_available"],
+            plan["draft_available"],
+            plan["deficit"],
+            eviction_plan.planned_freed,
+            eviction_plan.remaining_deficit,
+            [getattr(node, "id", None) for node in eviction_plan.victims],
+        )
+
     def shutdown(self):
         """Best-effort auto-detach of storage backend on process shutdown.
 
@@ -1021,84 +1495,130 @@ class HiRadixCache(RadixCache):
             self.dec_node_lock_ref(node)
         return node
 
-    def _sync_cp_write_required_host_slots(self, result) -> int:
-        """Return the max host slots any CP rank needs before write backup.
+    def _evict_cp_host_plan(
+        self, plan: dict, *, node_id: int, phase: str
+    ) -> bool:
+        eviction_plan = plan["eviction_plan"]
+        if any(v > 0 for v in eviction_plan.remaining_deficit):
+            logger.warning(
+                "[CP_HICACHE_FALLBACK][cp_host_reservation_plan_insufficient] "
+                "reason=insufficient_evictable_host_slots node_id=%d phase=%s "
+                "required=%s target_avail=%s draft_avail=%s deficit=%s "
+                "planned_freed=%s remaining_deficit=%s victims=%s",
+                node_id,
+                phase,
+                plan["required"],
+                plan["target_available"],
+                plan["draft_available"],
+                plan["deficit"],
+                eviction_plan.planned_freed,
+                eviction_plan.remaining_deficit,
+                [getattr(node, "id", None) for node in eviction_plan.victims],
+            )
+            return False
 
-        CP ranks must make the reserve/evict/retry decision collectively.  A
-        local successful reservation is not enough: another CP rank may be full
-        and enter host eviction, whose implementation contains collective
-        all_reduce calls.  If successful ranks skip that branch they can instead
-        enter unrelated collectives (for example disagg polling), causing a
-        process-group mismatch.
-        """
+        if len(eviction_plan.victims) == 0:
+            return True
 
-        required_host_slots = (
-            int(result.required_host_slots)
-            if isinstance(result, HiCacheWriteFailure)
-            else 0
-        )
-        if (
-            getattr(self, "tp_group", None) is None
-            or getattr(self, "tp_world_size", 1) <= 1
-        ):
-            return required_host_slots
+        local_freed = 0
+        victim_ids = []
+        for victim in eviction_plan.victims:
+            victim_id = getattr(victim, "id", None)
+            if not self._cp_host_leaf_is_plannable_victim(victim):
+                raise RuntimeError(
+                    "CP HiCache deterministic eviction plan diverged before "
+                    f"reservation: node_id={node_id} phase={phase} victim_id={victim_id}"
+                )
 
-        required = torch.tensor(required_host_slots, dtype=torch.int64, device="cpu")
-        torch.distributed.all_reduce(
-            required, op=torch.distributed.ReduceOp.MAX, group=self.tp_group
-        )
-        return int(required.item())
+            victim_ids.append(victim_id)
+            self._record_remove_event(victim)
+            local_freed += self.cache_controller.evict_cp_host(victim.cp_hicache)
+            victim.host_len = 0
+            victim.cp_hicache = None
+            victim.host_value = None
+            self._remove_host_leaf(victim)
 
-    def _release_cp_write_reservation_for_retry(self, result, node_id: int) -> None:
-        if isinstance(result, HiCacheWriteFailure):
-            return
         logger.info(
-            "[HiCache-write] release local CP reservation before collective retry: node_id=%d owned_positions=%d",
+            "[HiCache-evict] deterministic CP host eviction before write: "
+            "node_id=%d phase=%s victims=%s local_freed=%d planned_freed=%s",
             node_id,
-            result.metadata.owned_positions.numel(),
+            phase,
+            victim_ids,
+            local_freed,
+            eviction_plan.planned_freed,
         )
-        self.cache_controller.evict_cp_host(result.metadata)
+        return True
 
-    def _reserve_write_cp_indices_collectively(
+    def _reserve_write_cp_indices_no_collective(
         self, device_indices: torch.Tensor, node_id: int
     ):
+        plan = self._cp_build_write_admission_plan(
+            device_indices, node_id=node_id, phase="initial"
+        )
+        planned_required_slots = max(plan["current_compatible_need"], default=0)
+        if planned_required_slots > 0:
+            if not self._evict_cp_host_plan(plan, node_id=node_id, phase="initial"):
+                return HiCacheWriteFailure(required_host_slots=planned_required_slots)
+
         result = self.cache_controller.reserve_write_cp(
             device_indices=device_indices,
             node_id=node_id,
         )
-        required_host_slots = self._sync_cp_write_required_host_slots(result)
-        if required_host_slots <= 0:
+        debug_plan = plan if self._cp_capacity_debug_enabled() else None
+        self._cp_debug_compare_write_admission(
+            debug_plan,
+            result=result,
+            planned_required_slots=planned_required_slots,
+        )
+        if not isinstance(result, HiCacheWriteFailure):
             return result
 
-        self._release_cp_write_reservation_for_retry(result, node_id)
-        self._evict_host_for_physical_slots(
-            required_host_slots,
-            synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
+        retry_plan = self._cp_build_write_admission_plan(
+            device_indices, node_id=node_id, phase="retry_after_local_failure"
         )
+        retry_required_slots = max(
+            retry_plan["current_compatible_need"], default=int(result.required_host_slots)
+        )
+        if retry_required_slots <= 0:
+            raise RuntimeError(
+                "CP HiCache host reservation failed although deterministic "
+                "capacity planner predicted no eviction need: "
+                f"node_id={node_id} required_host_slots={result.required_host_slots}"
+            )
+        if not self._evict_cp_host_plan(
+            retry_plan, node_id=node_id, phase="retry_after_local_failure"
+        ):
+            return HiCacheWriteFailure(required_host_slots=retry_required_slots)
+
         logger.info(
-            "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
+            "[HiCache-write] write_backup CP retry after deterministic host eviction: "
+            "node_id=%d needed_slots=%d",
             node_id,
-            required_host_slots,
+            retry_required_slots,
         )
         result = self.cache_controller.reserve_write_cp(
             device_indices=device_indices,
             node_id=node_id,
         )
-        required_host_slots = self._sync_cp_write_required_host_slots(result)
-        if required_host_slots <= 0:
+        self._cp_debug_compare_write_admission(
+            retry_plan if self._cp_capacity_debug_enabled() else None,
+            result=result,
+            planned_required_slots=retry_required_slots,
+        )
+        if not isinstance(result, HiCacheWriteFailure):
             return result
 
-        self._release_cp_write_reservation_for_retry(result, node_id)
         logger.info(
-            "[HiCache-write] write_backup CP FAILED after collective retry: node_id=%d len=%d needed_slots=%d",
+            "[HiCache-write] write_backup CP FAILED after deterministic retry: "
+            "node_id=%d len=%d needed_slots=%d",
             node_id,
             len(device_indices) if device_indices is not None else 0,
-            required_host_slots,
+            int(result.required_host_slots),
         )
-        return HiCacheWriteFailure(required_host_slots=required_host_slots)
+        return result
 
-    def _reserve_write_cp_collectively(self, node: TreeNode):
-        return self._reserve_write_cp_indices_collectively(node.value, node.id)
+    def _reserve_write_cp(self, node: TreeNode):
+        return self._reserve_write_cp_indices_no_collective(node.value, node.id)
 
     def _attach_prepared_cp_backup(
         self, node: TreeNode, prepared: PreparedCpHiCacheBackup
@@ -1165,6 +1685,17 @@ class HiRadixCache(RadixCache):
         )
         self.cache_controller.evict_cp_host(prepared.metadata)
 
+    def _warn_cp_hicache_fallback(self, fallback_name: str, reason: str, **kwargs):
+        details = " ".join(f"{key}={value}" for key, value in kwargs.items())
+        if details:
+            details = " " + details
+        logger.warning(
+            "[CP_HICACHE_FALLBACK][%s] reason=%s%s",
+            fallback_name,
+            reason,
+            details,
+        )
+
     def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int:
         """Return the already-present radix prefix length without mutating the tree.
 
@@ -1212,13 +1743,31 @@ class HiRadixCache(RadixCache):
         return total_prefix_len
 
     def prepare_write_backup_for_req(self, req) -> None:
-        if (
-            self.disable
-            or not self._uses_cp_hicache
-            or self.cache_controller.write_policy == "write_back"
-            or getattr(req, "is_chunked", 0) > 0
-            or getattr(req, "cp_hicache_prepared_backup", None) is not None
-        ):
+        if self.disable or not self._uses_cp_hicache:
+            return
+        if self.cache_controller.write_policy == "write_back":
+            self._warn_cp_hicache_fallback(
+                "prepare_write_backup_skipped",
+                "write_back_policy",
+                rid=getattr(req, "rid", "<unknown>"),
+            )
+            return
+        if getattr(req, "is_chunked", 0) > 0:
+            self._warn_cp_hicache_fallback(
+                "prepare_write_backup_skipped",
+                "chunked_req",
+                rid=getattr(req, "rid", "<unknown>"),
+                is_chunked=getattr(req, "is_chunked", 0),
+            )
+            return
+        if getattr(req, "cp_hicache_prepared_backup", None) is not None:
+            prepared = getattr(req, "cp_hicache_prepared_backup", None)
+            self._warn_cp_hicache_fallback(
+                "prepare_write_backup_skipped",
+                "already_prepared",
+                rid=getattr(req, "rid", "<unknown>"),
+                node_id=getattr(prepared, "node_id", None),
+            )
             return
 
         token_ids = req.fill_ids
@@ -1239,17 +1788,26 @@ class HiRadixCache(RadixCache):
             req.req_pool_idx, start : len(keys)
         ].to(dtype=torch.int64, copy=True)
         if len(kv_indices) == 0:
+            self._warn_cp_hicache_fallback(
+                "prepare_write_backup_skipped",
+                "empty_kv_indices",
+                rid=getattr(req, "rid", "<unknown>"),
+                start=start,
+                key_len=len(keys),
+            )
             return
 
         node_id = TreeNode.counter
         TreeNode.counter += 1
-        result = self._reserve_write_cp_indices_collectively(kv_indices, node_id)
+        result = self._reserve_write_cp_indices_no_collective(kv_indices, node_id)
         if isinstance(result, HiCacheWriteFailure):
-            logger.info(
-                "[HiCache-write] prepare CP backup skipped after host reservation failure: node_id=%d rid=%s len=%d",
-                node_id,
-                getattr(req, "rid", "<unknown>"),
-                len(kv_indices),
+            self._warn_cp_hicache_fallback(
+                "prepare_write_backup_reservation_failed",
+                "host_reservation_failed",
+                node_id=node_id,
+                rid=getattr(req, "rid", "<unknown>"),
+                logical_len=len(kv_indices),
+                required_host_slots=result.required_host_slots,
             )
             return
 
@@ -1294,12 +1852,21 @@ class HiRadixCache(RadixCache):
             write_back,
         )
         if self._uses_cp_hicache:
-            result = self._reserve_write_cp_collectively(node)
+            if not write_back:
+                self._warn_cp_hicache_fallback(
+                    "post_forward_catch_up_backup",
+                    "write_backup_called_without_prepared_cp_backup",
+                    node_id=node.id,
+                    logical_len=len(node.value) if node.value is not None else 0,
+                )
+            result = self._reserve_write_cp(node)
             if isinstance(result, HiCacheWriteFailure):
-                logger.info(
-                    "[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d",
-                    node.id,
-                    len(node.value) if node.value is not None else 0,
+                self._warn_cp_hicache_fallback(
+                    "post_forward_catch_up_backup_failed",
+                    "host_reservation_failed",
+                    node_id=node.id,
+                    logical_len=len(node.value) if node.value is not None else 0,
+                    required_host_slots=result.required_host_slots,
                 )
                 return 0
 
@@ -1442,10 +2009,11 @@ class HiRadixCache(RadixCache):
         queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
         if self.tp_world_size > 1:
             # synchronize TP workers to make the same update to radix cache
-            torch.distributed.all_reduce(
+            self._cp_hicache_all_reduce(
                 queue_size,
                 op=torch.distributed.ReduceOp.MIN,
                 group=self.tp_group,
+                tag="writing_check_min",
             )
 
         finish_count = int(queue_size.item())
@@ -1783,8 +2351,11 @@ class HiRadixCache(RadixCache):
             if not synchronize_across_ranks:
                 return bool(local_done)
             done = torch.tensor(local_done, dtype=torch.int, device="cpu")
-            torch.distributed.all_reduce(
-                done, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
+            self._cp_hicache_all_reduce(
+                done,
+                op=torch.distributed.ReduceOp.MIN,
+                group=self.tp_group,
+                tag="host_evict_done_min",
             )
             return bool(done.item())
 
@@ -2133,8 +2704,11 @@ class HiRadixCache(RadixCache):
             dtype=torch.int,
         )
         if self.tp_world_size > 1:
-            torch.distributed.all_reduce(
-                qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
+            self._cp_hicache_all_reduce(
+                qsizes,
+                op=torch.distributed.ReduceOp.MIN,
+                group=self.tp_group,
+                tag="storage_queue_min",
             )
 
         n_revoke, n_backup, n_release = map(int, qsizes.tolist())
@@ -2181,10 +2755,11 @@ class HiRadixCache(RadixCache):
                 [1 - int(can_terminate), int(operation_terminated)],
                 dtype=torch.int,
             )
-            torch.distributed.all_reduce(
+            self._cp_hicache_all_reduce(
                 states,
                 op=torch.distributed.ReduceOp.MAX,
                 group=self.tp_group,
+                tag="prefetch_terminate_max",
             )
             can_terminate = states[0].item() == 0
             operation_terminated = states[1].item() == 1
@@ -2222,10 +2797,11 @@ class HiRadixCache(RadixCache):
             completed_tokens_tensor = torch.tensor(
                 min_completed_tokens, dtype=torch.int
             )
-            torch.distributed.all_reduce(
+            self._cp_hicache_all_reduce(
                 completed_tokens_tensor,
                 op=torch.distributed.ReduceOp.MIN,
                 group=self.tp_group,
+                tag="prefetch_completed_min",
             )
             min_completed_tokens = completed_tokens_tensor.item()
         fetched_token_ids = token_ids[:min_completed_tokens]

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -331,6 +331,12 @@ class SessionAwareCache(BasePrefixCache):
     def flush_write_through_acks(self) -> None:
         return self.inner.flush_write_through_acks()
 
+    def prepare_write_backup_for_req(self, req) -> None:
+        prepare_fn = getattr(self.inner, "prepare_write_backup_for_req", None)
+        if prepare_fn is not None:
+            return prepare_fn(req)
+        return None
+
     def check_hicache_events(self):
         return self.inner.check_hicache_events()