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