[Refactor] Choose fully-DP dense and DSA / MLA prefill CP layers' steps from declarations

[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

• A dense MLP on every rank (moe_dense_tp_size=1) runs on each rank's own slice of the tokens and owes no sum, like a MoE dispatched per DP shard. A global exclusion still kept every layer of such a model on the scatter-mode steps.
• DSA and MLA prefill CP ran DSACPLayerCommunicator, a subclass with its own three tables.
  • Under them attention TP is one and the dense MLP runs on every rank, so most layers only normalize.
  • A MoE on the TP group (DSA's interleave split without an a2a backend) gathers a CP extend's shards over the attention-CP group and reduce-scatters its output there. That reduce-scatter also completes the sum the MoE leaves.
• Which kind of prefill CP ran was asked separately by the FFN input selector, the output selector and the FFN exit. The FFN exit let a CP extend's MoE skip its sum because the batch ran the CP steps. The reduce-scatter was chosen without checking which group the MoE owes its sum over.

Modifications

• Dense layers on every rank take the declarations, with the FFN on the local rows. The layer after one takes those rows as its input. The steps are the ones the scatter modes chose.

  • Under two-batch overlap with attention TP > 1, a dense layer before a sparse one gathers its output for the split, so those layers keep the scatter-mode steps.
  • Input-scattered attention never runs with such an MLP (_input_can_be_scattered excludes it), so these layers build no input-scattered steps.

• DSA and MLA prefill CP take the declarations, as GQA prefill CP does.

  • A CP extend (_batch_shards_over_cp) runs the layer's CP steps; other batches run its ordinary ones.
  • Where the FFN is on the TP group, the CP steps gather over the attention-CP group in equal shards (_mlp_input_gather_attention_cp). The FFN output declares that it leaves its sum to the reduce-scatter that takes each rank's shard back (_reduce_scatter_over_cp).
  • Layers whose FFN runs on the rank's own rows only normalize.
  • Also declared under DSA and MLA CP: a MoE under attention DP (it dispatches with a2a), and a dense layer under two-batch overlap (with attention TP 1 its gather for the split moves nothing).

• Removed: DSACPLayerCommunicator and its tables (DSACPCommunicateSimpleFn, DSACPCommunicateWithAllReduceAndLayerNormFn, DSACPCommunicateSummableTensorPairFn).

  • The collectives and the KV prefetch stay in communicator_dsa_cp.
  • DeepseekV2DecoderLayer builds the base communicator under prefill CP.
  • A DSA or MLA prefill CP layer that the declarations do not cover fails at construction, since the scatter-mode steps have no attention-CP gather.
  • Comments in model_overrides/deepseek_v2.py, model_runner.py and decode_cuda_graph_runner.py that named the subclass are updated.

• One CpMoves, chosen at construction for the kind of prefill CP (_cp_moves()). It holds:

  • the gather after each rank completes its block;
  • the take-back of a complete output;
  • under DSA and MLA CP, the reduce-scatter over the attention-CP group and its group.

  GQA prefill CP keeps its padded MoE-CP gather and completes the sum within the layer.

• The output selector hands the FFN's sum to the reduce-scatter only when the sum is owed over the same ranks (_same_ranks), and rejects the layer otherwise. A complete output is only taken back, as its equal-length shard (_take_back_attention_cp_shard). BoundarySteps.ffn_output_move_completes_sum records whether the output move completes the sum, and the FFN exit publishes mlp_reduce_scatter from it.

• Notes:

  • On a CP extend, a layer whose FFN runs on the rank's own rows no longer publishes mlp_reduce_scatter, and should_use_reduce_scatter answers False for it. It owes no sum, so no reader behaves differently: should_skip_mlp_all_reduce (dense TP 1) and should_add_replicated_moe_output (the a2a combine is complete).
  • Newly rejected at construction: DSA or MLA prefill CP with a MoE on the TP group whose sum is owed over other ranks than the attention-CP group, for example DSA interleave with --moe-dp-size equal to the attention CP size. On main the subclass reduce-scattered over attention CP an output the MoE had already completed.
  • A DeepSeek-family model with --enable-prefill-cp but neither DSA nor MLA CP used to build the subclass, which failed at construction. It now builds the base communicator and takes the GQA prefill CP steps. No GPU run covers it.
  • The two-batch-overlap exclusion is narrowed to attention TP > 1. So DP attention with attention TP 1, DeepEP, two-batch overlap and moe_dense_tp_size=1 now runs every layer on the declared path. With attention TP 1 the gather for the split moves nothing, and the split still reads the scatter modes. No GPU run covers this configuration at this PR.

• Tests:

  • new test_declared_attention_cp.py: two CP ranks run prepare_mlp and the FFN exit rank by rank. The reduce-scatter completes the sum the MoE leaves, and each rank gets its shard's complete value. A complete output is only taken back, with no sum and no mlp_reduce_scatter. A MoE that owes its sum over other ranks is rejected at construction.
  • test_declared_decoder_boundary.py adds dense layers on every rank, and DSA CP: a CP extend leaves its sum to the reduce-scatter, DSA CP asks its own predicates for a CP extend, and its dense layers run on their shard. Its planning helper publishes an exec config with overlap off.

Accuracy Tests

No behavior change is intended, apart from the construction-time changes in Notes.

• Dense MLP on every rank:

  • call sequences through the real constructor, with a fully-DP dense axis: 218,112 cases. All are identical or differ only in the entry or a probe, except 480. Those are sparse layers under GQA prefill CP with attention TP 1 that now take the declared path: the old path ran an all-reduce over the one-rank attention-TP group before the norm, and the new one only normalizes. The values are the same.
  • the seven other comparisons are identical, and the related unit files pass. Two mutants are caught.
  • GPU (B200 × 4), DeepSeek-V2-Lite with moe_dense_tp_size=1 and DP attention: TP2 × DP2, TP4 × DP2 and TP4 × DP4 are bit-identical to the parent. A runtime audit in TP4 × DP2 counted 0 / 27 declared layers per rank on the parent, 27 / 27 on the change.
  • DeepEP TP4 × DP2 × EP4 fails in prefill on both trees, in main's MLA ragged prefill (q.shape[0] (10) does not match qo_indptr[-1] (9)).

• DSA and MLA CP on the declarations:

  • call sequences, with a CP-flavor axis: 219,264 cases. All are identical or differ only in the entry or a probe, except 480. Those are DSA CP extends whose FFN runs on the rank's own rows, and they differ only in the mlp_reduce_scatter / should_use_reduce_scatter value (see Notes).
  • The DSA interleave TP-group MoE (gather plus reduce-scatter) makes the same calls as the removed subclass.
  • the seven other comparisons are identical, and 174 related unit files pass. Five mutants are caught.

• CpMoves:

  • all eight call-sequence comparisons are identical to the parent, including all 219,264 constructor cases. The reachable configurations run the functions they ran before.
  • 176 related unit files pass. Three mutants are caught. A fourth, marking the complete-output take-back as completing the sum, is equivalent: that flag is never read on that path.

• GPU (H200 × 8), removed subclass vs declarations, requests sent one at a time. Long prompts are CP extends and short ones are not; all go through the decode CUDA graph.

  Configuration	Result
  MLA CP, 4-layer DeepSeek-V3, attention CP 4 (tp 4), zigzag, DeepEP + DeepGEMM	bit-identical
  MLA CP, same, attention CP 2 × DP 2 (tp 4, dp 2)	bit-identical
  DSA interleave CP, 5-layer GLM-5.2 (bf16, no a2a), attention CP 8	bit-identical
  DSA zigzag CP, same model with DeepEP	rejected by both trees: zigzag strategy does not support DSA indexer key gather (as on main)

  A runtime audit on the change:

  • MLA CP: 4 / 4 declared layers per rank, and 8 prepare_mlp calls on CP extends.
  • DSA interleave: 5 / 5 declared layers per rank, and 10 prepare_mlp calls on CP extends. The attention-CP gather and the reduce-scatter each ran 4 times, i.e. 2 MoE layers × 2 long prompts.

• These results were measured before the stack was rebased onto main after [Refactor] Choose a dense layer's boundaries under attention DP from both sides' declarations #41257.

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

Speed Tests and Profiling

No kernel or collective changes for the configurations that ran before.

On a DSA or MLA CP extend, each read of the batch's steps evaluates the DSA / MLA CP predicates (is_dsa_enable_prefill_cp() / is_mla_cp_enabled(), then dsa_use_prefill_cp / is_mla_cp_active). They are not cached, and there are about ten such reads per layer. Other batches, decode included, return before them. This was not measured separately.

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 #36359958185
Latest PR Test (Extra): 🚫 Run #36359958091
Latest PR Test (AMD ROCm 10): ⏳ Run #36359958200