[BugFix] Fix prefill memory leak in PD + GDN (#12994)
This commit is contained in:
@@ -612,7 +612,12 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
else:
|
||||
self.send_kv_chunk(self.chunked_req)
|
||||
# chunked request keeps its rid but will get a new req_pool_idx
|
||||
self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
|
||||
if self.tp_worker.model_runner.mambaish_config is not None:
|
||||
self.req_to_token_pool.free(
|
||||
self.chunked_req.req_pool_idx, free_mamba_cache=False
|
||||
)
|
||||
else:
|
||||
self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
|
||||
self.running_batch.batch_is_full = False
|
||||
|
||||
def send_kv_chunk(
|
||||
|
||||
Reference in New Issue
Block a user