[Refactor] Nemotron-H: build each layer's boundaries from its stage and the previous one

[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 Nemotron-H layer is a single stage: a Mamba / attention mixer, or an MLP / MoE FFN. Which stage completed a reduction was decided in the model, from the layer types on both sides.

• input_is_reduced, feeds_mlp_layer and skip_reduce carried that decision.
• An FFN whose input was already complete folded it into the residual in bf16, then handed prepare_mlp a zero update to reduce.
• The aux hidden-state snapshot reduced a copy after its own pattern check.
• The MTP MoE layer was marked as the last layer after construction.

Modifications

• nemotron_h_utils declares each stage (StageDecl):
  • a mixer computes on the attention's rows. It owes its attention-TP sum to an FFN stage after it always, or to a mixer after it when the fused kernel takes it;
  • an FFN computes on the TP group's rows. A MoE dispatched by an a2a backend computes on this rank's own rows, as the scatter modes had it. An FFN may leave its sum to a mixer after it.
• New boundary_layout.stage_edges builds a layer's two boundaries from its own declaration and the previous stage's output. LayerCommunicator(stage=LayerStage(...)) runs them through make_boundary. make_layer_communicator(layer_norm, *, pattern, layer_idx) builds it from the pattern.
• New LayerCommunicator.mixer_exit, the mixer's counterpart of ffn_exit. It decides once, from what the stage's output declares, whether the mixer skips its all-reduce (MixerExit.skips_reduction), and finish hands on the output. The FFN stages run through ffn_exit.
• An FFN whose input is complete adds it and normalizes. At the start of the layer stack it reads its input with no residual: AddAndNorm.update_and_read_ffn_input accepts residual=None. An FFN on each rank's own rows (a MoE dispatched by an a2a backend) first takes this rank's slice of the complete input and of the residual (_mlp_input_slice).
• Aux snapshot: the layer loop's snapshot reads from the boundary's declaration whether the value there owes an attention-TP sum (_owes_attention_sum_at).
• MTP: the MTP MoE layer takes its last-layer fact from the MTP pattern it is built with. The assignment after construction is removed.
• Removed from the model: input_is_reduced, the feeds_mlp_layer attribute and helper, the skip_reduce decision, the zero tensor, and the nn.Identity() placeholders. A single-stage layer passes None for the norm it does not use.
• token_axis_sizes becomes public for the model's declarations.
• The guard that no FFN behind ffn_exit reduces over attention TP now also checks subclasses of ffn_exit users, and the FFN a stage holds in mixer. Before, it did not see Nemotron-H.
• Notes:
  • Same reductions on real transitions. On the transitions real checkpoints use (mixer → FFN, FFN → mixer, mixer → mixer), each boundary runs the same reductions on the same tensors. Two things differ:
    • an FFN → mixer all-reduce that the fused kernel does not take may run in the next stage's input instead of in the FFN. It is the same all-reduce on the same group and tensor, as for the other models on ffn_exit;
    • at attention TP 1, the snapshot no longer all-reduces over the size-1 group.
  • Synthetic patterns with an FFN after an FFN, or at the start of the stack: the all-reduce of the zero update is gone, and the add now happens inside the fused add + norm instead of in bf16 before it. These boundaries are not bit-identical to main.
  • An a2a (for example DeepEP) MoE stage whose input is complete, i.e. at the start of the stack or right after another FFN stage, with attention TP > 1: it takes this rank's slice of the input and of the residual, then adds and normalizes, where main reduce-scattered the zero update. Like the other synthetic transitions it is not bit-identical to main. The shipped Nemotron-H / Nemotron-3 patterns, including the MTP pattern *E, put every MoE after a Mamba or attention layer.
  • A single-stage layer rejects attention CP and LayerNorm SP at construction. Neither runs with Nemotron-H: LayerNorm SP supports Qwen3 only, and no Mamba2 backend handles CP.
• Tests:
  • new test_nemotron_h_stages.py: adjacent layers agree on their shared boundary; what each kind of boundary carries; the residual follows the input onto a finer slice; the MixerExit decision table.
  • test_a_complete_value_is_sliced_onto_each_rank: a complete input into an FFN on each rank's slice is sliced and read, with and without a residual.
  • test_nemotron_h_aux_capture.py adds patterns where the FFN defers.
  • test_nemotron_h_mtp_reduction.py and the ffn_exit guard test are updated.

Accuracy Tests

No behavior change is intended on the transitions real checkpoints use.

• Call sequences, parent vs change:

  • identical: 256,128 constructor cases, 14,976 MHC cases, 53,248 prepare_mlp, 50,688 prepare_attn, 92,160 FFN-exit selection, the pair-function table and the CuteDSL exit. The communicator interface checks pass on both trees.
  • layout moves, 30,720 cases: 1,152 differ, all prepare_mlp into the scattered layout with no residual, which is now supported (read the input only). A real model always passes a residual there, except an FFN at the stack start.
  • Nemotron-H stages, 9,664 cases: no value differs and none errors.
    • 6,736 keep even the call order.
    • 48 make the same (op, group, input) calls with one reduction moved past the next layer's start. They are all FFN → mixer without the fused kernel, on M-M*- at TP 2.
    • 2,144 only drop the parent's reduction of a zero update, all on synthetic patterns with FFN → FFN or an FFN at the start (*E-M, E*E-, -, --, -M).
    • 736 only drop the parent's snapshot all-reduce over a size-1 group at TP 1.

• Unit tests:

  • the 89 related unit files pass;
  • test_nemotron_h_aux_capture.py passes 30 / 30 subtests, and the same expectations pass on the parent;
  • 125 related unit files run one process per file: only the three GPU-only files fail, on the parent as well.

• Mutants: ten are caught. They include:

  • the mixer not handing on the fused sum;
  • the residual not following to the slice;
  • the FFN not leaving its sum to a mixer;
  • the snapshot not completing;
  • an FFN at the stack start read with a residual;
  • the stack never entered;
  • NemotronHMLP's down projection reducing over attention TP, which the guard now catches.

• GPU (H200): bit-identical to the parent. Nemotron-3-Nano-30B-A3B-FP8, 4 greedy requests × 32 tokens, decode CUDA graph on:

  • TP2 with fusion on: the audit shows fused decisions on both trees;
  • TP2 with fusion off: the change's audit shows the FFN → mixer sum completed in the next stage's input;
  • TP4 × DP2;
  • PP2 × TP2.

  GSM8K (200 questions), TP4 × DP2: parent 0.960, change 0.975.

  The GPU runs used a patch that differs from this commit only in where two independent attributes are set in LayerCommunicator.__init__, plus one unit test. The call-sequence comparisons and the 125-file run are on this commit.

• Not covered on GPU:

  • an a2a (DeepEP) MoE stage;
  • reduce_scatterv deferred into the next mixer under DP attention with attention TP 1;
  • MTP. The checkpoint has no MTP layers, so MTP is covered by unit tests.

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

Speed Tests and Profiling

No kernel changes, and on real patterns no new collectives. An FFN → mixer all-reduce that the fused kernel does not take may run at the start of the next stage instead of at the end of the FFN. On synthetic FFN → FFN patterns, the zero tensor and its all-reduce are gone. Each mixer and FFN stage now wraps the scope it already opened in a small MixerExit / FfnExit object 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 #36360079782
Latest PR Test (Extra): 🚫 Run #36360079635
Latest PR Test (AMD ROCm 10): 🚫 Run #36360079806