[Scheduler] Fix final chunked-prefill abort commit race

[jeremyzhang866]

Motivation

Fixes #34149.
Related prior report: #34112.

While testing chunked-prefill cancellation on latest main, we independently reproduced the same user-visible symptom reported in #34112: a cancelled request can expose one token and then finish with abort.

The #34112 attachment used SGLang v0.5.9. Its negative batch-level output IDs were old FutureMap relay handles and are not the failure criterion or target of this PR. Current main no longer uses that ScheduleBatch.output_ids representation.

Our deterministic current-main reproduction isolates a final-prefill overlap race. An abort can arrive after the final chunked-prefill EXTEND is launched but before its delayed result is committed. At that point scheduler.chunked_req is already clear, Req.output_ids is empty, and the result is pending in result_queue.

Root cause

There were two commit boundaries to protect:

1. abort_request() records req.to_finish = FINISH_ABORT(), but the prefill result processor previously appended the sampled token and token-derived metadata before update_finish_state() consumed the pending finish.
2. In the CUDA/HIP overlap loop, scheduling the next batch happens before processing the delayed final-prefill result. Without an admission guard, the pending-finish prefill request can enter optimistic decode, increment decode_batch_idx / the KV committed watermark, and perform an optimistic decode step and its associated allocation before the result is dropped (which may reserve multiple slots or pages under speculative decoding).

Filtering only in the output streamer is too late: request output, reasoning state, logprobs, hidden states, sampling metadata, grammar state, and KV accounting may already be mutated.

Modifications

• When a delayed result is actually pending, exclude pending-finish prefill requests from the prefill-to-decode merge before optimistic decode preparation.
• Keep decode requests in mixed batches on their existing partial-output path.
• In the final-prefill result processor, promote the pending finish with new_accepted_len=0 and reuse normal finished-request cleanup.
• Do not commit the dropped token's output ID, reasoning state, hidden state, customized info, sampling mask, returned logprob values, or grammar state.
• Still advance packed hidden-state and logprob cursors so later live requests in the same batch read their own slices.
• Preserve synchronous, MLX, and PP scheduling semantics by applying the merge guard only while result_queue contains a delayed result.
• Leave PD-prefill transfer lifecycle unchanged; it uses a separate scheduler/result path.

Regression coverage

The scripted test deterministically waits for this window:

• the request has passed through chunked prefill;
• the final EXTEND is launched;
• scheduler.chunked_req is None;
• no middle chunks remain in flight;
• Req.output_ids is empty;
• the final result is pending.

It then aborts and verifies:

• finish reason is FINISH_ABORT;
• no token is appended;
• no optimistic decode step occurs;
• the request-pool row and KV object are released;
• reclaimable KV capacity returns to its baseline after the scheduler drains.

Additional CPU tests cover both sides of the scheduler gate (pending queue vs. already-processed result), mixed prefill/decode membership, and a two-request packed-logprob case where the aborted segment advances the cursor without attaching metadata to the aborted request.

Accuracy Tests

Cancellation control flow only; normal token generation is unchanged.

Validated on top of main commit 7605529bdf98ded2815e48c7054b1bc4533256c9:

PYTHONPATH=python .venv/bin/python -m pytest -q \
  test/registered/unit/managers/test_scheduler_chunked_req_gate.py \
  test/registered/unit/managers/test_batch_result_processor_hidden_states.py \
  test/registered/unit/managers/test_batch_result_processor_spec_grammar.py

11 passed, 2 subtests passed

SGLANG_USE_MLX=1 PYTHONPATH=python .venv/bin/python \
  test/registered/scripted_runtime/test_scripted_runtime_core.py \
  TestScriptedRuntimeCore.test_abort_at_last_chunk_does_not_append_output

Ran 1 test in 129.234s
OK

Also passed:

• direct CI-style execution of the two modified CPU test files;
• Python compile checks for all touched files;
• registered-test and package-registration repository checks;
• Black 26.1.0, Ruff 0.15.1, isort 7.0.0, codespell 2.4.1, and git diff --check.

The packed-logprob test was mutation-checked: it fails both if the aborted request attaches returned logprobs and if its packed cursor contribution is skipped.

Speed Tests and Profiling

Not applicable. The new checks are scheduler control-flow predicates; the result-drop branch runs only for a request that is already finishing.

Topology notes

The race is not caused by TP, DP, or PP. TP/DP use the same per-request result state. PP and MLX process their pending result before the relevant last_batch is reused, so the queue-gated admission change is a no-op there while the result-drop protection remains valid. PD-prefill is intentionally out of scope because it has a separate transfer and result lifecycle.

The registered end-to-end regression is TP=1 / DP=1 / PP=1. Multi-rank topology behavior was source-reviewed but not claimed as a multi-GPU runtime validation.

Checklist

• Rebased onto current main.
• Added deterministic registered regression coverage.
• Added focused scheduler and packed-metadata unit coverage.
• Ran relevant unit and scripted-runtime tests.
• Ran formatting, compile, and test-registration checks.
• Documentation and performance benchmarks are not applicable to this abort-only correctness fix.

---

CI States

Latest PR Test (Base): ❌ Run #37111470582
Latest PR Test (Extra): ❌ Run #37111470433
Latest PR Test (AMD ROCm 10): ❌ Run #37111470553

[jeremyzhang866]

@cctry @hnyls2002 Could you please take a look when you have time? This PR fixes the final chunked-prefill abort/result-commit race in the regular overlap scheduler and adds a deterministic scripted-runtime regression. The change is intentionally scoped to the non-PD path; local unit, grammar/spec, formatting, registration, and MLX end-to-end race tests pass.

[jeremyzhang866]

/tag-run-ci-label

[jeremyzhang866]

@cctry @hnyls2002 The updated head b1da9c26f9 is now Ready, rebased onto current main, and the full lint workflow is green. The GPU/CPU test workflows shown in red were gated before the Ready transition (PR Draft: true), so no test shard ran. As the fork author I cannot add the required run-ci label (can_tag_run_ci_label=false). Could one of you add run-ci or run /tag-and-rerun-ci when reviewing? The PR body contains the local scripted/CPU validation and the exact gate diagnosis.

[jeremyzhang866]

@xiezhq-hermann Could you please take a look at the current head a7380b2f71 when you have time? This PR fixes the regular-overlap final-prefill cancel/result-commit race and verifies that KV and request-pool state are fully reclaimed. Lint and the local CPU/MLX race regressions pass; the official CPU/CUDA/AMD jobs are still skipped only because the PR is missing the run-ci label. If the scope looks right, could you review it and, if you have permission, help trigger run-ci?

[jeremyzhang866]

@cctry @xiezhq-hermann A gentle follow-up on the current head 04af626fed — when you have a chance, could you help review the scheduler/result-commit scope? The PR remains mergeable against current main (a53d3636ce), and lint is green. The official CPU/CUDA/AMD jobs have not run because the PR still has no run-ci label. If the scope looks right, could one of you also help trigger the required CI? Thanks!

[jeremyzhang866]

@cctry @xiezhq-hermann Update: the branch has now been rebased onto main at 2be9773a21; the current head is b178911c3b. The patch remained identical across the rebase, and compile, registered-test checks, Ruff, Black, and isort all pass locally. Could one of you please trigger the combined tag-and-rerun CI command for this current head so the CPU/CUDA/AMD jobs actually run? Thanks!

[jeremyzhang866]

@cctry @xiezhq-hermann Final update after syncing with current main: the new head is a3702437b7. The patch scope and patch ID are unchanged; compile, registered-test checks, Ruff, Black, isort, and diff checks pass. The PR still has no run-ci label, so no CPU/CUDA/AMD test shard has run. @cctry, could you please post the standard combined tag-and-rerun CI command as a separate comment for this head? @xiezhq-hermann, when convenient, could you review the scheduler/result-commit changes? Thanks!

[jeremyzhang866]

@cctry @xiezhq-hermann Update: this PR is now rebased onto current main at 7605529bdf; the new head is 141a3e22e0, and GitHub reports it mergeable again. The only rebase conflict was the existing hidden-state test helper, resolved by preserving main's RuntimeContext setup while retaining this PR's logprob processor injection. Validation on the rebased head: focused CPU tests 11 passed, 2 subtests passed; the deterministic MLX scripted-runtime abort race test passed in 129.234s; Black, Ruff, isort, codespell, compile, CI-registration, package-registration, and diff checks all pass. The official GitHub Lint workflow also passed on this head. The CPU/CUDA/AMD shards remain gated because the PR still has no run-ci label. When convenient, could you please review the scheduler/result-commit scope and help trigger the gated CI if the scope looks right? Thanks!

[JulianZJN]

I went through the current head against the regular overlap loop and the prefill result processor.

The two protections are at the right commit boundaries:

• while result_queue holds the delayed final-prefill result, a prefill request with to_finish is excluded before the prefill-to-decode merge, so it cannot take an optimistic decode step;
• process_batch_result_prefill consumes the pending finish with an accepted length of zero before attaching the sampled token or token-derived state.

The mixed-batch exception also looks necessary: requests already in decoding_reqs retain the existing partial-output semantics. The packed-result handling is important as well—the aborted segment still advances hidden-state and logprob cursors, and the two-request test would catch a slice shift into the following live request.

I do not see a code-level blocker in this scope. The remaining validation gap is the hardware path. The original symptom was reproduced under CUDA overlap on RTX 4090, while the deterministic end-to-end regression reported here ran through MLX; the official CPU/CUDA/AMD shards are still gated on run-ci. Before merge, I would want the current head exercised by the gated one-GPU shard or the original CUDA reproducer, retaining the assertions on zero post-abort tokens, no decode_batch_idx advance, and complete KV/request-pool reclamation.

[jeremyzhang866]

@JustinTong0323 Could you please take a look when you have time? This PR complements #35255: that PR makes a deferred abort follow a request after it leaves the chunked-prefill slot, while this one prevents a delayed final-prefill result from committing a token or entering optimistic decode after the abort has won. The current head 141a3e22e0 is mergeable; focused CPU/MLX regressions and lint pass. If the scope looks right, could you also run /tag-and-rerun-ci so the one-GPU path is exercised? Thanks!

[JulianZJN]

Follow-up to my earlier review: the CUDA validation gap is now closed for the scope claimed by this PR.

I independently ran a matched red/green validation on one RTX 4090 D against the exact current head 141a3e22e0c6f03882d640174323086596a38c8b:

• unpatched base 7605529bdf98ded2815e48c7054b1bc4533256c9: the first CUDA run reproduced the target failure, chunked-prefill result appended a token after abort: [38297];
• exact PR head: the same deterministic scripted CUDA regression passed 10/10 independent server/scheduler lifecycles;
• clean synthetic merge onto main at eac91ac362f1ea9caf383cabfc2a57bf7e07b367: the focused CPU set passed (11 passed, 2 subtests passed) and the CUDA regression passed 5/5. Current main (d9f6861359) is only one docs-only commit beyond that tested revision.

The scripted oracle is substantive: it checks zero post-abort output tokens, no decode_batch_idx or KV committed-watermark advance, request-pool and KV release, restoration of reclaimable KV capacity, an empty result queue, and a fully idle scheduler after drain. The packed hidden-state/logprob cursor tests also pass.

The older #34112 attachment was not runnable as a standalone artifact because its backends.sglang.* dependency was absent. That does not weaken the causal result for the current #34149 race: the same test fails on the exact base at the intended delayed final-prefill boundary and passes on the patch.

Scope remains TP=1 / DP=1 / PP=1 regular overlap scheduling on CUDA; PD-prefill and multi-rank runtime validation are not claimed. Within that scope, I no longer see a code or hardware-validation blocker and approve this change.

The required pr-test-finish check is still red only because call-gate / pr-gate rejected the missing run-ci label and all substantive shards were skipped. @cctry @xiezhq-hermann @hnyls2002, could one of you trigger the standard tag-and-rerun CI command for this head?

[jeremyzhang866]

@JustinTong0323, quick follow-up while #35255 is in active CI: I rechecked current main at a90d770c40. It still promotes the final-prefill request without excluding a pending to_finish while its result is delayed, and the prefill result path still appends next_token_id before honoring that pending finish. So the #34153 protection has not been absorbed by main.

#35255 handles deferred abort delivery after a request leaves the chunked-prefill slot; this PR protects the downstream result-commit and optimistic-decode boundaries after that abort wins. The exact head 141a3e22e0 remains mergeable and has JulianZJN matched RTX 4090 red/green validation. Could you review it and add run-ci, or let me know if you prefer a single rebase after #35255 lands?

[jeremyzhang866]

@xiezhq-hermann @cctry Could you please take another look at #34153 when convenient?

I rechecked current main at a90d770c40. The final chunked-prefill abort/result-commit race is still present: a request with a pending to_finish can enter optimistic decode while its EXTEND result is delayed, and the result path can append next_token_id before applying the abort.

#35255 is complementary rather than a replacement: it fixes deferred abort delivery after the request leaves the chunked-prefill slot, while #34153 protects the downstream result-commit boundary after that abort wins.

The exact head 141a3e22e0 has focused CPU/MLX coverage and JulianZJN matched RTX 4090 red/green validation. The PR remains mergeable, but no run-ci CPU/GPU shard has run yet. Could one of you review the scheduler scope and add run-ci?

[jeremyzhang866]

@hnyls2002 Could you please review the updated head 30ae497fad when convenient?

I rebased the single commit onto current main at 7ef49e8f7d. The production hunks are unchanged from the previously reviewed and CUDA-validated head; range-diff shows only auxiliary_host_output=None added to a unit-test result stub for the new main API. Focused CPU tests pass (8 passed, 2 subtests passed), along with compile, Black, isort, the CI Ruff rules, registered-test checks, and diff checks.

The scripted race oracle is registered for CUDA/AMD. Current main now runs the local Mac path with overlap disabled, so I am not treating the old hardware result as validation of this new SHA. #35255 remains open and complementary. Could you review the scheduler/result-commit scope and add run-ci for this head?

[JulianZJN]

Re-reviewed the rebased current head f74c17adf0e8.

The two production protections remain at the correct boundaries for regular overlap scheduling: pending-finish prefill requests are excluded before optimistic decode while a delayed result is queued, and the result processor consumes the pending finish without committing the sampled token or token-derived state. The new beam-search branch is not a blocker for this logic: beam requests delegate finish/abort ownership to BeamCoordinator.commit_prefill(), which checks req.to_finish and aborts the group before committing staged selection.

For the claimed non-beam TP1/DP1/PP1 scope, I do not see a new code-level blocker after the rebase. My matched CUDA red/green run was against the earlier exact head 141a3e22e0c6, so this approval does not represent exact-SHA GPU validation of f74c17adf0e8; the current official run-ci shard is still required.

[jeremyzhang866]

@cctry @xiezhq-hermann Could you please take one more look at the rebased current head f74c17adf0e8 when convenient?

The Scheduler/result-commit protections remain unchanged in scope, and the exact head has been re-reviewed with no new code-level blocker identified. Lint is green, but the official CPU/GPU shards are still skipped because the PR has no run-ci label. Could one of you review the Scheduler changes and, if the scope looks right, trigger /tag-and-rerun-ci for this head? Thanks!

[jeremyzhang866]

Update on the current head:

• Rebases cleanly onto upstream/main at 756d0e0a85 ([Model] Add K2 Horizon FP8 checkpoint support (#38033)).
• New PR head: e6acb48c916470fe9b4ade613c0a261f46d0a553.
• The patch remains focused on the final chunked-prefill abort race. Current main contains the deferred-abort cleanup from Fix: abort handling for dispatched requests after client disconnect #35255, but still does not contain the final-prefill result-commit and optimistic-decode protections covered here.
• Passed locally: py_compile, registered-test checks, targeted Ruff (F401,F821,UP037), Ruff format, and git diff --check.
• The focused pytest files could not be collected in this local environment because PyTorch is not installed; no test pass is claimed from that run.

@cctry @xiezhq-hermann @hnyls2002, could you review the rebased head and, if the scope looks right, trigger the required run-ci validation?

[jeremyzhang866]

@cctry Since you merged #35255, could you please take a look at #34153 when convenient?

The two PRs cover adjacent abort boundaries: #35255 handles deferred abort delivery after a request leaves the chunked-prefill slot, while #34153 protects the delayed final-prefill result from committing a token or entering optimistic decode after the abort wins.

#34153 is now rebased onto the latest main (756d0e0a85), with current head e6acb48c916470fe9b4ade613c0a261f46d0a553. The patch remains focused and the local static/registration checks pass. The PR is still missing run-ci, which I cannot add from the fork; if the scope looks right, could you review it and help trigger the required CI?

[jeremyzhang866]

Update on #34153:

I checked the branch against the current sgl-project/sglang:main. The PR branch is now 121 commits behind main, so I will do one final rebase before asking for the official validation; no further production-scope changes are planned.

The patch remains focused on the delayed final-prefill result boundary: when abort wins, the result must not commit a token or enter optimistic decode. This is complementary to #35255, which handles deferred abort delivery after the request leaves the chunked-prefill slot.

After the rebase I will report the new head and rerun the available static/targeted checks. Could a maintainer please review the rebased head and, if the scope still looks right, add the run-ci label or run /tag-and-rerun-ci for the required official CPU/CUDA validation?

[jeremyzhang866]

Rebased onto current main and kept the change limited to the delayed final-prefill abort path. The focused CPU tests pass (14 passed, 2 subtests passed), and the PR is mergeable. Could you please take a look and help trigger run-ci when convenient?

[jeremyzhang866]

@hnyls2002 Could you take a look at this PR when you have a chance? Thanks!

[jeremyzhang866]

@hnyls2002 Could you take a look at this PR when you have a chance? Thanks!