Batch CP HiCache host admission before write reservation

CP shared-KV prefill batching can present multiple requests to HiCache write-through in one forward. Keeping host free-room admission per request causes repeated small L2 evictions and defeats the free-room policy, while request metadata and radix attachment still need per-request ownership.\n\nThis change adds a batch prepare path that aggregates required host tokens by CP owner lane, performs one host admission/eviction step for the batch, and then reserves/submits each request independently. SessionAwareCache and the scheduler prefer the batch API when available.\n\nConstraint: Radix nodes, host slots, draft slots, rollback, and ack semantics remain per request.\nRejected: Merge reservations or radix nodes across requests | would complicate rollback and writing_check ack ownership.\nRejected: Add collective synchronization for host admission | local owner-lane logic already provides deterministic capacity accounting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce per-request host free-room eviction in the scheduler path without profiling host-full workloads.\nTested: local py_compile for touched Python files\nTested: remote g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 117 passed, 5 warnings\nNot-tested: full ETE bs>1 HiCache run under production traffic

python/sglang/srt/managers/scheduler.py

@@ -2654,6 +2654,18 @@ class Scheduler(
         ):
             return
 
+        prepare_batch_fn = getattr(
+            self.tree_cache, "prepare_write_backups_for_reqs", None
+        )
+        if prepare_batch_fn is None:
+            inner_cache = getattr(self.tree_cache, "inner", None)
+            prepare_batch_fn = getattr(
+                inner_cache, "prepare_write_backups_for_reqs", None
+            )
+        if prepare_batch_fn is not None:
+            prepare_batch_fn(batch.reqs)
+            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)

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -313,6 +313,13 @@ class PreparedCpHiCacheBackup:
     attached: bool = False
 
 
+@dataclass(frozen=True)
+class CpWriteBackupCandidate:
+    req: object
+    kv_indices: torch.Tensor
+    node_id: int
+
+
 @dataclass(frozen=True)
 class CpHiCacheCapacitySnapshot:
     target_capacity: Tuple[int, ...]
@@ -1407,10 +1414,10 @@ class HiRadixCache(RadixCache):
         )
         return EvictResult(num_tokens_evicted=num_evicted)
 
-    def _cp_build_write_admission(
-        self, device_indices: torch.Tensor, *, node_id: int, phase: str
+    def _cp_build_write_admission_from_required(
+        self, required: Tuple[int, ...], *, node_id: int, phase: str
     ) -> CpWriteAdmission:
-        required = self._cp_required_host_tokens_by_rank(device_indices)
+        required = tuple(int(v) for v in required)
         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)
@@ -1418,12 +1425,13 @@ class HiRadixCache(RadixCache):
         trigger_ratio = float(
             getattr(self, "hicache_host_free_room_trigger_ratio", 0.0) or 0.0
         )
+        page_size = int(getattr(self, "page_size", 1) or 1)
         target_deficit = tuple(
             _free_room_deficit(
                 required=req,
                 available=avail,
                 capacity=capacity,
-                page_size=self.page_size,
+                page_size=page_size,
                 target_ratio=target_ratio,
                 trigger_ratio=trigger_ratio,
             )
@@ -1438,7 +1446,7 @@ class HiRadixCache(RadixCache):
                     required=req,
                     available=avail,
                     capacity=capacity,
-                    page_size=self.page_size,
+                    page_size=page_size,
                     target_ratio=target_ratio,
                     trigger_ratio=trigger_ratio,
                 )
@@ -1463,6 +1471,16 @@ class HiRadixCache(RadixCache):
             eviction_plan=eviction_plan,
         )
 
+    def _cp_build_write_admission(
+        self, device_indices: torch.Tensor, *, node_id: int, phase: str
+    ) -> CpWriteAdmission:
+        required = self._cp_required_host_tokens_by_rank(device_indices)
+        return self._cp_build_write_admission_from_required(
+            required,
+            node_id=node_id,
+            phase=phase,
+        )
+
     def shutdown(self):
         """Best-effort auto-detach of storage backend on process shutdown.
 
@@ -2038,18 +2056,25 @@ class HiRadixCache(RadixCache):
         return True
 
     def _reserve_write_cp_indices_no_collective(
-        self, device_indices: torch.Tensor, node_id: int
+        self,
+        device_indices: torch.Tensor,
+        node_id: int,
+        *,
+        admission_checked: bool = False,
     ):
-        admission = self._cp_build_write_admission(
-            device_indices, node_id=node_id, phase="initial"
-        )
-        if any(v > 0 for v in admission.deficit_by_owner):
-            if not self._evict_cp_host_for_write_admission(
-                admission, node_id=node_id, phase="initial"
-            ):
-                return HiCacheWriteFailure(
-                    required_host_slots=max(admission.eviction_plan.remaining_deficit)
-                )
+        if not admission_checked:
+            admission = self._cp_build_write_admission(
+                device_indices, node_id=node_id, phase="initial"
+            )
+            if any(v > 0 for v in admission.deficit_by_owner):
+                if not self._evict_cp_host_for_write_admission(
+                    admission, node_id=node_id, phase="initial"
+                ):
+                    return HiCacheWriteFailure(
+                        required_host_slots=max(
+                            admission.eviction_plan.remaining_deficit
+                        )
+                    )
 
         result = self.cache_controller.reserve_write_cp(
             device_indices=device_indices,
@@ -2416,16 +2441,9 @@ class HiRadixCache(RadixCache):
         if hasattr(self, "evictable_host_leaves"):
             self._update_host_leaf_status(parent)
 
-    def prepare_write_backup_for_req(self, req) -> 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
+    def _build_prepare_write_backup_candidate_for_req(
+        self, req
+    ) -> Optional[CpWriteBackupCandidate]:
         if getattr(req, "cp_hicache_prepared_backup", None) is not None:
             prepared = getattr(req, "cp_hicache_prepared_backup", None)
             self._warn_cp_hicache_fallback(
@@ -2485,7 +2503,32 @@ class HiRadixCache(RadixCache):
 
         node_id = TreeNode.counter
         TreeNode.counter += 1
-        result = self._reserve_write_cp_indices_no_collective(kv_indices, node_id)
+        return CpWriteBackupCandidate(
+            req=req,
+            kv_indices=kv_indices,
+            node_id=node_id,
+        )
+
+    def _submit_prepare_write_backup_candidate(
+        self,
+        candidate: CpWriteBackupCandidate,
+        *,
+        admission_checked: bool = False,
+    ) -> None:
+        req = candidate.req
+        kv_indices = candidate.kv_indices
+        node_id = candidate.node_id
+        if admission_checked:
+            result = self._reserve_write_cp_indices_no_collective(
+                kv_indices,
+                node_id,
+                admission_checked=True,
+            )
+        else:
+            result = self._reserve_write_cp_indices_no_collective(
+                kv_indices,
+                node_id,
+            )
         if isinstance(result, HiCacheWriteFailure):
             self._warn_cp_hicache_fallback(
                 "prepare_write_backup_reservation_failed",
@@ -2519,6 +2562,81 @@ class HiRadixCache(RadixCache):
             result.metadata.owned_positions.numel(),
         )
 
+    def prepare_write_backup_for_req(self, req) -> 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
+
+        candidate = self._build_prepare_write_backup_candidate_for_req(req)
+        if candidate is None:
+            return
+        self._submit_prepare_write_backup_candidate(candidate)
+
+    def prepare_write_backups_for_reqs(self, reqs) -> None:
+        if self.disable or not self._uses_cp_hicache:
+            return
+        reqs = list(reqs)
+        if len(reqs) == 0:
+            return
+        if self.cache_controller.write_policy == "write_back":
+            self._warn_cp_hicache_fallback(
+                "prepare_write_backups_skipped",
+                "write_back_policy",
+                batch_size=len(reqs),
+            )
+            return
+
+        candidates: List[CpWriteBackupCandidate] = []
+        for req in reqs:
+            candidate = self._build_prepare_write_backup_candidate_for_req(req)
+            if candidate is not None:
+                candidates.append(candidate)
+        if len(candidates) == 0:
+            return
+
+        admission_checked = False
+        if len(candidates) > 1:
+            required = self._cp_zero_counts()
+            for candidate in candidates:
+                required = self._cp_add_counts(
+                    required,
+                    self._cp_required_host_tokens_by_rank(candidate.kv_indices),
+                )
+            admission = self._cp_build_write_admission_from_required(
+                required,
+                node_id=candidates[0].node_id,
+                phase="batch_prepare",
+            )
+            if any(v > 0 for v in admission.deficit_by_owner):
+                if self._evict_cp_host_for_write_admission(
+                    admission,
+                    node_id=candidates[0].node_id,
+                    phase="batch_prepare",
+                ):
+                    admission_checked = True
+                else:
+                    self._warn_cp_hicache_fallback(
+                        "prepare_write_backups_batch_admission_failed",
+                        "host_reservation_failed",
+                        batch_size=len(candidates),
+                        required_by_owner=required,
+                        remaining_deficit=admission.eviction_plan.remaining_deficit,
+                    )
+            else:
+                admission_checked = True
+
+        for candidate in candidates:
+            self._submit_prepare_write_backup_candidate(
+                candidate,
+                admission_checked=admission_checked,
+            )
+
     def _node_host_len(self, node: TreeNode) -> int:
         if self._uses_cp_hicache:
             return node.host_len

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -337,6 +337,14 @@ class SessionAwareCache(BasePrefixCache):
             return prepare_fn(req)
         return None
 
+    def prepare_write_backups_for_reqs(self, reqs) -> None:
+        prepare_fn = getattr(self.inner, "prepare_write_backups_for_reqs", None)
+        if prepare_fn is not None:
+            return prepare_fn(reqs)
+        for req in reqs:
+            self.prepare_write_backup_for_req(req)
+        return None
+
     def check_hicache_events(self):
         return self.inner.check_hicache_events()