[Fix] Keep residual rows and sums consistent across appends and pipeline handoffs

[ch-wan]

This PR is part of a stack (oldest at bottom):

• [Refactor] Declare every model's stages from a shared function and drop the meta-device neighbour build #43460
• [Refactor] Declare pipeline neighbours from a shared declaration function #43459
• [Refactor] Decide an exit's fixed facts when the stage is built #43458
• [Refactor] Bind a layer stack in one pass over each stage's data flow #43457
• [Refactor] Decide input-scattered eligibility and stage rejections in one place #43456
• [Refactor] Make residual read and update capabilities explicit and checked #43455
• -> [Fix] Keep residual rows and sums consistent across appends and pipeline handoffs #43454
• [Refactor] Record what a bound stage path does instead of reading it back #43453

Motivation

Five defects in how a stage boundary places the residual and completes sums, each reachable on main:

1. After an append that ends on an attention which always leaves its sum, the next append rebuilt that attention's incoming boundary from a declaration without its history, took the attention as the start of the stack, and sliced a residual that was already on this rank's attention-TP slice a second time. ZAYA1's one-stage appends hit this with an a2a MoE and TP > 1: a size mismatch, or on a small batch wrong values without an error.
2. A stack (or a pipeline rank) ending on such attentions after an FFN that runs on its own rows (a2a MoE, or a dense FFN fully data-parallel) left the residual on the slice, while the next rank and the final norm expect the attention's rows.
3. DeepSeek on a later pipeline rank read the batch size and device off the owed-sum handle it received on an input-scattered batch: AttributeError on every later rank.
4. The FFN before a pipeline handoff resolved the rows its residual is on while it runs with two different rules, one for its entry and one for its exit. With input-scattered attention the handoff sent the attention-TP slice where the next rank expects every row; with attention DP and a dense FFN fully data-parallel, binding failed with NotImplementedError when a rank's last FFN followed such a dense layer.
5. to_pp() sent a declared partial sum unreduced for the next rank to complete. The pipeline sends one slice per attention-TP rank and gathers the slices on the receiver, so the receiver reassembled a mixture of different ranks' partials. Nemotron-H cut between a mixer and the FFN after it, and DeepSeek with input-scattered attention and a reduce-scatter boundary reduction, produced wrong outputs under PP with attention TP > 1; with an all-reduce boundary reduction, input-scattered attention summed the received output a second time.

Modifications

• Bind an append after an attention with the residual rows it ran on. The stack keeps the boundary each pending stage is bound with and hands it to the stage after it. The attention-after-attention step that keeps the residual on the slice reduce-scatters over the full TP group, so binding rejects it under attention DP instead of summing across ranks that hold other sequences.
• Return the residual on the attention's rows before a stack's last attention. That FFN returns the residual on the attention's rows, as an FFN does when it ends the stack.
• Read a DeepSeek pipeline rank's batch from positions, not hidden states.
• Keep an FFN's residual on one set of rows across a pipeline handoff. The entry rewrite is removed; entry and exit use the FFN's own resolution.
• Complete every sum an output owes before a pipeline handoff. to_pp() completes every sum before transport and from_pp() declares none; the receiving stage takes the path it already binds for an input whose sum is complete.
• Test that a pipeline cut keeps a stack's rows and values. For stacks of two to four stages (an attention that always leaves its sum, a mixer whose exit may leave its sum to an FFN, a dense FFN, a sparse FFN), cut at every stage boundary under four parallel configurations and two boundary reductions (6,384 cuts): within each rank both sides of every boundary agree on the residual's rows, the handoff carries the residual on the attention's rows, the stages away from the cut bind as they do uncut, and a cut rejects only what the stack rejects uncut. A second test runs the sender, to_pp(), a transport that sends one slice per attention-TP rank and gathers them back, from_pp() and the receiver on simulated attention-TP ranks, and checks the values against the uncut run, that every owed sum completes once and that the receiver sums nothing.

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:

• Each fix has a regression test that fails on its parent commit.
• (1) and (2): running ZAYA1-shaped stacks (attention, a2a MoE, attention, FFN; and attention, FFN, attention with a final norm) on simulated attention-TP ranks with fake collectives, the parent fails with a size mismatch or gives wrong values on small batches; this PR matches a per-token reference exactly.
• (3)-(5) on H200, PP2 against PP1 of the same tree (DeepSeek on this PR's tree; the other two on the stack's head, whose PP path is this PR's):
  • DeepSeek-V2-Lite (synthetic 4 layers, q_lora), TP2, input-scattered attention, rs with 2+2 and 1+3 layer splits and ar with 2+2: PP2 is bit-identical to PP1 (before: crash at the receiver, then wrong outputs, up to 2.4 in prompt logprob).
  • Nemotron-H (6 layers), TP2, cut between a mixer and its FFN: the PP2 difference drops from 0.93 / 2.21 (prefill logprob, first token diverging) to 0.011 / 0.016, the level of the FFN-to-mixer cut that never carried a partial.
  • TP4, attention DP2, dense FFN fully data-parallel, PP2: the 2+2 split is bit-identical to PP1. The 1+3 split differs by 0.098 because the dense FFN's output is written into the residual before the handoff (bf16 add, then norm on the next rank) instead of a fused add+norm; with the fused kernel replaced by add-then-norm in both runs, PP2 and PP1 are bit-identical. Main shows the same 0.098 on the 1+3 split: its PP2 output is bit-identical to this PR's (the 2+2 split fails to construct on main).
• The plan comparison differs from the parent only in the configurations each fix names.

Speed Tests and Profiling

The sum a PP rank used to leave to the next rank now completes before the handoff (one all-reduce or reduce-scatter over attention TP on the sending rank instead of one on the receiving rank). No other path changes.

Measured on the merge candidate against main (H200, dummy weights, decode with CUDA graph, three runs per tree, a difference counted only if it exceeds the same-tree spread): Qwen3-8B at TP1 and TP2, DeepSeek-V2-Lite at TP2, and DeepSeek-V2-Lite at TP4 with attention DP2, the dense FFN fully data-parallel and PP2 (1+26 layers) show no difference in prefill throughput (4k tokens, batch 1 and 8), decode throughput (batch 1, 16, 64), peak allocated / reserved memory per rank, or steady-state memory. DeepSeek-V2-Lite with input-scattered attention at TP2 PP2 runs only on this stack (main crashes), so it has no baseline.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): 🚫 Run #38080746857
Latest PR Test (Extra): 🚫 Run #38080746610
Latest PR Test (AMD ROCm 10): ❌ Run #38080746826

[mintlify]

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

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

💡 Tip: Enable Automations to automatically generate PRs for you.