Fix cache_tokens calculate issue when retracted

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -61,8 +61,9 @@ def prepare_for_prebuilt(self: ScheduleBatch):
                     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

python/sglang/srt/managers/schedule_batch.py

@@ -572,6 +572,8 @@ def __init__(
 
         # 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 @@ def reset_for_retract(self):
         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 @@ def prepare_for_extend(self):
 
             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