[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:
co-authored by
hzh0425
flyerming
parent
900aad5f72
commit
efa71ce5ab
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user