From 80797c2a1558e333751c7fe1f62e7727dade63c9 Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Mon, 17 Nov 2025 21:22:20 +0800 Subject: [PATCH] [PD][bug fix] fix memleak when last_batch is none (#13144) Signed-off-by: Xuchun Shang Co-authored-by: Shangming Cai --- python/sglang/srt/disaggregation/prefill.py | 51 +++++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 69eb22f30..952374ed5 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -597,27 +597,38 @@ class SchedulerDisaggregationPrefillMixin: return transferred_rids def process_prefill_chunk(self: Scheduler) -> None: + chunked_req_to_exclude = set() + if self.chunked_req: + chunked_req_to_exclude.add(self.chunked_req) + self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) + if self.enable_overlap: + # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved + self.chunked_req.tmp_end_idx = min( + len(self.chunked_req.fill_ids), + len(self.chunked_req.origin_input_ids), + ) + else: + self.send_kv_chunk(self.chunked_req) + # chunked request keeps its rid but will get a new 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 + if self.last_batch and self.last_batch.forward_mode.is_extend(): - if self.chunked_req: - # Move the chunked request out of the batch so that we can merge - # only finished requests to running_batch. - self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req) - self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) - if self.enable_overlap: - # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved - self.chunked_req.tmp_end_idx = min( - len(self.chunked_req.fill_ids), - len(self.chunked_req.origin_input_ids), - ) - else: - self.send_kv_chunk(self.chunked_req) - # chunked request keeps its rid but will get a new 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) + if self.last_batch.chunked_req: + # In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req. + # We need to discard it. + chunked_req_to_exclude.add(self.last_batch.chunked_req) + + last_bs = self.last_batch.batch_size() + self.last_batch.filter_batch( + chunked_req_to_exclude=list(chunked_req_to_exclude) + ) + if self.last_batch.batch_size() < last_bs: self.running_batch.batch_is_full = False def send_kv_chunk(