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

Refactor: Moving extend_logprob_start_len calculation out of prepare_for_extend (#16105)

Browse Source
This commit is contained in:
Cheng Wan
2025-12-30 12:38:33 +08:00
committed by GitHub
parent 9e263c2162
commit 60f1ca6925
7 changed files with 25 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_policy.py
+4 -4
View File
@@ -454,7 +454,7 @@ class PrefillAdder:
if _rem_tokens <= 0:
_rem_tokens = self.rem_chunk_tokens
truncated = req.extend_input_len > _rem_tokens
req.extend_input_len = min(req.extend_input_len, _rem_tokens)
req.set_extend_input_len(min(req.extend_input_len, _rem_tokens))
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
self.can_run_list.append(req)
self._update_prefill_budget(
@@ -559,7 +559,7 @@ class PrefillAdder:
# Chunked prefill
trunc_len = self.rem_chunk_tokens
req.extend_input_len = trunc_len
req.set_extend_input_len(trunc_len)
req.fill_ids = req.fill_ids[:trunc_len]
self.can_run_list.append(req)
self.new_chunked_req = req
@@ -608,7 +608,7 @@ class PrefillAdder:
req.last_host_node, req.host_hit_length
)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
prefix_len = len(req.prefix_indices)
req.cache_protected_len = prefix_len
@@ -651,7 +651,7 @@ class PrefillAdder:
)
# Chunked prefill
req.extend_input_len = trunc_len
req.set_extend_input_len(trunc_len)
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + trunc_len]
self.can_run_list.append(req)