Fix flush cache (#627)

This commit is contained in:
Lianmin Zheng
2024-07-15 19:56:55 -07:00
parent 56f5fc4ab5
commit 41d1f67704
5 changed files with 30 additions and 20 deletions

View File

@@ -98,6 +98,8 @@ class TokenToKVPool:
self.can_use_mem_size += len(free_index)
def clear(self):
self.prefetch_buffer = torch.empty(0, device="cuda", dtype=torch.int32)
self.mem_state.fill_(True)
self.can_use_mem_size = self.size