[HiCache]Fix hybrid model move_indices (#22940)

Co-authored-by: hzh0425 <hzh0425@apache.org>
Co-authored-by: flyerming <flyerming@163.com>
This commit is contained in:
huangtingwei
2026-04-21 00:15:39 -07:00
committed by GitHub
co-authored by hzh0425 flyerming
parent 900aad5f72
commit efa71ce5ab
2 changed files with 49 additions and 14 deletions
@@ -683,7 +683,9 @@ class HiCacheController:
return
op = CacheOperation.merge_ops(self.write_queue)
host_indices, device_indices = self.move_indices(op)
host_indices, device_indices = self.move_indices(
op.host_indices, op.device_indices
)
self.write_queue.clear()
start_event = device_module.Event()
@@ -723,8 +725,7 @@ class HiCacheController:
)
return device_indices
def move_indices(self, op: CacheOperation):
host_indices, device_indices = op.host_indices, op.device_indices
def move_indices(self, host_indices: torch.Tensor, device_indices: torch.Tensor):
# move indices to GPU if using kernels, to host if using direct indexing
if self.io_backend == "kernel":
if not host_indices.is_cuda:
@@ -752,7 +753,9 @@ class HiCacheController:
producer_id = self.layer_done_counter.update_producer()
op = CacheOperation.merge_ops(self.load_queue)
host_indices, device_indices = self.move_indices(op)
host_indices, device_indices = self.move_indices(
op.host_indices, op.device_indices
)
self.load_queue.clear()
producer_event = self.layer_done_counter.events[producer_id]
producer_event.start_event.record()