[BugFix] Fix prefill memory leak in PD + GDN (#12994)

This commit is contained in:
Ziming Huang
2025-11-12 00:48:07 +08:00
committed by GitHub
parent d52800dbf5
commit 5de25f786e

View File

@@ -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(