chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix cache_tokens calculate issue when retracted (#11900)

Browse Source 
Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com>
Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
This commit is contained in:
Mike Qiu
2025-11-17 21:25:38 +08:00
committed by GitHub
parent 80797c2a15
commit 9ba3597d2a
2 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
View File

@@ -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

10
python/sglang/srt/managers/schedule_batch.py
View File

@@ -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