[Minor] Improve the style and fix flaky tests (#1584)

This commit is contained in:
Lianmin Zheng
2024-10-06 00:10:48 -07:00
committed by GitHub
parent 2422de5193
commit 9244f27f0a
7 changed files with 15 additions and 7 deletions

View File

@@ -31,10 +31,13 @@ class ReqToTokenPool:
self.size = size
self.max_context_len = max_context_len
self.device = device
self.free_slots = list(range(size))
self.req_to_token = torch.empty(
(size, max_context_len), dtype=torch.int32, device=device
)
self.free_slots = list(range(size))
def available_size(self):
return len(self.free_slots)
def alloc(self, need_size: int) -> List[int]:
if need_size > len(self.free_slots):