[HiCache][perf] Queue and batch write-through backups off the request-finish path

python/sglang/srt/environ.py

@@ -779,6 +779,8 @@ class Envs:
     SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
     SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
     SGLANG_HICACHE_SKIP_HOST_DUPLICATE_RECLAIM = EnvBool(False)
+    # Defer write-through backups to the end of the batch and back them up merged.
+    SGLANG_ENABLE_HICACHE_BATCHED_BACKUP = EnvBool(False)
     SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
     # File-backend LRU eviction (opt-in; sizes accept SI/IEC suffixes, "0" disables).
     SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE = EnvStr(None)

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

@@ -541,21 +541,25 @@ def write(
         node_id: int = -1,
         extra_pools: Optional[list[PoolTransfer]] = None,
         flush: bool = True,
+        node_ids: Optional[list[int]] = None,
     ) -> Optional[torch.Tensor]:
+        """Queue a D2H backup; flush=False defers the start_writing() submit.
+        node_ids, when given, replaces node_id as the op's ack list."""
         allocation = self.allocate_host_transfers(device_indices, extra_pools)
         if allocation is None:
             return None
         host_indices, pool_transfers = allocation
 
-        self.write_queue.append(
-            CacheOperation(
-                host_indices,
-                device_indices,
-                node_id,
-                priority,
-                pool_transfers=pool_transfers or None,
-            )
+        operation = CacheOperation(
+            host_indices,
+            device_indices,
+            node_id,
+            priority,
+            pool_transfers=pool_transfers or None,
         )
+        if node_ids is not None:
+            operation.node_ids = list(node_ids)
+        self.write_queue.append(operation)
         if flush:
             self.start_writing()
         return host_indices

python/sglang/srt/mem_cache/rust_tree_core/adapter.py

@@ -969,6 +969,9 @@ def is_backuped(self, node_id: NodeId) -> bool:
     def is_root(self, node_id: NodeId) -> bool:
         return self._binding.is_root(node_id)
 
+    def get_parent_node_id(self, node_id: NodeId) -> Optional[NodeId]:
+        return self._binding.parent_node_id(node_id)
+
     def get_last_hash_value(self, node_id: NodeId) -> Optional[str]:
         return self._binding.get_last_hash_value(node_id)
 

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -567,6 +567,11 @@ def is_root(self, node_id: NodeId) -> bool:
         """Whether the node is the tree root."""
         return self.node_by_id(node_id) is self.root_node
 
+    def get_parent_node_id(self, node_id: NodeId) -> Optional[NodeId]:
+        """The parent node id, or None for the root."""
+        parent = self.node_by_id(node_id).parent
+        return None if parent is None else parent.id
+
     supports_rotation_base = True
 
     def rotation_base_of(self, node_id: NodeId) -> Optional[int]:

python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py

@@ -205,6 +205,11 @@ def is_root(self, node_id: NodeId) -> bool:
         """Whether the node is the tree root."""
         ...
 
+    @abstractmethod
+    def get_parent_node_id(self, node_id: NodeId) -> Optional[NodeId]:
+        """The parent node id, or None for the root."""
+        ...
+
     # Logical-page KV sharding: whether this core stamps and honors
     # UnifiedTreeNode.rotation_base. A core that does not cannot serve a
     # sharded allocator (it would never decline a cross-base graft), and

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -151,6 +151,16 @@ class _OngoingWriteThrough(NamedTuple):
     publish_node_ids: list[NodeId]
 
 
+class _NodeBackupSpec(NamedTuple):
+    """One node's pending D->H backup: its transfers and the lock it holds."""
+
+    node_id: NodeId
+    lock_params: Optional[DecLockRefParams]
+    device_value: torch.Tensor
+    comp_xfers: dict[ComponentType, list[PoolTransfer]]
+    publish_node_ids: list[NodeId]
+
+
 class _OngoingLoadBack(NamedTuple):
     """Tracks an in-flight H→D load-back operation."""
 
@@ -276,6 +286,7 @@ def __init__(
         )
         self.pp_rank = params.pp_rank
         self.pp_size = params.pp_size
+        self._batched_backup = envs.SGLANG_ENABLE_HICACHE_BATCHED_BACKUP.get()
         self.work_list: list[torch.distributed.Work] = []
 
         # HiCache D↔H defaults (overridden by init_hicache)
@@ -395,6 +406,8 @@ def _reset_full(self) -> None:
         # Reset Controller.
         self.session.slots.clear()
         self.ongoing_write_through: dict[int, _OngoingWriteThrough] = {}
+        # Locked write-through nodes awaiting flush_pending_backups, in queue order.
+        self.queued_backups: dict[NodeId, Optional[DecLockRefParams]] = {}
         self.ongoing_load_back: dict[int, _OngoingLoadBack] = {}
         self.enable_storage = False
         self.prefetch_loaded_tokens_by_reqid: dict[CacheRequestHandle, int] = {}
@@ -1310,6 +1323,14 @@ def _apply_cache_action(self, action: CacheAction | ComponentAction) -> None:
         elif isinstance(action, BackupKV):
             if self.linker is not None:
                 self.linker.offload_nodes(action.node_ids)
+            elif (
+                self._batched_backup
+                and self.buffer_pipeline is None
+                and not self.is_write_back
+            ):
+                # Write-through backups are queued and merged at flush; write-back
+                # backs up only its target, so it stays at insert time.
+                self._queue_write_through_backup(action)
             else:
                 self._execute_and_commit_kv_backup(action)
         else:
@@ -1615,6 +1636,86 @@ def _execute_and_commit_kv_backup(
             written = len(host_indices)
         return written
 
+    def _build_node_backup_spec(
+        self, node_id: NodeId, lock_params: Optional[DecLockRefParams]
+    ) -> Optional[_NodeBackupSpec]:
+        """Read what a node still needs backed up; None (lock released) when nothing."""
+        device_value, comp_xfers = self.tree_core.build_backup_spec(node_id)
+        # Overlapping chain actions may revisit nodes with Full KV already
+        # backed up. Skip only when no transfer remains.
+        if device_value.numel() == 0 and not comp_xfers:
+            self._release_queued_backup(node_id, lock_params)
+            return None
+        return _NodeBackupSpec(
+            node_id,
+            lock_params,
+            device_value,
+            comp_xfers,
+            self._backup_publish_node_ids(node_id, comp_xfers),
+        )
+
+    def _backup_specs(
+        self, specs: list[_NodeBackupSpec], *, write_back: bool = False
+    ) -> Optional[int]:
+        """Full KV tokens written, or None if host alloc fails (locks stay held)."""
+        device_value, comp_xfers, groups = self._merge_backup_transfers(specs)
+        host_indices = self._execute_kv_backup(
+            -1,
+            device_value,
+            comp_xfers,
+            self._build_backup_sidecar(device_value, comp_xfers),
+            node_ids=[spec.node_id for spec in specs],
+        )
+        if host_indices is None:
+            return None
+        for merged, transfers in groups:
+            sizes = [len(t.device_indices) for t in transfers]
+            for transfer, host_slice in zip(
+                transfers, merged.host_indices.split(sizes)
+            ):
+                transfer.host_indices = host_slice
+        sizes = [len(spec.device_value) for spec in specs]
+        for spec, host_slice in zip(specs, host_indices.split(sizes)):
+            self.tree_core.commit_backup(spec.node_id, host_slice, spec.comp_xfers)
+            lock_params = spec.lock_params
+            if not write_back and lock_params is None:
+                lock_params = self.inc_lock_ref(spec.node_id).to_dec_params()
+            self._track_write_through_node(
+                spec.node_id, lock_params, publish_node_ids=spec.publish_node_ids
+            )
+        return len(host_indices)
+
+    @staticmethod
+    def _merge_backup_transfers(
+        specs: list[_NodeBackupSpec],
+    ) -> tuple[
+        torch.Tensor,
+        dict[ComponentType, list[PoolTransfer]],
+        list[tuple[PoolTransfer, list[PoolTransfer]]],
+    ]:
+        """Concatenate the nodes' transfers per (component, pool); each merged
+        transfer comes back paired with the per-node transfers it covers."""
+        if len(specs) == 1:
+            return specs[0].device_value, specs[0].comp_xfers, []
+        parts: dict[tuple[ComponentType, PoolName], list[PoolTransfer]] = {}
+        for spec in specs:
+            for component_type, transfers in spec.comp_xfers.items():
+                for transfer in transfers:
+                    parts.setdefault((component_type, transfer.name), []).append(
+                        transfer
+                    )
+        comp_xfers: dict[ComponentType, list[PoolTransfer]] = {}
+        groups = []
+        for (component_type, name), transfers in parts.items():
+            merged = PoolTransfer(
+                name=name,
+                device_indices=torch.cat([t.device_indices for t in transfers]),
+            )
+            comp_xfers.setdefault(component_type, []).append(merged)
+            groups.append((merged, transfers))
+        device_value = torch.cat([spec.device_value for spec in specs])
+        return device_value, comp_xfers, groups
+
     @staticmethod
     def _backup_publish_node_ids(
         node_id: NodeId, comp_xfers: dict[ComponentType, list[PoolTransfer]]
@@ -1635,8 +1736,10 @@ def _build_backup_sidecar(self, device_value, comp_xfers):
             CacheTransferPhase.BACKUP_HOST, kv_xfer, comp_xfers
         )
 
-    def _execute_kv_backup(self, node_id, device_value, comp_xfers, sidecar_xfers):
-        """Execute Backup action."""
+    def _execute_kv_backup(
+        self, node_id, device_value, comp_xfers, sidecar_xfers, *, node_ids=None
+    ):
+        """Allocate host slots and queue a single-node or batched D->H backup."""
         kv_tokens = len(device_value)
         anchor_entry = self.cache_controller.mem_pool_host.anchor_entry
         if (
@@ -1652,7 +1755,11 @@ def _execute_kv_backup(self, node_id, device_value, comp_xfers, sidecar_xfers):
         aux_xfers.extend(sidecar_xfers)
         # Defer submission so the next flush can merge pending node backups.
         return self.cache_controller.write(
-            device_value, node_id=node_id, extra_pools=aux_xfers or None, flush=False
+            device_value,
+            node_id=node_id,
+            extra_pools=aux_xfers or None,
+            flush=False,
+            node_ids=node_ids,
         )
 
     def _track_write_through_node(
@@ -3320,6 +3427,8 @@ def writing_check(
 
         if write_back:
             # Blocking: submit what is still queued, then wait for every ack.
+            if self.queued_backups:
+                self._drain_queued_backups()
             cc.start_writing()
             while self.ongoing_write_through:
                 for ack in cc.ack_write_queue:
@@ -3527,11 +3636,91 @@ def check_hicache_events(self) -> None:
             self.storage_metrics_collector.log_storage_metrics(storage_metrics)
 
     def flush_pending_backups(self) -> None:
-        """Submit pending D2H backups as a merged operation."""
+        """Back queued write-through nodes up and submit every pending D2H backup."""
         if self.linker is not None or self.cache_controller is None:
             return
+        if self.queued_backups:
+            self._drain_queued_backups()
         self.cache_controller.start_writing()
 
+    def _queue_write_through_backup(self, action: BackupKV) -> None:
+        for node_id in action.node_ids:
+            if node_id in self.queued_backups or node_id in self.ongoing_write_through:
+                continue
+            self.queued_backups[node_id] = self.inc_lock_ref(node_id).to_dec_params()
+
+    def _drain_queued_backups(self) -> None:
+        """Back queued nodes up ancestors first, merged as far as SWA windows allow;
+        like a failed BackupKV chain, a node whose parent stays unbacked is dropped."""
+        queued, self.queued_backups = self.queued_backups, {}
+        batch: list[_NodeBackupSpec] = []
+        batch_ids: set[NodeId] = set()
+        covered: set[NodeId] = set()
+        for node_id, lock_params in self._ancestors_first(queued).items():
+            spec = self._build_node_backup_spec(node_id, lock_params)
+            if spec is not None and not covered.isdisjoint(spec.publish_node_ids):
+                # A node already in the batch sits in this one's SWA window; commit
+                # the batch so the rebuilt spec sees it backed up, as node by node would.
+                self._backup_queued_batch(batch)
+                batch, batch_ids, covered = [], set(), set()
+                spec = self._build_node_backup_spec(node_id, lock_params)
+            if spec is None:
+                continue
+            if not self._parent_backed_up(node_id, batch_ids):
+                self._release_queued_backup(node_id, lock_params)
+                continue
+            batch.append(spec)
+            batch_ids.add(node_id)
+            covered.update(spec.publish_node_ids)
+        self._backup_queued_batch(batch)
+
+    def _ancestors_first(
+        self, queued: dict[NodeId, Optional[DecLockRefParams]]
+    ) -> dict[NodeId, Optional[DecLockRefParams]]:
+        """Queued nodes in an order that puts every unbacked ancestor first; a split
+        can queue a parent behind its child or leave it unqueued (no lock held)."""
+        ordered: dict[NodeId, Optional[DecLockRefParams]] = {}
+        for node_id, lock_params in queued.items():
+            chain = []
+            parent = self.tree_core.get_parent_node_id(node_id)
+            while parent not in ordered and self._is_unbacked_node(parent):
+                chain.append(parent)
+                parent = self.tree_core.get_parent_node_id(parent)
+            for ancestor in reversed(chain):
+                ordered[ancestor] = queued.get(ancestor)
+            ordered.setdefault(node_id, lock_params)
+        return ordered
+
+    def _is_unbacked_node(self, node_id: Optional[NodeId]) -> bool:
+        """A non-root node with no host backup yet."""
+        return (
+            node_id is not None
+            and not self.tree_core.is_root(node_id)
+            and not self.tree_core.is_backuped(node_id)
+        )
+
+    def _parent_backed_up(self, node_id: NodeId, batch_ids: set[NodeId]) -> bool:
+        parent = self.tree_core.get_parent_node_id(node_id)
+        return parent in batch_ids or not self._is_unbacked_node(parent)
+
+    def _release_queued_backup(
+        self, node_id: NodeId, lock_params: Optional[DecLockRefParams]
+    ) -> None:
+        if lock_params is not None:
+            self.dec_lock_ref(node_id, lock_params)
+
+    def _backup_queued_batch(self, batch: list[_NodeBackupSpec]) -> None:
+        if not batch or self._backup_specs(batch) is not None:
+            return
+        # The merged host allocation failed: retry node by node, ancestors first.
+        for spec in batch:
+            if not self._parent_backed_up(spec.node_id, set()):
+                self._release_queued_backup(spec.node_id, spec.lock_params)
+                continue
+            retry = self._build_node_backup_spec(spec.node_id, spec.lock_params)
+            if retry is not None and self._backup_specs([retry]) is None:
+                self._release_queued_backup(spec.node_id, spec.lock_params)
+
     def ready_to_load_host_cache(self) -> int:
         """Notify the cache controller to start the KV cache loading."""
         if self.linker is not None:

rust/sglang-radix-tree/src/python_bindings.rs

@@ -1816,6 +1816,12 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
             .map_err(node_access_error)
     }
 
+    /// The parent node's external handle, or None for the root.
+    fn parent_node_id(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Option<NodeId>> {
+        py.detach(|| self.core().parent_node_id(node_id))
+            .map_err(node_access_error)
+    }
+
     /// The node's last page hash, or None when it was never hashed.
     fn get_last_hash_value(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Option<String>> {
         py.detach(|| self.core().get_last_hash_value(node_id))
@@ -2289,15 +2295,6 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
         py.detach(|| self.core().inspect_contains_node(node_id))
     }
 
-    fn inspect_get_parent_node_id(
-        &self,
-        py: Python<'_>,
-        node_id: NodeId,
-    ) -> PyResult<Option<NodeId>> {
-        py.detach(|| self.core().inspect_get_parent_node_id(node_id))
-            .map_err(node_access_error)
-    }
-
     fn inspect_get_child_node_ids(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Vec<NodeId>> {
         py.detach(|| self.core().inspect_get_child_node_ids(node_id))
             .map_err(node_access_error)
@@ -3182,6 +3179,11 @@ macro_rules! tree_core_binding {
                 catch_native_panic(|| self.inner.is_root(py, node_id))
             }
 
+            /// The parent node's external handle, or None for the root.
+            fn parent_node_id(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Option<NodeId>> {
+                catch_native_panic(|| self.inner.parent_node_id(py, node_id))
+            }
+
             /// The node's last page hash, or None when it was never hashed.
             fn get_last_hash_value(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Option<String>> {
                 catch_native_panic(|| self.inner.get_last_hash_value(py, node_id))
@@ -3532,15 +3534,6 @@ macro_rules! tree_core_binding {
                 catch_native_panic(|| Ok(self.inner.inspect_contains_node(py, node_id)))
             }
 
-            #[cfg(feature = "inspection")]
-            fn inspect_get_parent_node_id(
-                &self,
-                py: Python<'_>,
-                node_id: NodeId,
-            ) -> PyResult<Option<NodeId>> {
-                catch_native_panic(|| self.inner.inspect_get_parent_node_id(py, node_id))
-            }
-
             #[cfg(feature = "inspection")]
             fn inspect_get_child_node_ids(&self, py: Python<'_>, node_id: NodeId) -> PyResult<Vec<NodeId>> {
                 catch_native_panic(|| self.inner.inspect_get_child_node_ids(py, node_id))