[Fix] Step-3.5: build the shared expert without TP under all-to-all MoE backends

[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

Step-3.5 adds the outputs of its routed experts and its shared expert, then runs one TP all-reduce on the sum. That is correct when both parts are TP partial sums. Under DeepEP, fuseep and the other all-to-all MoE backends, the combine already returns each token's routed output in full, on the rank that owns the token, but the shared expert stayed TP-sharded:

• The all-reduce summed the complete routed output again, so it was counted TP times.
• The shared expert's partial sums were computed on the tokens each rank holds, so the all-reduce added outputs of different tokens.

DeepSeek-V2 and GLM-4-MoE avoid both by building the shared expert with TP size 1 under these backends.

Modifications

• Build share_expert with tp_size=1 under the backends GLM-4-MoE uses for this: DeepEP, Mooncake, NIXL, MORI, Ascend fuseep, the FlashInfer and FlashInfer MegaMoE all-to-all, and the FlashInfer CUTLASS FP4 all-gather.
• Step3p5MoEMLP takes the shared expert's output and adds it in place, without allocating a third tensor: into the routed output before reduce_moe_output() on the TP path; on the DeepEP path, where nothing is left to sum, the combined routed output is added into the shared expert's output.
• The decoder no longer forces the MoE block to skip its reduction and then all-reduces by hand. The FFN exit's flags reach reduce_moe_output() as for the other MoE models.

Accuracy Tests

H200, stepfun-ai/Step-3.5-Flash cut to its first five layers (three dense, two MoE):

• --tp-size 2, python -m sglang.benchmark.one_batch --correctness-test: identical to the parent commit, and to a second run of the parent.
• --tp-size 4, greedy sglang.Engine runs with per-step logits: all 99 steps bitwise identical to the parent.
• --tp-size 4 --ep-size 4 --moe-a2a-backend deepep --moe-runner-backend deep_gemm, against the TP4 run as the reference. The DeepEP forward does not apply moe_router_scaling_factor (3.0 here), which this PR does not change, so the comparison also applies it with a test hook in both trees:
  • This PR: max |Δlogit| 0.17 (logits up to 9.2), all 99 argmax steps agree.
  • Parent: max |Δlogit| 8.09, and the first generated token differs for all three prompts.
  • This PR without the hook: max |Δlogit| 1.36, all 99 argmax steps agree; the remaining difference is the missing scaling factor.
• 109 affected unit test files (2,697 tests) pass at the top of this stack, the same as on main.

Speed Tests and Profiling

Without an all-to-all backend the same kernels and collectives run. Under the all-to-all backends each rank now runs the whole shared expert on its own tokens instead of a TP shard of it on the same tokens, as in DeepSeek-V2 and GLM-4-MoE; 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 #36667714896
Latest PR Test (Extra): 🚫 Run #36667714664
Latest PR Test (AMD ROCm 10): ❌ Run #36667714975

[ch-wan]

/tag-and-rerun-ci

[mintlify]

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

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Sep 30, 2026, 4:14 AM

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