Repository navigation
[Refactor] Bind a layer stack in one pass over each stage's data flow - #43457
Merged
Merged
Conversation
ch-wan
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Ying1123,
ispobock and
merrymercy
as code owners
October 10, 2026 00:47
This was referenced Oct 10, 2026
ch-wan
force-pushed
the
cheng/refactor/stage-binding-one-pass
branch
from
October 10, 2026 03:23
0cbefd3 to
21ee9e6
Compare
ch-wan
force-pushed
the
cheng/refactor/stage-construction-checks
branch
2 times, most recently
from
October 10, 2026 06:21
9a51db6 to
25ca41c
Compare
ch-wan
force-pushed
the
cheng/refactor/stage-binding-one-pass
branch
from
October 10, 2026 06:21
21ee9e6 to
3617891
Compare
ch-wan
force-pushed
the
cheng/refactor/stage-construction-checks
branch
from
October 10, 2026 19:41
25ca41c to
d56382c
Compare
ch-wan
requested review from
JustinTong0323,
fzyzcjy,
sogalin,
wisclmy0611 and
zijiexia
as code owners
October 10, 2026 19:41
Base automatically changed from
cheng/refactor/stage-construction-checks
to
main
October 10, 2026 19:41
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Automations to automatically generate PRs for you. |
A stack bound its appends one at a time. The last stage of each append waited for the next append's first stage, and every connection between two appends was built twice: once as the waiting stage's outgoing connection, and again as the next stage's incoming one, rebuilt by looking back through the producer's declaration. The waiting stage's incoming connection was rebuilt the same way. Capture across an attention was carried by replacing the attention's bound paths once its FFN had bound, and an attention's view of its consumer's input rows was set on its plan after the plan was built. The stack now binds when it closes, as before, but in one pass: its stages are chained in append order, each connection between two of them is built once, as one stage's outgoing and the next one's incoming connection, and a branch starts from the incoming connection of the stage it branches from. An attention that always leaves its sum binds after the FFN that follows it in its append, so its plan is built with that FFN's capture check, and with its consumer's input rows. Each stage records its producer without that producer's history.
A connection resolved both of its stages for every variant: the producer with its consumer, for the exit, and the consumer again on its own, for the entry. So every stage was resolved twice per variant. An attention's resolution also left out the rows its residual is on while it runs; each connection out of it found them in the connection into it, and a connection without one fell back to the attention's rows. A stage is now resolved once per variant, given what arrives at its entry and the stage after it, into one record: its input and output contracts, what arrived, the rows its residual is on while it runs, and the rows its exit returns it on. An attention's rows come from what arrives at it there. The edge into a stage and the edge out of it are both read from that record, and the connections along a chain of stages are built from the records in one pass. A stage on another rank is resolved the same way, from the stack input as what reaches it.
What arrives at a stage after an FFN that exits on the attention's rows, or after a mixer, copied the producer's sum group and whether it may carry its sum to the next stage. An entry reads the group only of a sum it owes, and never reads whether the producer may carry one: a sum carried on arrives with the value. Only a mixer's declared sum is owed across such an edge, so what arrives now says just that, and the tests that checked the copy check the producer's output and the sum owed.
ch-wan
force-pushed
the
cheng/refactor/stage-binding-one-pass
branch
from
October 10, 2026 19:41
3617891 to
1529fca
Compare
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.
This PR is part of a stack (oldest at bottom):
Motivation
A stack bound its appends one at a time. The last stage of each append waited for the next append's first stage, and every connection between two appends was built twice (once as the waiting stage's outgoing connection, once as the next stage's incoming one, rebuilt by looking back through declaration history). Each stage was resolved twice per batch variant, and an attention's resolution left out the rows its residual is on while it runs, which each connection had to find elsewhere. That double resolution is where the two disagreeing rules fixed in #43454 came from.
Modifications
Accuracy Tests
How the stack is checked:
test_aiter_allreduce_fusion_amdcases that fail on NVIDIA on main too; the test files that changed most recently pass on every commit. The fulltest/registered/unitsuite (926 files, one process per file) shows no failure on the stack's head that main does not have.516fd1c1b5unless marked as on the merge candidate; rebasing onto current main had no conflicts, and the CPU checks were run again after it.961bcf481d, H200, asymmetric dummy weights unless noted): DeepSeek-V2-Lite input-scattered PP2 is bit-identical to PP1 (main crashes on it); Nemotron-H cut between a mixer and its FFN differs from PP1 only at the rounding level of the FFN-to-mixer cuts (main: 0.93 and 2.21); Llama, Qwen3-MoE, GPT-OSS and DeepSeek-V2-Lite are bit-identical to main at PP1 and PP2; GSM8K (real DeepSeek-V2-Lite weights, 250 questions, TP4 with attention DP2 and the dense FFN fully data-parallel) gives 0.360 at PP2, 0.364 at PP1 and 0.368 on main at PP1, medians of three runs whose run-to-run spread is 0.012–0.016. The last revision of the stack changed only comments, docstrings, an equivalent condition and explicitly set attributes after these runs; its bound plans are identical.For this PR: the plan comparison shows no difference. Where a configuration is rejected for more than one reason, the message may now name another of them (all stages resolve before any binds); which scenarios are rejected does not change. The third commit differs only in arrival fields of arrivals that owe nothing, which no entry reads. Counting resolutions per stage per variant: 3.75 (decoder layers) and 6.62 (single-stage appends) on main, 1.00 here.
Speed Tests and Profiling
Construction resolves each stage once per variant instead of several times; no forward-path change.
Checklist
CI States
Latest PR Test (Base): 🚫 Run #38080821606
Latest PR Test (Extra): 🚫 Run #38080821429
Latest PR Test (AMD ROCm 10): ❌ Run #38080821654