From 862a909a085f68f2e87a14e5bb7f8c930612c955 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 2 Sep 2026 02:09:01 -0700 Subject: [PATCH] [Fix] Lock PP dynamic-chunk profiling requests before releasing through the tree cache (#37509) --- python/sglang/srt/managers/schedule_batch.py | 5 ++--- python/sglang/srt/managers/scheduler.py | 2 +- python/sglang/srt/managers/scheduler_pp_mixin.py | 7 +++++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 0552fbde2..2323aa7af 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1371,9 +1371,8 @@ class Req(ReqDllmMixin): appending only the new output tokens. Falls back to a full rebuild when the in-place append is invalid: - - aliasing: scheduler_pp_mixin assigns full_untruncated_fill_ids = - origin_input_ids directly, so extending in place would write output - tokens into the origin; + - aliasing: full_untruncated_fill_ids is origin_input_ids itself, so + extending in place would write output tokens into the origin; - lengths disagree: fresh req (array still empty), retraction (output_ids reset to empty), or set_finish_with_abort (origin replaced by a 1-token stub). diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4af4753fb..eb3f566a0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1253,7 +1253,7 @@ class Scheduler( self.profile_and_init_predictor() except Exception as e: logger.warning( - f"[PP Dynamic Chunk] Failed to profile prefill latency: {e}. " + f"[PP Dynamic Chunk] Failed to profile prefill latency: {e!r}. " "Dynamic chunking will be disabled." ) self.enable_dynamic_chunking = False diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 009dd2ce2..b8994cd30 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -634,8 +634,11 @@ class SchedulerPPMixin: origin_input_ids=input_ids, sampling_params=sampling_params, ) - req.full_untruncated_fill_ids = req.origin_input_ids - req.logprob_start_len = -1 + # Walk the same match -> lock -> alloc lifecycle as a scheduled + # request so release_kv_cache can release it symmetrically. + req.init_next_round_input(self.tree_cache) + lock = self.tree_cache.inc_lock_ref(req.last_node) + req.swa_uuid_for_lock = lock.swa_uuid_for_lock req.set_extend_range( len(req.prefix_indices), len(req.full_untruncated_fill_ids) )