[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()
 

test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py

@@ -1174,6 +1174,89 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
         controller.move_hybrid_indices.assert_called_once()
         self.assertEqual([indices.device.type for indices in captured], ["cpu", "cpu"])
 
+    def _chain_write_controller(self, captured):
+        """A hybrid controller over a real HostPoolGroup whose KV pool records
+        every backup it receives."""
+
+        class ChainHostPool:
+            layout = "page_first"
+            page_size = 4
+            device = "cpu"
+            size = 64
+            logical_size = 64
+            size_per_token = 2
+            can_use_write_back_jit = True
+
+            def __init__(self):
+                self.next_free = 0
+
+            def alloc(self, need_size):
+                start = self.next_free
+                self.next_free += need_size
+                return _indices(start, start + need_size)
+
+            def backup_from_device_all_layer(
+                self, device_pool, host_indices, device_indices, io_backend
+            ):
+                captured.append((host_indices, device_indices))
+
+        controller = HybridCacheController.__new__(HybridCacheController)
+        controller.write_queue = []
+        controller.io_backend = "kernel"
+        controller.mem_pool_host = HostPoolGroup(
+            [
+                PoolEntry(
+                    name=PoolName.KV,
+                    host_pool=ChainHostPool(),
+                    device_pool=None,
+                    layer_mapper=lambda layer_id: layer_id,
+                    is_primary_index_anchor=True,
+                )
+            ]
+        )
+        controller.mem_pool_device = None
+        controller.ack_write_queue = []
+        with mock.patch.object(transfer_module, "device_module", _FakeDeviceModule):
+            controller.l2_transfer_engine = L2TransferEngine("kernel")
+        return controller
+
+    def test_hybrid_write_without_flush_merges_chain_into_one_submit(self):
+        captured = []
+        controller = self._chain_write_controller(captured)
+
+        with mock.patch.object(transfer_module, "device_module", _FakeDeviceModule):
+            first = controller.write(_indices(4, 8), node_id=1, flush=False)
+            second = controller.write(_indices(12, 16), node_id=2, flush=False)
+            self.assertEqual(len(controller.write_queue), 2)
+            self.assertEqual(captured, [])
+            self.assertEqual(controller.ack_write_queue, [])
+
+            controller.start_writing()
+            # Flushing a drained queue must not submit another copy or ack.
+            controller.start_writing()
+
+        self.assertEqual(controller.write_queue, [])
+        self.assertEqual(len(captured), 1)
+        host_indices, device_indices = captured[0]
+        self.assertEqual(host_indices.tolist(), torch.cat([first, second]).tolist())
+        self.assertEqual(
+            device_indices.tolist(), list(range(4, 8)) + list(range(12, 16))
+        )
+        self.assertEqual(len(controller.ack_write_queue), 1)
+        self.assertEqual(controller.ack_write_queue[0].node_ids, [1, 2])
+        self.assertEqual(controller.ack_write_queue[0].num_tokens, 8)
+
+    def test_hybrid_write_flushes_by_default(self):
+        captured = []
+        controller = self._chain_write_controller(captured)
+
+        with mock.patch.object(transfer_module, "device_module", _FakeDeviceModule):
+            controller.write(_indices(4, 8), node_id=1)
+
+        self.assertEqual(controller.write_queue, [])
+        self.assertEqual(len(captured), 1)
+        self.assertEqual(controller.ack_write_queue[0].node_ids, [1])
+
     def test_write_back_jit_cache_controller_keeps_host_indices_on_cpu(self):
         captured = {}
 

test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py

@@ -68,6 +68,7 @@ class TestUnifiedPPSyncBatching(unittest.TestCase):
         cache.writing_check = MagicMock()
         cache.loading_check = MagicMock()
         cache.cache_controller = SimpleNamespace(
+            start_writing=MagicMock(),
             ack_write_queue=[
                 SimpleNamespace(
                     finish_event=SimpleNamespace(query=MagicMock(return_value=ready))