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
|
||||
|
||||
Reference in New Issue
Block a user