[metamergebot]

Based on main. The new test follows the checkpoint_disagg_prefill / checkpoint_kv_cache names from #41520 and the release-time is_insert=False from #42194.

Motivation

In the overlap event loop of a disaggregated prefill instance, process_batch_result blocks in copy_done.synchronize() until the oldest forward's sampled tokens have been copied to the host. While the scheduler thread sits in that wait it receives nothing: new requests, bootstrapped requests, aborts, pause/continue and weight updates all queue up in ZMQ, and the second forward slot stays empty even when work is ready. With long prefills this turns intake into a stop-and-go pattern and leaves the GPU idle between forwards.

This change lets the prefill overlap loop keep ingesting requests while a forward's result is still pending, so a second batch can be prepared and launched behind the first one, and results are processed as soon as they are complete on every TP rank.

Modifications

• event_loop_overlap_disagg_prefill keeps at most two batches in flight. On each pass it ingests requests, tries to prepare a batch if a slot is free, and then checks whether the oldest queued result is ready (is_disagg_prefill_batch_result_ready, a non-blocking copy_done.query() reduced with MIN over the TP CPU group so all ranks agree). If it is not ready, the loop goes back to intake instead of blocking.

    +---> [1] Ingest requests; collect bootstrap arrivals
    |             |
    |             v
    |     [2] Prepare a batch if a slot is free and
    |         preparation is not marked as failed
    |             |
    |             v
    |     [3] Oldest result ready on all TP ranks?
    |             |
    +--- No ------+
    |            Yes
    |             v
    +------ Process the oldest result
  

• A failed preparation while a result is pending sets a marker so the loop does not retry admission on every poll; new input or newly bootstrapped requests clear it immediately, and it is always cleared once a result finishes. Storage prefetch retry deadlines only advance on scheduling steps, not on polls, and process_prefill_chunk runs once per launched batch (ScheduleBatch.prefill_chunk_processed).

• ingest_requests(stop_at_pause=...) keeps receiving and broadcasting on every pass but, while a result is pending, retains a PauseGenerationReqInput and everything behind it until the next scheduling boundary, so a pause cannot retract a batch whose result has not been processed. All TP ranks retain the same suffix. Timeout aborts are not re-polled while dispatch is held, and _poll_timeout_aborts skips requests whose abort is already pending (req.to_finish).

• collect_inflight_reqs includes batches queued in result_queue while polling is enabled, so aborts, running timeouts and weight-version tracking see a forward that has been launched but is not yet last_batch.

• Retract pauses on a prefill instance set the chunked request's tmp_end_idx before draining the pending result, and an optimistic prefill chunk is released and requeued only when no middle chunk is still in flight; otherwise the retract drains every result and nothing would requeue it.

• The feature is gated by a new environment variable, SGLANG_ENABLE_DISAGG_PREFILL_CONTINUOUS_INPUT_POLLING (default off; named with the ENABLE_ verb per the env-var naming convention in .agents/skills/env-var-conventions), introduced by this pull request in environ.py. It is only enabled for non-speculative autoregressive generation with attention DP = CP = 1; other shapes log a warning and keep the blocking loop. With the flag off, result handling stays exactly as before. At TP = 1 a bounded time.sleep(PREFILL_INPUT_POLL_INTERVAL_S) per polling pass lets bootstrap, transfer and cache threads reacquire the GIL; the 100 µs value is an initial heuristic (documented at the constant), not a measured optimum. _dispatch_event_loop_once resets the flag so a PD role switch that selects a different loop does not inherit it.

Accuracy Tests

The change reorders scheduler work within a prefill instance and does not touch the model forward or sampling. As an end-to-end check, a flag-on vs flag-off PD serving comparison was run with public components only: one node with four NVIDIA GB300 GPUs (the shape of the 4-gpu-gb300 CI runner), Qwen/Qwen3-8B, --disaggregation-transfer-backend nixl, the Python mini load balancer and default server arguments, as TP1 (prefill on GPU 0, decode on GPU 1) and TP2 (prefill on GPUs 0-1, decode on GPUs 2-3) stacks, each freshly started. Procedure: GSM8K through sglang.test.run_eval (200 examples, greedy); json_schema requests racing plain long-prompt requests; streaming requests disconnected mid-prefill and requests aborted through /abort_request on both servers while forwards were pending; /pause_generation in in_place and in retract mode on both servers, issued into a burst of long requests and followed by /continue_generation; and a check that the prefill queues drain afterwards. The figures from that run are not reproduced in this description.

Speed Tests and Profiling

The pending-result wait is replaced by non-blocking query() polls plus one small MIN all-reduce on the TP CPU group per pass; with the flag off there is no new work. The same stacks were benchmarked with python -m sglang.bench_serving --backend sglang --dataset-name random --random-range-ratio 1.0 in a burst, a long chunked-prefill, a steady-rate and a light-rate scenario, with the prefill scheduler's CPU use sampled at idle and under load and py-spy stack samples of the prefill scheduler taken under load in both modes. The figures from those runs are not reproduced in this description.

Checklist

• Format your code according to the Format code with pre-commit: ruff, ruff-format, isort, codespell and scripts/lint/check_registered_tests.py pass on the changed files.
• Add unit tests according to the Run and add unit tests: test/registered/unit/disaggregation/test_prefill_result_polling.py drives the real prefill overlap loop with controlled intake and copy events (two-slot pipelining, deferred pause in both modes, shutdown during a wait, storage-retry pacing, chunked and optimistic prefill hand-offs, unsupported shapes falling back, and a two-process Gloo check that both ranks agree on FIFO completion). Existing scheduler tests gain cases for pending-abort suppression, inflight inventory with a queued result, idle-step accounting under polling, and single-result retract. The changed test files and the unchanged callers of the touched scheduler methods (mlx/test_scheduler_mixin.py, mlx/test_attention_patching.py, test_batch_result_processor_mamba_boundary.py, test_priority_scheduling_disaggregation.py, test_scheduler_hicache_events.py, test_scheduler_weight_version_tracking.py, test_decode_queue_cleanup.py, test_pd_role_switch.py) pass locally on CPU, one process per file.
• Update documentation according to Write documentations: the env var is documented next to its definition in environ.py.
• Provide accuracy and speed benchmark results: see the two sections above for what was run.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): 🚫 Run #37129843019
Latest PR Test (Extra): ❌ Run #37129842972
Latest PR Test (AMD ROCm 10): ❌ Run #37129843035

[metamergebot]

This port clears our bar on 0249d589790661f4f73c909a6617e9884cb69fbb: every base-a and base-b job is green.
base-c and the accelerator suites are reported but not gated, because they run
on scarce hardware and upstream merges through them routinely.

gate 18✓✗1✂1 lint✓ build✓ A✓ B✓ C 0/10⏳ other✓ · accel ✗amd,npu

@merrymercy @Jialin — could one of you take a look and merge it if you are happy? Nothing else is waiting on our side.

[metamergebot]

This port clears our bar on f4d659d730c1baf16b8ee1db56f79bacd38ab896: every base-a and base-b job is green.
base-c and the accelerator suites are reported but not gated, because they run
on scarce hardware and upstream merges through them routinely.

gate✓ lint✓ build✓ A✓ B✓ C 0/10⏳ other✓ · accel ✗amd,npu

@merrymercy @Jialin — could one of you take a look and merge it if you are happy? Nothing else is waiting on our side.

[metamergebot]

This port clears our bar on 275a31d695590a351078c920412aa86e1861c495: every base-a and base-b job is green.
base-c and the accelerator suites are reported but not gated, because they run
on scarce hardware and upstream merges through them routinely.

gate 20✓✗1 lint✓ build✓ A✓ B✓ C 0/10⏳ other✓ · accel ✗amd,npu,xpu

@merrymercy @Jialin — could one of you take a look and merge it if you are happy? Nothing else is waiting on our side.