[Refactor] Choose every layer's boundaries from declarations and remove the scatter-mode selection

[ch-wan]

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

• [Refactor] Replace LayerScatterModes with LayerFacts and remove ScatterMode #41443
• [Refactor] Give a layer its CuTe DSL kernels at construction instead of a subclass #41442
• [Refactor] Build the boundary into any stage with one construction #41441
• [Refactor] Declare each stage's residual read and update, and give each stage its own entry #41440
• [Refactor] Split the layer communicator into a package (move only) #41439
• -> [Refactor] Choose every layer's boundaries from declarations and remove the scatter-mode selection #41438
• [Refactor] Step-3.5: complete the dense MLP's sum through ffn_exit #41437
• [Fix] LongCat-Flash under attention DP: branch and merge the dense FFNs through the communicators #41436
• [Refactor] Choose MoE layers' boundaries from declarations when moe_dp_size equals attn_cp_size #41435
• [Refactor] Falcon-H1: complete the FFN's sum through ffn_exit #41434
• [Fix] Falcon-H1: count the Mamba mixer's output once under tensor parallelism #41433
• [Fix] Run MoE layers under attention DP and GQA prefill CP on the declared DP × CP gather #41432
• [Refactor] Move CuTe DSL-fused layers onto the declared boundaries and FFN-exit kernel entries #41431
• [Refactor] Nemotron-H: build each layer's boundaries from its stage and the previous one #41430
• [Refactor] Build each decoder boundary from the declarations of its two sides #41429
• [Refactor] Publish the LoRA token layout from the batch's FFN input rows #41428
• [Refactor] Let the MoE declare whether its skipped reduction is one TP all-reduce #41427
• [Refactor] Run two-batch-overlap layers on the declared boundaries #41426
• [Refactor] Run MHC layers on the shared boundary steps with MHC's residual operations #41425
• [Refactor] Choose fully-DP dense and DSA / MLA prefill CP layers' steps from declarations #41424
• [Fix] Gather a dense FFN's input across attention DP and CP in one DP sum #41423
• [Fix] Keep one copy of CP-replicated rows in the DP gather #41422
• [Refactor] Choose a GQA prefill CP extend's steps from declarations #41421
• [Refactor] Run every batch of a layer from one BoundarySteps #41420
• [Refactor] Choose the LayerNorm SP region's and input-scattered batches' steps from declarations #41419
• [Refactor] Give the fused prepare_mlp kernels an explicit contract #41418
• [Refactor] Choose the boundaries of plain-TP dense layers and MoE layers from declarations #41417

Motivation

Two kinds of code still depended on scatter modes:

• Readers outside the communicator.
  • The NPU MLA / DSA attention and the DSA NPU indexer checked layer_input_mode == SCATTERED and attn_mode == TP_ATTN_FULL to know whether the layer input arrives on each attention-TP rank's slice.
  • ForwardBatch.moe_num_token_non_padded read sparse_mlp_scatter_mode().
• Configurations that _declared_sides still sent back to the scatter-mode steps. Some of them crash or hang there:
  • attention CP without --enable-prefill-cp;
  • attention DP × CP with moe_dp_size == attn_cp_size;
  • an a2a MoE under attention DP × GQA prefill CP.

Once both are gone, the selection from scatter modes (_select_boundaries_from_scatter_modes and the get_fn mode tables) has no caller left.

Modifications

1. Answer layout questions from the declarations

• LayerCommunicator.input_on_attention_tp_slices says whether the layer input arrives on each attention-TP rank's slice, from the declared input rows.
  • DeepseekV2AttentionMLA.forward / forward_prepare, forward_mha_prepare_npu, forward_mla_prepare_npu, forward_dsa_prepare_npu and the DSA NPU indexer take this bool instead of layer_scatter_modes.
  • DeepSeek-V2, GigaChat3.5, GLM-4-MoE-Lite and GLM5-Next pass it.
• moe_cp_gathers_sparse_moe_input(forward_batch) answers the MoE-CP question, and sparse_moe_gathers_over_moe_cp() is its construction-time half. moe_num_token_non_padded uses it together with is_moe_input_scattered_across_dp_ranks().

2. Build every attention CP layer's steps from the declarations

• _declared_sides never returns None, and _cp_on_declarations is removed. CP steps are built only when a prefill CP can shard a batch's tokens (_generic_prefill_cp_shards_tokens()).
• Attention CP without prefill CP runs like a prefill CP's batches that do not shard their tokens.
  • Without attention DP, an FFN may leave its sum as it does without CP.
  • Under attention DP, the FFN completes its own sum, because the DP reduce-scatter runs over the TP group, which spans the CP ranks.
  • On the scatter-mode steps this shape reduce-scattered across the CP ranks, and an idle rank skipped the DP gather the other ranks ran. It now runs.
• Refused at construction. _refuse_uncovered_cp_moe checks sparse layers and raises NotImplementedError naming the combination:
  • a MoE on the TP group with moe_dp_size == attn_cp_size under attention DP, with or without prefill CP. Its CP extend failed in the DP reduce-scatter.
  • a MoE dispatched per DP shard (a2a, the fp4 all-gather, DWDP) under attention DP and GQA prefill CP. It failed gathering over attention TP.
  • a MoE on the TP group with moe_dp_size == attn_cp_size under DSA or MLA prefill CP. Building its CP steps already failed.
• A non-first layer built without the previous layer's facts is refused. Every in-tree init_new call passes those facts.
• MHC: MHCLayerCommunicator refuses at construction, as the scatter-mode path did:
  • a MoE gathered over the MoE-CP group (moe_dp_size < attn_cp_size), whether or not prefill CP is on (the scatter-mode path refused it through MOE_FULL);
  • input-scattered attention under attention CP (the scatter-mode tables had no step for it).
• Input-scattered attention may run under attention CP without prefill CP: _input_can_be_scattered now checks for prefill CP instead of attn_cp_size == 1.

3. Remove the selection from scatter modes

• Removed:
  • _select_boundaries_from_scatter_modes, _post_init_communicate (including MHC's override) and _select_mlp_input;
  • the mode tables CommunicateSimpleFn.get_fn and CommunicateSummableTensorPairFn.get_fn;
  • mlp_input_kind / MlpInputKind, _mlp_input_order, _mlp_input_gather, _mlp_input_attn_tp_all_reduce and _complete_scattered_input;
  • scatter_mode_layouts, plus CommunicateContext.layouts, is_same_layout and is_same_group_size;
  • the scatter-mode fallbacks of input_rows and _branch_rows;
  • the FFN exit's case for a DSA / MLA CP extend on a layer without CP steps, which no layer is now.
• Kept:
  • LayerScatterModes, because the declarations still read the layer facts it carries;
  • the CommunicateSummableTensorPairFn moves that tbo_split_moves and the FFN output move use.

Tests

• test_declared_decoder_boundary.py adds tests for:

  • which CP configurations the steps cover;
  • CP without prefill CP;
  • a layer without the previous layer's facts being refused;
  • MHC refusing a MoE gathered over MoE-CP, with and without prefill CP.

  Its LoRA and TBO cases now read the declarations, and the cases that compared declarations against scatter modes are removed with the modes.

• test_communicator_layout.py is removed along with the selection it tested.

• test_num_token_non_padded_localization.py patches the new predicates.

Notes

• Newly refused at startup: the three combinations above. On main, the a2a MoE under attention DP × GQA prefill CP crashes at attention TP 2 (TP8 × DP2 × CP2 with DeepEP). At attention TP 1 (TP4 × DP2 × CP2 × EP4 with DeepEP) main runs, but its logprobs drift from a TP1 reference by up to 2.37 on a request that stays within 0.09 without CP. The refusal also covers the fp4 all-gather and DWDP, which were not run.
• Behavior fix: attention DP × CP without prefill CP, which hung on the scatter-mode steps, now runs.
• NPU callers: DeepseekV2AttentionMLA callers that never passed the scatter modes (Kimi-K3, Kimi-Linear, LongCat-Flash, Bailing MoE V3 / Linear, Hunyuan-V4 and others) now get input_on_attention_tp_slices=False. With SGLANG_USE_AG_AFTER_QLORA on NPU, they used to raise AttributeError on None. At attention TP 1 the NPU readers also skip the one-rank gather (a copy) they ran before.
• Removed without aliases: CommunicateSimpleFn.get_fn, CommunicateSummableTensorPairFn.get_fn, scatter_mode_layouts, mlp_input_kind and MlpInputKind.

Accuracy Tests

Commit 1: layout questions

• Call-sequence comparisons: all ten are identical to the parent.
• Unit tests: the related unit tests pass. The 126-file CPU unit suite has no new failures.
• Not run: the NPU paths. Their new predicate was compared with the old one by reading: they select the same layers, except for the one-rank gather at attention TP 1.

Commit 2: every CP layer on the declarations

• The parent's scatter-mode steps on H200 (Qwen3-30B-A3B), which are the basis for the refusals:

  • TP4 × CP2 without prefill CP runs.
  • TP8 × DP2 × CP2 with --moe-dp-size 2 crashes in dp_reduce_scatter_tensor.
  • TP8 × DP2 × CP2 with DeepEP crashes in the attention-TP all-gather, while TP8 × DP2 with DeepEP runs on the same machine.

• Call-sequence comparison through the real constructor, with two new axes: attention CP without prefill CP, and moe_dp_size == attn_cp_size. 96,096 cases differ from the parent, and every one falls into a known class:

  • refusals: 34,560;
  • under attention DP without prefill CP, the FFN completes its own sum and returns over DP, where the parent reduce-scattered over the TP group across CP: 21,888;
  • empty batches (a collective on 0 rows, or another 0-row tensor as the residual): 17,280;
  • an idle DP rank joining the other DP ranks' gather: 8,064;
  • reductions over one rank at attention TP 1: 7,200;
  • dense layers in a model that is refused at its sparse layers: 5,376;
  • a dense MLP on every rank, with no sum owed: 1,728.

  The other nine comparisons are identical.

• MHC comparison with CP axes (off, without prefill CP, prefill CP) and moe_dp_size, 74,880 cases: 24,874 differ, and every one falls into a known class:

  • both trees refuse at construction, with different messages: 11,434;
  • the same refusals as for any layer: 9,216;
  • input-scattered attention under attention CP without prefill CP, which the parent refused: 2,784. This PR refuses them too (the MHC comparison at the stack's tip, before and after that refusal, differs in exactly these cases: 4,032 of 74,880, all now refusals);
  • no reduce-scatter across CP under attention DP without prefill CP: 1,440.

• Mutants and unit tests: 11 mutants are caught. The 126-file CPU unit suite has no new failures.

• GPU (H200 × 8): Qwen3-30B-A3B bf16, fa3, zigzag, with a runtime audit on each run. These runs were made before the MHC refusal was folded in; that part only touches MHC construction and splits one predicate in two.

  • TP4 × CP2 without prefill CP: bit-identical to the parent (4 / 4). All 48 layers take the declared entry; on the parent, none did.
  • TP8 × DP2 × CP2 without prefill CP: the parent hangs and is killed by the watchdog. This PR runs. Against TP1, the first differing tokens are [0, 0, 9, 16] and the max prefill top-5 |Δlogprob| is [0.518, 0.515, 0.08, 0.131], within the range of supported shapes.
  • TP8 × DP2 × CP2 with --moe-dp-size 2, and with DeepEP: startup raises NotImplementedError naming the combination.

Commit 3: selection removed

• Call-sequence comparisons: the six that still apply are identical to the parent:

  • declared boundaries (575,040 cases);
  • MHC (74,880);
  • CuTe DSL boundaries (89,088);
  • FFN exit selection (92,160);
  • Nemotron-H stages (9,664);
  • CuTe DSL exit (64).

  Five comparisons that built communicators from the removed mode tables are retired with them.

• Unit tests: they pass. The 125-file CPU unit suite has no new failures.

• GPU (H200 × 8): parent vs this PR, requests one at a time, bit-identical, 4 / 4 each:

  • Qwen3-30B-A3B: GQA prefill CP TP4 × CP2, TP4 × CP2 without prefill CP, and TP8 × DP2 × CP2;
  • GLM-5.2 (5 layers), DSA interleave prefill CP 8, with 4 attention-CP gathers and reduce-scatters per rank;
  • DeepSeek-V3 (4 layers), MLA prefill CP 4.

  On TP8 × DP2 × CP2, one run of this PR first differed on 2 of 4 requests. Five parent runs and three more runs of this PR were all bit-identical.

Checked at the stack's tip against main (H200 × 8 and B200 × 4, greedy, requests one at a time)

• A dense-only model under attention DP × prefill CP with moe_dp_size == attn_cp_size (Qwen3-0.6B TP8 × DP2 × CP2, --moe-dp-size 2): main raises in the DP gather; the stack matches a TP1 reference's text on all four requests (max |Δlogprob| 0.08).
• DSA attention CP without prefill CP (5-layer GLM-5.2): at attention CP 8 (attention TP 1), main and the stack differ only on the long request, by no more than main differs from itself between runs. At attention CP 4 (attention TP 2), main raises; the stack's distance to a run without CP matches main's attention CP 8 run.
• MHC with input-scattered attention under attention CP (GLM-5.3-Flash TP4, attention CP 2): main and the stack both refuse at startup.

Not covered

• The fp4 all-gather and DWDP under attention DP × GQA prefill CP (refused).
• Input-scattered attention under attention CP for a model without MHC fails at runtime on main and on the stack alike (ragged prefill shape check); not changed here.
• NPU paths.

The whole stack was also compared with main at its tip; the results are in #41443.

Speed Tests and Profiling

No kernel changes. Collectives change only in two places:

• attention DP × CP without prefill CP, where the FFN now completes its own sum (the parent hung);
• NPU at attention TP 1, which skips the one-rank gather.

DeepSeek-family layers now read input_on_attention_tp_slices, a set lookup, once per layer per forward.

Decode latency, time to first token and memory for the whole stack against main are in #41443.

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 #36360236998
Latest PR Test (Extra): 🚫 Run #36360236669
Latest PR Test (AMD ROCm 10): 🚫 Run #36360236978