diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9892392e2..53cd2ddf5 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -835,6 +835,10 @@ class SchedulerDisaggregationPrefillMixin: def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.result_queue = deque() + def pop_and_process_last_batch() -> None: + tmp_batch, tmp_result = self.result_queue.popleft() + self.process_batch_result(tmp_batch, tmp_result) + while True: # Receive requests recv_reqs = self.recv_requests() @@ -847,22 +851,30 @@ class SchedulerDisaggregationPrefillMixin: batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch + # Process the previous batch after planning the current batch but + # before launching it. CP shared-KV overlap planning may need the + # immediately previous prepared batch as a virtual/pending prefix, + # but launching the current batch would create another prepared plan. + # This keeps the non-PP overlap depth at one prepared batch. + if self.last_batch: + pop_and_process_last_batch() + elif batch is None: + # When the server is idle, do self-check and re-init some states + self.self_check_during_idle() + # Launch the current batch if batch: self._register_per_layer_transfers(batch) batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) + if len(self.result_queue) > 1: + raise RuntimeError( + "Overlap disagg prefill expected at most one queued " + f"batch result, got {len(self.result_queue)}" + ) else: batch_result = None - # Process the last batch - if self.last_batch: - tmp_batch, tmp_result = self.result_queue.popleft() - self.process_batch_result(tmp_batch, tmp_result) - elif batch is None: - # When the server is idle, do self-check and re-init some states - self.self_check_during_idle() - self.process_disagg_prefill_inflight_queue() # Run sample of the current batch diff --git a/test/registered/unit/disaggregation/test_overlap_disagg_prefill_event_loop.py b/test/registered/unit/disaggregation/test_overlap_disagg_prefill_event_loop.py new file mode 100644 index 000000000..6a308a90a --- /dev/null +++ b/test/registered/unit/disaggregation/test_overlap_disagg_prefill_event_loop.py @@ -0,0 +1,96 @@ +"""Unit tests for overlap disaggregation prefill event loop ordering.""" + +import unittest +from types import SimpleNamespace +from typing import Any, cast + +from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=6, suite="stage-a-test-cpu") + + +class FakeBatch: + def __init__(self, label: str): + self.label = label + + def copy(self): + return self + + +class TestOverlapDisaggPrefillEventLoop(CustomTestCase): + def test_processes_previous_result_after_planning_before_current_launch(self): + class StopLoop(Exception): + pass + + batch1 = FakeBatch("batch-1") + batch2 = FakeBatch("batch-2") + batches = iter([batch1, batch2]) + events = [] + + def get_next_batch(): + batch = next(batches) + events.append(f"get:{batch.label}") + return batch + + def run_batch(batch): + result = f"result-{batch.label}" + events.append(f"run:{batch.label}") + return result + + def process_batch_result(batch, result): + events.append(f"process:{batch.label}") + + def launch_batch_sample_if_needed(result): + events.append(f"sample:{result}") + if result == "result-batch-2": + raise StopLoop() + + scheduler = cast( + Any, + SimpleNamespace( + last_batch=None, + recv_requests=lambda: [], + process_input_requests=lambda reqs: None, + waiting_queue=[], + disagg_prefill_bootstrap_queue=SimpleNamespace( + pop_bootstrapped=lambda: [] + ), + get_next_disagg_prefill_batch_to_run=get_next_batch, + _register_per_layer_transfers=lambda batch: None, + run_batch=run_batch, + process_batch_result=process_batch_result, + process_disagg_prefill_inflight_queue=lambda: events.append("inflight"), + launch_batch_sample_if_needed=launch_batch_sample_if_needed, + self_check_during_idle=lambda: None, + ), + ) + + with self.assertRaises(StopLoop): + SchedulerDisaggregationPrefillMixin.event_loop_overlap_disagg_prefill( + scheduler + ) + + self.assertEqual( + events, + [ + "get:batch-1", + "run:batch-1", + "inflight", + "sample:result-batch-1", + "get:batch-2", + "process:batch-1", + "run:batch-2", + "inflight", + "sample:result-batch-2", + ], + ) + self.assertEqual(len(scheduler.result_queue), 1) + queued_batch, queued_result = scheduler.result_queue[0] + self.assertIs(queued_batch, batch2) + self.assertEqual(queued_result, "result-batch-2") + + +if __name__ == "__main__": + unittest.main()