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

[Generative Score API] Fix on prefill-only scheduler running batch loss track problem (#14320)

Browse Source 
Co-authored-by: Wenyan Yao <wenyao@linkedin.com>
Co-authored-by: Sundara Raman Ramachandran <sundar24295@gmail.com>
This commit is contained in:
haNa-meister
2026-03-11 13:15:50 -07:00
committed by GitHub
co-authored by Wenyan Yao Sundara Raman Ramachandran
parent a54d71e967
commit 252ef90fc2
1 changed files with 14 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+14 -4
View File
@@ -2059,14 +2059,21 @@ class Scheduler(
self.running_batch.batch_is_full = False self.running_batch.batch_is_full = False
# Merge the new batch into the running batch. # Merge the new batch into the running batch.
# For prefill-only batch, we can avoid going through decoding step. if not self.last_batch.is_empty():
if not self.last_batch.is_empty() and not self.last_batch.is_prefill_only:
if self.running_batch.is_empty(): if self.running_batch.is_empty():
self.running_batch = self.last_batch self.running_batch = self.last_batch
else: else:
# Merge running_batch with prefill batch # Merge running_batch with prefill batch
self.running_batch.merge_batch(self.last_batch) self.running_batch.merge_batch(self.last_batch)
# For prefill-only batch, filter out finished requests since they
# won't go through the decode step. This keeps running_batch accurate
# for load reporting (num_running_reqs via /get_load).
# Runs outside the last_batch block so stale requests are cleaned
# even when no new batches arrive (e.g. traffic stops).
if self.running_batch.is_prefill_only:
self.running_batch.filter_batch()
if self.dllm_config is not None: if self.dllm_config is not None:
new_batch = self.get_new_batch_dllm() new_batch = self.get_new_batch_dllm()
else: else:
@@ -2085,8 +2092,11 @@ class Scheduler(
# Run prefill first if possible # Run prefill first if possible
ret = new_batch ret = new_batch
else: else:
# Run decode # Run decode (skip for prefill-only batches)
if not self.running_batch.is_empty(): if (
not self.running_batch.is_empty()
and not self.running_batch.is_prefill_only
):
self.running_batch = self.update_running_batch(self.running_batch) self.running_batch = self.update_running_batch(self.running_batch)
ret = self.running_batch if not self.running_batch.is_empty() else None ret = self.running_batch if not self.running_batch.is_empty() else None
else: else: