From 96158fa110535b9964beceb28476125bfe3db287 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Tue, 16 Jun 2026 02:47:35 +0800 Subject: [PATCH] Bound overlap prefill to one pending CP batch CP shared-KV batch planning needs the immediately previous prepared batch to remain visible as a virtual prefix, but launching another batch before processing that result can leave multiple prepared plans racing against radix insertion. The event loop now processes the previous result after planning the current batch and before launching it, preserving one-batch lookback without accumulating deeper overlap state. Constraint: CP HiCache prepared batch views are inserted during process_batch_result, not at forward launch. Rejected: Process previous result before planning current batch | loses visibility of the previous pending prepared plan needed by current planning. Confidence: medium Scope-risk: moderate Directive: Do not increase non-PP overlap depth for CP shared-KV without adding pending-radix reservation semantics. Tested: Remote cjy-glm5-new pytest test/registered/unit/disaggregation/test_overlap_disagg_prefill_event_loop.py test/registered/unit/disaggregation/test_overlap_disagg_decode_event_loop.py Not-tested: Full E2E latency impact of reduced overlap depth. (cherry picked from commit 7df8723eee8203e9e334b229178e3f8bac61396a) --- python/sglang/srt/disaggregation/prefill.py | 28 ++++-- .../test_overlap_disagg_prefill_event_loop.py | 96 +++++++++++++++++++ 2 files changed, 116 insertions(+), 8 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_overlap_disagg_prefill_event_loop.py 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()