Repository navigation
[VLM] Retire aborted disaggregated prefill results - #36988
Merged
Merged
Conversation
mickqian
requested review from
ByronHsu,
Duyi-Wang,
HaiShaw,
ShangmingCai,
hnyls2002 and
sogalin
as code owners
August 29, 2026 07:14
3 of 5 tasks
4 tasks done
Open
2 tasks done
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
3 of 5 tasks
3 of 4 tasks
ShangmingCai
added a commit
to ShangmingCai/sglang
that referenced
this pull request
Sep 15, 2026
`Req.set_finish_with_abort()` records the rejection in `to_finish` rather than `finished_reason`, and swaps the prompt for a one-token stub. `finished()` reads `finished_reason`, so the request still reports itself as live and the disaggregation admission queues let it through. sgl-project#36988 added a post-forward retire path that catches it once the prefill result comes back, so the client does get its 400 -- but only after a bootstrap handshake, a metadata buffer, an RDMA sender init and a forward pass on the stub, and after the decode worker has entered it into the prealloc queue and waited out the transfer timeout. Only disaggregation is affected. In NULL mode the stub costs one cheap forward pass and the batch boundary promotes `to_finish` into the 400, which is what `set_finish_with_abort` intends, so that path is left alone. The two PD admission doors reject the request instead of enqueuing it, and the scheduler exposes the retirement itself because the session, beam and output state it has to unwind live there. `_add_request_to_queue` and its call sites are unchanged. Neither door may retire a re-entering request: "Abort method 3" marks a *running* request and `filter_batch` does not drop it, since `finished()` is still False, so a retracted, preempted or resumed request can arrive carrying a pending abort while its queue still owes a release. The two doors detect that differently, because only one of them can. A preempted prefill request still holds its metadata buffer or its host retraction backup, so `is_unadmitted_reject()` sees it. A retracted decode request holds none of those markers by the time it is requeued -- `release_req` nulled `kv.req_pool_idx`, `reset_for_retract` nulled `kv.mamba_pool_idx`, the decode metadata buffer lives on `DecodeRequest` rather than on `Req`, and `add()` itself clears `retraction_mb_id` -- while still owning the host pages `release_req` allocated, which only `retraction_restore` or `retraction_discard` free. So the decode door gates on its own `is_retracted` and `is_rebootstrap` flags rather than sniffing state. `retraction_backup` is additionally unset for a `seqlen <= 1` retraction, so no resource check could have covered that door. Skipping scheduling skips two unwinds that scheduling would have done, so both happen at the door: `StreamingSession.find_active_slot`'s pre-abort detach (without it a streaming session stays in-flight and rejects every later request on it) and `beam_coordinator.retire_group` (the group is counted in ahead of the checks that reject). The trace span is closed with the real reason instead of being auto-closed at GC time. Measured on 1P1D (2x H200, Qwen3-0.6B, mooncake) against main at 354ee46, with the prefill worker's KV pool sized so its scheduler `max_req_input_len` (1018) rejects a 1500-token prompt while the TokenizerManager's `context_length` check (4096) does not. Both workers rejecting: before after prefill HTTP 400 (0.0s) 400 (0.0s) decode HTTP timeout @300S 400 (0.0s) prefill bootstrap + metadata buf reserved not reached prefill batches for the stub 1 0 decode prealloc queue entry yes no No correctness change: both arms pass every probe. The stale-KV and wrong-status symptoms this was first written against are pre-sgl-project#36988 and do not reproduce on main. Still open: when only the prefill worker rejects -- an asymmetrically sized fleet -- the decode worker holds a valid request and still waits out the full transfer timeout. Notifying the peer before any send would make that fail fast; it is not part of this change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
A VLM encoder or embedding-validation failure can abort a disaggregated prefill request after KV blocks, hybrid Mamba state, HiCache state, or metadata buffers have already been allocated. The completed GPU result still returns later. The old result path could leave those allocations owned by the aborted request; the next idle
InvariantCheckerthen reported the exact leaked KV/Mamba state and restarted the scheduler.The invariant remains fatal. This change fixes ownership cleanup before the invariant is reached instead of suppressing it.
What
The cleanup is generic to disaggregated prefill, but this PR is scoped to the failure path exposed by malformed VLM encoder results.
Validation
test_prefill_abort_result_cleanup.py: final, middle-chunk, delayed-result, sender-failure, and Mamba-before-KV casesNo exception is caught around
InvariantChecker; any remaining accounting bug still terminates loudly.CI States
Latest PR Test (Base): ✅ Run #33930569184
Latest PR Test (Extra): ❌ Run #33930569004
Latest PR Test (AMD ROCm 7.2): ❌ Run #33930569160