[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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user