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:
laoyao0822
2026-06-19 06:14:20 +00:00
committed by leavelet
parent 512fe92a83
commit 96158fa110
2 changed files with 116 additions and 8 deletions
@@ -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()