Fix: Complete int32 to int64 conversion (#4465)
This commit is contained in:
@@ -622,11 +622,10 @@ class HostKVCache(abc.ABC):
|
||||
self.mem_state = torch.zeros(
|
||||
(self.size,), dtype=torch.uint8, device=self.device
|
||||
)
|
||||
self.free_slots = torch.arange(self.size, dtype=torch.int32)
|
||||
self.can_use_mem_size = self.size
|
||||
|
||||
# A lock for synchronized operations on memory allocation and state transitions.
|
||||
self.lock = threading.RLock()
|
||||
self.clear()
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_size_per_token(self):
|
||||
@@ -656,7 +655,7 @@ class HostKVCache(abc.ABC):
|
||||
def clear(self):
|
||||
self.mem_state.fill_(0)
|
||||
self.can_use_mem_size = self.size
|
||||
self.free_slots = torch.arange(self.size, dtype=torch.int32)
|
||||
self.free_slots = torch.arange(self.size, dtype=torch.int64)
|
||||
|
||||
@synchronized
|
||||
def get_state(self, indices: torch.Tensor) -> MemoryStateInt:
|
||||
|
||||
Reference in New Issue
Block a user