From 568aa5fcdbc7afa88eddb28e1af22dc60a5f6014 Mon Sep 17 00:00:00 2001 From: Yongji Wu <30348494+libertyeagle@users.noreply.github.com> Date: Sat, 13 Jun 2026 00:17:31 -0700 Subject: [PATCH] Fix missing draft KV pool transfers in HybridCacheController (#27953) --- .../hybrid_cache/hybrid_cache_controller.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) 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 868ae7aba..e69160368 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 @@ -410,6 +410,13 @@ class HybridCacheController(BaseHiCacheController): self.io_backend, pool_transfers=resolved_pool_transfers, ) + if self.has_draft and host_indices.numel() > 0: + self.mem_pool_host_draft.backup_from_device_all_layer( + self.mem_pool_device_draft, + host_indices, + device_indices, + self.io_backend, + ) finish_event.record() self._record_transfer_indices_on_stream( self.write_stream, @@ -484,6 +491,18 @@ class HybridCacheController(BaseHiCacheController): self.io_backend, pool_transfers=resolved_pool_transfers, ) + if ( + self.has_draft + and host_indices.numel() > 0 + and i < self.mem_pool_host_draft.layer_num + ): + self.mem_pool_host_draft.load_to_device_per_layer( + self.mem_pool_device_draft, + host_indices, + device_indices, + i, + self.io_backend, + ) producer_event.complete(i) self._record_transfer_indices_on_stream( self.load_stream,