[HiCache][Perf] fix: batch HiCache D2H submits per step for hybrid pools (#39050)

python/sglang/srt/managers/scheduler.py

@@ -4743,6 +4743,9 @@ class Scheduler(
         elif batch.forward_mode.is_idle():
             self.batch_result_processor.process_batch_result_idle(batch, result)
 
+        # Submit this batch's queued host backups before the next scheduler step.
+        self.tree_cache.flush_pending_backups()
+
         self._record_step_counters(batch, result)
 
         self.metrics_reporter.log_batch_result_stats(batch, result)

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -550,6 +550,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
         """
         raise NotImplementedError()
 
+    def flush_pending_backups(self) -> None:
+        """
+        Submit queued host backups.
+        Caches without deferred backups have nothing to flush.
+        """
+        pass
+
     def take_events(self):
         return [] if self.kv_events is None else self.kv_events.take()
 

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

@@ -327,7 +327,10 @@ class HybridCacheController(BaseHiCacheController):
         priority: Optional[int] = None,
         node_id: int = -1,
         extra_pools: Optional[list[PoolTransfer]] = None,
+        flush: bool = True,
     ) -> Optional[torch.Tensor]:
+        """Queue a D2H backup; flush=False leaves it queued so the caller can
+        merge several nodes into one start_writing() submit."""
         host_indices = self.mem_pool_host.alloc(len(device_indices))
         if host_indices is None:
             return None
@@ -349,7 +352,8 @@ class HybridCacheController(BaseHiCacheController):
                 pool_transfers=pool_transfers or None,
             )
         )
-        self.start_writing()
+        if flush:
+            self.start_writing()
         return host_indices
 
     def _move_op_indices(

python/sglang/srt/mem_cache/pool_host/mamba.py

@@ -448,12 +448,6 @@ class MambaPoolHost(HostKVCache):
             return
         if io_backend == "kernel":
             item_size = MambaPoolHost._item_size_per_index(src_layers[0])
-            # Mamba JIT kernel expects all index tensors on CUDA.
-            # When can_use_write_back_jit is True on the HostPoolGroup,
-            # start_writing() keeps host_indices on CPU (for MLA staged kernel).
-            # Move dst_indices to CUDA here to satisfy the kernel's requirement.
-            if dst_indices.device.type != "cuda":
-                dst_indices = dst_indices.to(src_indices.device, non_blocking=True)
             transfer_kv_mamba_lf_pf(
                 src_ptrs=src_ptrs,
                 dst=dst,
@@ -549,6 +543,11 @@ class MambaPoolHost(HostKVCache):
         self, device_pool, host_indices, device_indices, io_backend="kernel"
     ):
         if self.layout in ["page_first", "page_first_direct"]:
+            if io_backend == "kernel" and host_indices.device != device_indices.device:
+                # The mamba JIT kernel wants both index tensors on the device;
+                # the staged MHA/MLA write path hands us CPU host indices.
+                # Convert once here rather than per conv/temporal tensor.
+                host_indices = host_indices.to(device_indices.device, non_blocking=True)
             # no ssm state on conv-only models: a 0-size batched memcpy errors
             if self.temporal_state_elem_size > 0:
                 self._copy_tensor_all_layers_lf_pf(

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1594,8 +1594,9 @@ class UnifiedRadixCache(BasePrefixCache):
                 return None
         aux_xfers = [x for xfers in comp_xfers.values() for x in 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
+            device_value, node_id=node_id, extra_pools=aux_xfers or None, flush=False
         )
 
     def _track_write_through_node(
@@ -3136,7 +3137,8 @@ class UnifiedRadixCache(BasePrefixCache):
             return
 
         if write_back:
-            # Blocking: wait for all pending write-backs
+            # Blocking: submit what is still queued, then wait for every ack.
+            cc.start_writing()
             while self.ongoing_write_through:
                 for ack in cc.ack_write_queue:
                     ack.finish_event.synchronize()
@@ -3312,6 +3314,9 @@ class UnifiedRadixCache(BasePrefixCache):
 
         # Reap the previous round's PP-sync sends before issuing new ones.
         self._drain_async_work()
+        # Backups queued outside process_batch_result: the chunked-prefill stash
+        # in get_next_batch_to_run, abort_request, and the PD prefill release.
+        self.flush_pending_backups()
 
         (
             write_finish_count,
@@ -3346,6 +3351,12 @@ class UnifiedRadixCache(BasePrefixCache):
                 storage_metrics.prefetch_stats = self.prefetch_outcome_stats_snapshot()
             self.storage_metrics_collector.log_storage_metrics(storage_metrics)
 
+    def flush_pending_backups(self) -> None:
+        """Submit pending D2H backups as a merged operation."""
+        if self.linker is not None or self.cache_controller is None:
+            return
+        self.cache_controller.start_writing()
+
     def ready_to_load_host_cache(self) -> int:
         """Notify the cache controller to start the KV cache loading."""
         if self.linker is not None:

python/sglang/srt/session/streaming_session.py

@@ -601,6 +601,9 @@ class StreamingSession(BasePrefixCache):
     def check_hicache_events(self):
         return self.inner.check_hicache_events()
 
+    def flush_pending_backups(self) -> None:
+        self.inner.flush_pending_backups()
+
     def take_events(self):
         return self.inner.take_events()