[Fix] Release consumed residual contributions

[ch-wan]

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

• [Refactor] Build seven more decoders from stage boundaries #41779
• [Refactor] Simplify layer boundary internals #41778
• [Fix] Step-3.5: build the shared expert without TP under all-to-all MoE backends #41754
• [Refactor] Send producer-written residual streams through to_pp #41753
• [Fix] Reduce the Sarvam dense MLP through its row-parallel projection #41752
• [Refactor] Leave decoder stacks through final_norm(skip_empty=True) #41751
• [Refactor] Rename layer boundary protocols, contracts and helpers #41750
• -> [Fix] Release consumed residual contributions #41749

Motivation

Under attention DP, prefill CUDA graph capture uses about 62 MiB more memory than before the layer boundary refactor (#41547–#41557). The FFN all-reduce of a dense layer is now deferred to the next layer's input, so the unreduced output — gathered over the DP ranks, 2N rows for N local tokens — is stored in the residual stream's pending contribution. The next stage's prepare() consumes it, but the model loop still holds the returned handle until the next layer returns, and the handle keeps the tensor alive. Inside the graph memory pool that block cannot be reused.

Modifications

• Contribution.release() drops a contribution's tensors.
• ResidualStream.write() releases the contribution it replaces. It is called only where a pending contribution is consumed: a successful prepare(), fold() and branch merge. complete_output(), snapshot() and export (to_pp(), the final norm) do not release. Stale handles were already rejected by the stream, so nothing can read a released contribution.
• Test: a consumed partial is freed while the caller still holds its handle. It fails on the parent commit.

Accuracy Tests

B200, --tp-size 2 --dp-size 2 --enable-dp-attention, python -m sglang.benchmark.one_batch --correctness-test: the printed prefill logits and all three generations are identical to the parent commit for Qwen3-0.6B and for Qwen1.5-MoE-A2.7B.

Speed Tests and Profiling

Prefill CUDA graph capture memory, same B200 configuration (58 capture shapes, 4–8192 tokens):

Tree	Capture memory
Before the refactor (6582425854)	1124.0 MiB
Parent of this PR (f731e82f09)	1186.0 MiB
This PR	1124.0 MiB

Qwen1.5-MoE-A2.7B, same flags: 2.51 GB on the parent (the same on a repeat run), 2.26 GB with this PR, 2.26 GB before the refactor.

For Qwen3-0.6B the graph pool size after every shape equals the pre-refactor tree. Kernels and collectives are unchanged, so no latency run was made.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #36623301447
Latest PR Test (Extra): ❌ Run #36623301195
Latest PR Test (AMD ROCm 10): ❌ Run #36623301173

[ch-wan]

/rerun-test test_residual_stream.py

[github-actions]

Results for /rerun-test test_residual_stream.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layer_boundary/test_residual_stream.py