fix: pause_generation should not populate running_batch on prefill nodes (#20273)

This commit is contained in:
Lawrence Wu
2026-04-03 16:16:06 -07:00
committed by GitHub
parent 5118295f7b
commit 9593d434c4
2 changed files with 214 additions and 2 deletions
+9 -2
View File
@@ -3278,7 +3278,14 @@ class Scheduler(
self.last_batch.filter_batch(
chunked_req_to_exclude=list(chunked_req_to_exclude)
)
if not self.last_batch.is_empty():
# Skip merge for disagg prefill: completed prefill requests are
# already in disagg_prefill_inflight_queue. Merging them into
# running_batch leaks them, since the prefill event loop never
# calls update_running_batch to clean them up.
if (
not self.last_batch.is_empty()
and self.disaggregation_mode != DisaggregationMode.PREFILL
):
if self.running_batch.is_empty():
self.running_batch = self.last_batch
else:
@@ -3287,7 +3294,7 @@ class Scheduler(
self.last_batch = None
self.cur_batch = None
if recv_req.mode == "retract":
if recv_req.mode == "retract" and not self.running_batch.is_empty():
self.running_batch.filter_batch(v1_spec_info_filtered=True)
if len(self.running_batch.reqs) != 0:
retracted_reqs = self.running_batch.retract_all(self.server_args)