Skip to content

[Refactor] Bind a layer stack in one pass over each stage's data flow - #43457

Merged
ch-wan merged 3 commits into
mainfrom
cheng/refactor/stage-binding-one-pass
Oct 10, 2026
Merged

ch-wan merged 3 commits into
mainfrom
cheng/refactor/stage-binding-one-pass

Conversation

@ch-wan

@ch-wan ch-wan commented Oct 10, 2026 •

Copy link
Copy Markdown
Collaborator

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

  • Bind a layer stack's stages in one pass over its connections. The stack still binds when it closes, but chains its stages in append order and builds each connection once, as one stage's outgoing and the next one's incoming connection; 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 its consumer's input rows. Each stage records its producer without the producer's history.
  • Resolve each stage's residual data flow once per batch variant. A stage is resolved once per variant, from what arrives at its entry and the stage after it, into one record: input and output contracts, the arrival, the rows the residual is on while the stage runs, and the rows its exit returns it on. Both edges of a stage are read from that record.
  • Leave out what a stage's entry does not read of its producer's output. The arrival no longer copies the producer's sum group and carry permission when nothing is owed.

Accuracy Tests

How the stack is checked:

  • Unit tests: the 39 layer-boundary test files pass at the head of every PR in the stack, apart from the two AMD-only test_aiter_allreduce_fusion_amd cases that fail on NVIDIA on main too; the test files that changed most recently pass on every commit. The full test/registered/unit suite (926 files, one process per file) shows no failure on the stack's head that main does not have.
  • Plan comparison: an out-of-tree harness binds every stage plan for seven stack shapes (decoder layers with and without an output transform, MHC, Nemotron-H patterns, Kimi K3, Qwen4-exp, NemotronNAS-style attention-only layers, single-stage appends) over 23 parallel configurations, four boundary reductions and one or two pipeline ranks (3,928 scenarios), and compares each bound plan leaf by leaf with the parent commit's.
  • GPU: dummy weights are seeded per TP rank and scaled (0.05, norms 1.0), so that a sum taken twice or a mixture of rank partials changes the output; the default dummy weights give every rank the same shard and would hide both.
  • GPU results below were measured on these commits as based on main 516fd1c1b5 unless marked as on the merge candidate; rebasing onto current main had no conflicts, and the CPU checks were run again after it.
  • On the merge candidate (head of this stack against main 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

@ch-wan
ch-wan force-pushed the cheng/refactor/stage-binding-one-pass branch from 0cbefd3 to 21ee9e6 Compare October 10, 2026 03:23
@ch-wan
ch-wan force-pushed the cheng/refactor/stage-construction-checks branch 2 times, most recently from 9a51db6 to 25ca41c Compare October 10, 2026 06:21
@ch-wan
ch-wan force-pushed the cheng/refactor/stage-binding-one-pass branch from 21ee9e6 to 3617891 Compare October 10, 2026 06:21
@ch-wan
ch-wan force-pushed the cheng/refactor/stage-construction-checks branch from 25ca41c to d56382c Compare October 10, 2026 19:41
Base automatically changed from cheng/refactor/stage-construction-checks to main October 10, 2026 19:41
@mintlify

mintlify Bot commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

Preview deployment for your docs. Learn more about Mintlify Previews.

Project Status Preview Updated
sglang-doc 🟢 Ready View Preview Oct 10, 2026, 7:44 PM

💡 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
ch-wan force-pushed the cheng/refactor/stage-binding-one-pass branch from 3617891 to 1529fca Compare October 10, 2026 19:41
@ch-wan
ch-wan merged commit a956aa1 into main Oct 10, 2026
20 of 25 checks passed
@ch-wan
ch-wan deleted the cheng/refactor/stage-binding-one-pass branch October 10, 2026 19:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant