@@ -64,6 +64,7 @@ def handle_finished_reqs(self, finished_reqs: List[InferReq]) -> List[InferReq]:
64
64
# 如果开启了cpu cache,将达到finished状态的请求开启将gpu kv cache 卸载到 cpu cache中的操作。
65
65
# 当 kv cache 卸载完成后,才会进行请求的真实退出操作。
66
66
true_finished_reqs = []
67
+ cpu_stream = g_infer_context .get_cpu_kv_cache_stream ()
67
68
for req in finished_reqs :
68
69
# 只有 group_req_id 和 request_id 相同的请求才会被卸载到 cpu cache 中。
69
70
# 这个限制是为了兼容 diverse 模式下的请求处理。
@@ -87,11 +88,10 @@ def handle_finished_reqs(self, finished_reqs: List[InferReq]) -> List[InferReq]:
87
88
# 发起将请求的 kv cache 卸载到 cpu cache 中的任务
88
89
# if self.backend.is_master_in_dp:
89
90
# mark_start("blueswhen offload_kv_to_cpu")
90
- torch . cuda . synchronize ( )
91
+ cpu_stream . wait_stream ( g_infer_context . overlap_stream )
91
92
trans_task = self ._start_kv_cache_offload_task (
92
- req = req , cpu_kv_cache_stream = g_infer_context . get_cpu_kv_cache_stream ()
93
+ req = req , cpu_kv_cache_stream = cpu_stream
93
94
)
94
- torch .cuda .synchronize ()
95
95
# if self.backend.is_master_in_dp:
96
96
# mark_end("blueswhen offload_kv_to_cpu")
97
97
@@ -100,6 +100,7 @@ def handle_finished_reqs(self, finished_reqs: List[InferReq]) -> List[InferReq]:
100
100
else :
101
101
true_finished_reqs .append (req )
102
102
103
+ cpu_stream .synchronize ()
103
104
return true_finished_reqs
104
105
else :
105
106
return finished_reqs
@@ -217,7 +218,7 @@ def fill_cpu_cache_to_reqs(self, reqs: List[InferReq]):
217
218
218
219
need_token_num = match_tokens - req .cur_kv_len
219
220
# 多匹配了一定数量的token 才进行复制操作,不然操作效率不高
220
- if need_token_num > 128 :
221
+ if need_token_num >= 64 :
221
222
if need_token_num <= idle_token_num :
222
223
if self .backend .radix_cache is not None :
223
224
g_infer_context .radix_cache .free_radix_cache_to_get_enough_token (need_token_num = need_token_num )
0 commit comments