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,