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-16 02:47:35 +08:00
committed by leavelet
parent 512fe92a83
commit 96158fa110
2 changed files with 116 additions and 8 deletions

View File

@@ -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