[Fix] use torch.cat instead of torch.concat to prevent entering the Autograd backends. (#4466)
This commit is contained in:
@@ -172,7 +172,7 @@ class TokenToKVPoolAllocator:
|
||||
return
|
||||
|
||||
if self.is_not_in_free_group:
|
||||
self.free_slots = torch.concat((self.free_slots, free_index))
|
||||
self.free_slots = torch.cat((self.free_slots, free_index))
|
||||
else:
|
||||
self.free_group.append(free_index)
|
||||
|
||||
@@ -183,7 +183,7 @@ class TokenToKVPoolAllocator:
|
||||
def free_group_end(self):
|
||||
self.is_not_in_free_group = True
|
||||
if self.free_group:
|
||||
self.free(torch.concat(self.free_group))
|
||||
self.free(torch.cat(self.free_group))
|
||||
|
||||
def clear(self):
|
||||
# The padded slot 0 is used for writing dummy outputs from padded tokens.
|
||||
@@ -739,7 +739,7 @@ class HostKVCache(abc.ABC):
|
||||
@synchronized
|
||||
def free(self, indices: torch.Tensor) -> int:
|
||||
self.mem_state[indices] = MemoryStateInt.IDLE
|
||||
self.free_slots = torch.concat([self.free_slots, indices])
|
||||
self.free_slots = torch.cat([self.free_slots, indices])
|
||||
self.can_use_mem_size += len(indices)
|
||||
return len(indices)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user