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

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
parent a54d71e967
commit 252ef90fc2

View File

@@ -2059,14 +2059,21 @@ class Scheduler(
self.running_batch.batch_is_full = False
# 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() and not self.last_batch.is_prefill_only:
if not self.last_batch.is_empty():
if self.running_batch.is_empty():
self.running_batch = self.last_batch
else:
# Merge running_batch with prefill 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:
new_batch = self.get_new_batch_dllm()
else:
@@ -2085,8 +2092,11 @@ class Scheduler(
# Run prefill first if possible
ret = new_batch
else:
# Run decode
if not self.running_batch.is_empty():
# Run decode (skip for prefill-only batches)
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)
ret = self.running_batch if not self.running_batch.is_empty() else None
else: