Fix missing draft KV pool transfers in HybridCacheController (#27953)
This commit is contained in:
@@ -410,6 +410,13 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
self.io_backend,
|
self.io_backend,
|
||||||
pool_transfers=resolved_pool_transfers,
|
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()
|
finish_event.record()
|
||||||
self._record_transfer_indices_on_stream(
|
self._record_transfer_indices_on_stream(
|
||||||
self.write_stream,
|
self.write_stream,
|
||||||
@@ -484,6 +491,18 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
self.io_backend,
|
self.io_backend,
|
||||||
pool_transfers=resolved_pool_transfers,
|
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)
|
producer_event.complete(i)
|
||||||
self._record_transfer_indices_on_stream(
|
self._record_transfer_indices_on_stream(
|
||||||
self.load_stream,
|
self.load_stream,
|
||||||
|
|||||||
Reference in New Issue
Block a user