Add a new event loop (#1677)

This commit is contained in:
Lianmin Zheng
2024-10-16 01:33:20 -07:00
committed by GitHub
parent a5114b6f91
commit 9116b2896f
9 changed files with 161 additions and 25 deletions
+11 -5
View File
@@ -38,12 +38,16 @@ class ChunkCache(BasePrefixCache):
max_prefix_len = len(key)
return entry.value[:max_prefix_len], entry
def cache_finished_req(self, req: Req, token_ids: Optional[List[int]] = None):
def cache_finished_req(
self, req: Req, token_ids: Optional[List[int]] = None, free_delta: int = 0
):
if token_ids is None:
token_ids = (req.origin_input_ids + req.output_ids)[:-1]
token_id_len = len(req.origin_input_ids) + len(req.output_ids) - 1
else:
token_id_len = len(token_ids)
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids)
req.req_pool_idx, : token_id_len + free_delta
]
self.req_to_token_pool.free(req.req_pool_idx)
self.token_to_kv_pool.free(kv_indices)
@@ -53,10 +57,12 @@ class ChunkCache(BasePrefixCache):
def cache_unfinished_req(self, req: Req, token_ids: Optional[List[int]] = None):
if token_ids is None:
token_ids = req.fill_ids
token_id_len = len(req.fill_ids)
else:
token_id_len = len(token_ids)
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids)
req.req_pool_idx, :token_id_len
]
if req.rid not in self.entries: