diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7c7410828..b574b30b2 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 661bdf5c9..5d260f241 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 539b952e7..e6409fe2d 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 8b8eda90e..51f1a400f 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index e2d4b5e8b..dd3dc60ae 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 96ec95f76..82088cd8a 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index d0646adb8..119d2766e 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/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 = {} diff --git a/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py b/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py index 8a9da6c56..c3600af32 100644 --- a/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py +++ b/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))