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)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user