Repository navigation
[Disaggregation] Honor admission aborts before KV handoff - #38704
Closed
kflansburg wants to merge 10 commits into
Closed
kflansburg wants to merge 10 commits into
kflansburg wants to merge 10 commits into
Conversation
added 10 commits
September 9, 2026 09:06
Preserve the original admission error instead of allocating decode pages or admitting a rejected request to optimistic prefill. Keep PP abort consensus and shared queue cleanup consistent with pending to_finish state. Notify known Mooncake destinations only through a guarded before-send path. Keep general abort and post-forward KV lifetime behavior unchanged. Reuse the bounded bootstrap wait and existing failure cleanup instead of adding room tombstones. Tests: 21 focused AST-loaded CPU regressions pass, including real Req.set_finish_with_abort methods. pre-commit --files passes with GITHUB_BASE_REF=upstream/main. Standard unit imports are blocked by missing torch/numpy; pyright is not clean in this environment. No GPU accuracy or speed benchmarks were run.
Publish total logical KV coverage in spare cached_tokens slot 7 at the final metadata write, retaining the 16-int row and buffer order. Compare against decode's preallocated kv_committed_len, including cached prefixes and rebootstrap outputs. After transfer readiness, MIN-reduce length validity across the existing attention-TP group before committing any handoff output. Abort mismatches with HTTP 500 and use pop_transferred for KV cleanup. Preserve fake transfers and the explicit zero/unknown fail-open behavior for older prefill workers. Cross-PP failure consensus is unchanged. Validation: import-light CPU regression tests written first, observed failing, then 10 tests and 14 subtests passed; direct unittest execution also passed. Pre-commit passed for all changed files with GITHUB_BASE_REF=upstream/main. Broader wire and grammar suites were blocked during collection by missing runtime dependencies (orjson). Tests follow in a separate commit. Accuracy/speed: no model or GPU runs. Valid handoff output is unchanged; this adds one CPU TP all-reduce per completed handoff, not benchmarked.
Exercise Req.set_finish_with_abort and queue/transport methods through CPU AST extraction without importing the GPU runtime. Cover original error preservation, no receiver initialization or allocation, PP filtering and abort consensus, known-peer notification, bootstrap timeout cleanup, duplicate output suppression, and refusal to notify after any send, including metadata-only sends. Update existing request fixtures with pending-finish state. Initial regressions: 14 failed, 1 passed before production changes. Final focused suite: 21 passed. pre-commit --files and commit hooks pass. Standard suites could not collect because torch/numpy are unavailable; GPU, live transport, and multiprocess collectives were not tested.
Execute the production metadata, final chunk send, poll reduction, handoff, and pop_transferred methods with real CPU tensors and mocked transport boundaries. Nonfake request hosts exercise validation, including cached prefixes, rebootstrap outputs, short/long/delta/page-rounded lengths, fake bypass, legacy zero, stale slot reuse, staging, rank rejection, and exactly-once cleanup. Update only affected metadata fixtures to supply input coverage and the actual 16-int cached_tokens layout. The handoff fixture now mocks the collective. Validation: observed expected regression failures before the functional fix. Final pytest run: 10 passed, 14 subtests passed in 1.43s. Direct unittest run: 10 tests passed in 0.670s. Pre-commit --files passed with GITHUB_BASE_REF=upstream/main. The full wire and grammar suites could not collect without runtime dependencies (orjson); no GPU, RDMA, multiprocess, accuracy, or speed testing was performed.
Only completed FINISH_ABORT requests bypass PP bootstrap readiness. Leave pending to_finish admission errors in the ready intersection so all stages can receive decode destinations before pop_bootstrapped retires their rejected requests. This relies on consistent admission validation across stages, as with PP-synchronized token capacity and grammar timeout decisions. Keep existing completed-abort and transport-failure behavior. Do not add rejected-room protocol state: metadata arriving after the existing bootstrap timeout can still recreate an orphaned room. The fix removes the new premature PP cleanup path, not that pre-existing timeout limitation. Tests: updated CPU AST regressions failed 4 cases before the change; all 27 now pass. pre-commit --files passes. Pyright reports 144 errors in the scheduler file and is not clean. No GPU accuracy, speed, live transport, or multiprocess tests were run.
The length collective added by 0349c51800 could be skipped by a TP peer taking a rank-local HiCache PENDING or FAILED branch. Reproduced the unmatched collective with two spawned Gloo ranks and a bounded five-second timeout. In pop_transferred, reduce restore completion and validity after the existing transport poll reaches a terminal consensus. Pending restores block every rank before commit or cleanup, including when a peer restore failed. A restore failure does not bypass ongoing RDMA or staging scatter. Remove the rank-local terminal branches so all TP ranks take the same path into length validation or failure cleanup. Retain transport abort/drain handling and existing cross-PP behavior. This adds a two-flag CPU reduction for HiCache requests whose transport poll is terminal; no new reduction is added with decode HiCache disabled. Validation: focused CPU suite passed 11 tests and 14 subtests, including 13 two-rank scenarios and 24 queue ticks per rank with real Gloo collectives. Direct unittest execution also passed. Pre-commit --files passed with GITHUB_BASE_REF=upstream/main. Tests follow separately. No real H2D/RDMA, GPU accuracy, or speed benchmark was run.
…data Exercise actual PP ready intersection and bad union across three consistently rejected stages. Deliver metadata through the production Mooncake registration loop with mocked transport/parsing, in both stage orders, then assert original 400 preservation, notification, cleanup, and no optimistic admission or duplicate output. Keep completed AbortReq behavior covered on each stage. Rename the old timeout test to avoid claiming total leak freedom, and characterize the existing orphan-room recreation when metadata arrives after timeout. Verification: 4 failed and 23 passed before the functional fix; 27 passed afterward. GITHUB_BASE_REF=upstream/main pre-commit run --files passes. Tests remain CPU AST-loaded simulations, not live networking or multiprocess collectives.
Run the production pop_transferred, poll helpers, metadata handling, and local restore completion state machine in two spawned CPU processes. The Gloo reductions are real; transport, scatter completion, and H2D event queries are mocked. Cover READY/PENDING, READY/FAILED, FAILED/PENDING, failure with peer RDMA or scatter still active, network failure with pending H2D, and length mismatch waiting on H2D. Check matched collective shapes/counts, no metadata reads or cleanup while waiting, and terminal output or exactly-once cleanup. The matrix covers 13 staging/nonstaging scenarios and 24 ticks per rank. Bound collectives to five seconds, per-tick process synchronization to ten seconds, and parent result collection to 45 seconds. Before the fix, this test reproduced a five-second timeout inside the length collective. After the fix, the focused suite passed 11 tests and 14 subtests in 2.65s; direct unittest execution passed 11 tests in 1.934s. Pre-commit --files passed with GITHUB_BASE_REF=upstream/main. Exact test command: uv run --python 3.12 --with torch --with numpy --with pytest --with pytest-asyncio --no-project python -m pytest test/registered/unit/disaggregation/test_transfer_sequence_length.py -q
…pletes Emit the original PP=1 prefill admission error once while retaining the bounded bootstrap wait and existing room cleanup. This lets the proxy return the nonstreaming prefill error even when decode has already returned an HTTP 200 SSE error. Keep PP readiness consensus unchanged and generate no dummy token. Validation: 27 focused admission tests and file-scoped pre-commit passed. No valid inference path or numerical changes; no live accuracy or speed run.
Exercise the real request finish methods, generation streamer, accumulator, incremental detokenizer setup, and token-ID payload with empty output IDs. Check single output across repeated bootstrap polls, later metadata or timeout cleanup, and unchanged PP waiting. Validation: 42 focused admission and transfer tests plus 14 subtests passed, including two-rank Gloo coverage. File-scoped pre-commit passed with GITHUB_BASE_REF=upstream/main.
3 of 4 tasks
rwang5203
reviewed
Sep 15, 2026
| dispatch happens later, after preallocation and ``send_metadata`` (see | ||
| ``pop_preallocated``). | ||
| """ | ||
| if is_aborted(req): |
Collaborator
There was a problem hiding this comment.
An aborted retracted request returns here without queuing or releasing its host backup, permanently consuming host slots. Could you preserve the reentry cleanup path before applying this guard?
4 of 5 tasks
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.
Motivation
Req.set_finish_with_abort()records a pending error into_finishand replaces the input with a one-token placeholder. Disaggregated admission queues must treat that request as rejected rather than continue bootstrap, allocation, or optimistic prefill as though it were valid.This builds on #36988, which already retires aborted post-forward prefill results. It addresses remaining admission handling and adds an independent transferred-length check before decode commits a handoff.
Modifications
Validation
Validated head
d1843270b8:Result: 38 passed, 14 subtests passed. Coverage includes actual
set_finish_with_abortmethods, pending admission errors, PP bootstrap ordering, metadata-only sends, legacy/fake metadata behavior, cached prefixes, rebootstrap, and resource cleanup.The collective regression uses two real Gloo processes with bounded timeouts, covering 13 scenarios and 24 steps per rank. Other runtime components are mocked, and the focused tests load production definitions through AST extraction to avoid GPU backend imports.
File-scoped
pre-commitandgit diff --checkpassed. Broader runtime-importing suites could not collect in the local environment because dependency setup remained incomplete, ending at a missingggufimport. No broad-suite pass is claimed.Accuracy Tests
No model accuracy or live GPU/RDMA tests were run. The intended behavior for valid handoffs is unchanged; rejected admission requests retain their original error, and inconsistent handoff lengths fail before decode commit.
Speed Tests and Profiling
Not benchmarked. The length guard adds a CPU attention-TP reduction per completed handoff. Decode HiCache adds a restore-state reduction when the transport poll is terminal. No performance improvement claim is made.
Scope and Remaining Limits
Checklist
CI States
Latest PR Test (Base): ❌ Run #34387680907
Latest PR Test (Extra): ❌ Run #34387680619
Latest PR Test (AMD ROCm 10): ❌ Run #34387680833