From 9ba3597d2a00f4f328103c053740b95fbea668cf Mon Sep 17 00:00:00 2001 From: Mike Qiu Date: Mon, 17 Nov 2025 21:25:38 +0800 Subject: [PATCH] Fix cache_tokens calculate issue when retracted (#11900) Signed-off-by: Michael Qiu Co-authored-by: Mike_Qiu --- .../srt/disaggregation/decode_schedule_batch_mixin.py | 5 +++-- python/sglang/srt/managers/schedule_batch.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 41f73caf4..80d29e9bb 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -61,8 +61,9 @@ class ScheduleBatchDisaggregationDecodeMixin: seq_len - pre_len == req.extend_input_len ), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}" - req.cached_tokens += pre_len - req.already_computed - req.already_computed = seq_len + if not req.retracted_stain: + req.cached_tokens += pre_len - req.already_computed + req.already_computed = seq_len req.is_retracted = False pre_lens.append(pre_len) req.extend_logprob_start_len = 0 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 326b010b2..e1b986961 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -572,6 +572,8 @@ class Req: # For retraction self.is_retracted = False + # Indicates if the req has ever been retracted. + self.retracted_stain = False # Incremental streamining self.send_token_offset: int = 0 @@ -947,6 +949,7 @@ class Req: self.swa_uuid_for_lock = None self.extend_input_len = 0 self.is_retracted = True + self.retracted_stain = True self.input_token_logprobs = None self.temp_input_top_logprobs_val = None self.temp_input_top_logprobs_idx = None @@ -1313,8 +1316,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): multimodal_inputs.append(req.multimodal_inputs) - req.cached_tokens += pre_len - req.already_computed - req.already_computed = seq_len + # Only calculate cached_tokens once. Once retracted, the 'retracted_stain' + # flag will always True + if not req.retracted_stain: + req.cached_tokens += pre_len - req.already_computed + req.already_computed = seq_len req.is_retracted = False # Compute the relative logprob_start_len in an extend batch