[Refactor] Choose prepare_attn / prepare_mlp steps and fused kernels at construction

[ch-wan]

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

• [Refactor] Choose a dense layer's boundaries under attention DP from both sides' declarations #41257
• [Refactor] Pick the two-batch-overlap split's layout moves once and remove execute #41256
• [Refactor] Run the LayerNorm SP region's boundary steps in the communicator itself #41255
• [Refactor] Declare at construction the layers whose FFN completes its own reduction #41254
• [Refactor] Drive an FFN exit's flags and its completion from one selection #41253
• -> [Refactor] Choose prepare_attn / prepare_mlp steps and fused kernels at construction #41252

Motivation

prepare_mlp went through CommunicateWithAllReduceAndLayerNormFn.get_fn, whose gather re-derived on every call facts fixed when the communicator is built: whether the residual must be gathered, whether attention DP is on, and which of the two DP orders applies. Whether a fused kernel may run was checked per call too.

Other code recognized these functions by identity:

• the CuteDSL communicator overrode both prepare methods and recognized the base gather by function identity;
• the MHC communicator mapped get_fn's functions to its own by identity;
• Eagle3 aux capture predicted the post-attention kernel from the same identity.

Modifications

• mlp_input_kind() is the one table for this boundary. From the layer's layouts it gives NORM, GATHER, GATHER_MOE_CP, SCATTER or ATTN_TP_ALL_REDUCE. LayerCommunicator._select_mlp_input builds that kind's steps once and prepare_mlp runs them. The gather runs _mlp_input_without_dp, _mlp_input_dp_replicate or _mlp_input_dp_partial, picked by _mlp_input_order with the facts bound, then the MoE-CP gather for GATHER_MOE_CP. CommunicateWithAllReduceAndLayerNormFn and its get_fn are removed.
• MHC and DSA-CP override _select_mlp_input. MHC takes its own function for the same kind, with its MHC state bound; its prepare_mlp override is gone. DSA-CP keeps its own table.
• Fused kernels are entries chosen at construction and tried in order. Each entry checks only whether its kernel takes this batch.
  • prepare_mlp gets the aiter / flashinfer entry when the norm has the fused method.
  • prepare_attn gets the same, with the aiter quant variant chosen once, and hands what the previous layer left to these entries before the LayerNorm SP, input-scattered and empty-batch checks. None of those sees a pending output: LayerNorm SP runs only Qwen3ForCausalLM, which never defers; no layer defers in an input-scattered forward; and both kernels refuse an empty batch.
• CuteDSLFusionLayerCommunicator no longer overrides prepare_attn / prepare_mlp. Its finalize handoff and its all-reduce + RMSNorm become the first entries. The function-identity and residual-mode checks in its eligibility become the conditions under which each entry is chosen.
• Aux capture reads the construction-time capability (_mlp_input_may_fuse) instead of the function identity.
• Norms are construction inputs.
  • GigaChat3.5 used to swap both norms of the communicator DeepseekV2DecoderLayer had built, which would leave entries chosen for the old norms.
  • DeepseekV2DecoderLayer now builds its communicator in _build_layer_communicator, and GigaChat3.5 builds one with its own norms.

Accuracy Tests

No behavior change is intended. It was checked as follows:

• Call sequences. prepare_mlp and prepare_attn were run with every collective, kernel, buffer and platform flag stubbed and recorded; each tree selects through its own code. 53,248 prepare_mlp cases and 50,688 prepare_attn cases cover:

  • the base and CuteDSL communicators;
  • complete and pending inputs;
  • the finalize handoff.

  Everything matches the parent except:

  • kernel-availability queries are no longer made where no entry exists or before an empty-batch check (2,544 prepare_mlp and 1,560 prepare_attn cases);
  • the error text for a pending output without a residual (12,800 cases where both trees raise).

• Unit tests. 195 files under unit/models, unit/layers and unit/batch_overlap were run one process per file on parent and change; no file fails only on the change. New tests cover:

  • the kind of each layout pair, the steps each kind takes, and the order chosen per configuration;
  • MHC taking its own function for a kind;
  • the aux capability;
  • the CuteDSL entries preceding the base one.

• GPU (B200). 4 greedy requests × 32 tokens, decode CUDA graph on; text, token logprobs and top-5 logprobs are bit-identical to the parent in:

  • Qwen3-30B-A3B-FP8 TP2 with fusion on and off, GLM-4.7-Flash TP2, Qwen3.5-35B-A3B TP2;
  • DeepSeek-V2-Lite TP2 × DP2 (replicate order) and TP4 × DP2 (partial order), Qwen3-30B-A3B-FP8 TP4 × DP2 × EP2.

• Rebase onto main (after [sglang][lora] Support DP attention in LoRA backends #36389, LoRA under DP attention, which publishes the LoRA token layout at the start of prepare_attn and prepare_mlp):

  • Both methods keep that call ahead of their steps. MHC's prepare_mlp override, removed here, published it too; the base prepare_mlp now does so for MHC.
  • test_triton_dp_attention_unit.py builds its communicator with the construction-time MLP input.
  • The 121 related unit files pass at this commit except three that fail on main as well. The GPU results above were measured before the rebase; after the rebase the six-PR stack is bit-identical to main on B200 in Qwen3-0.6B TP4 × DP2, TP4 × DP4 and TP2 × DP2 and DeepSeek-V2-Lite TP2 × DP2, TP4 × DP2 and TP4 × DP4 × EP4.

• Not run on GPU: CuteDSL, which only enables on Blackwell MNNVL; it is covered by unit tests and the call-sequence comparison.

Speed Tests and Profiling

No kernel or collective changes. Per-call re-derivation of fixed facts moves to construction.

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 #36233355023
Latest PR Test (Extra): 🚫 Run #36233354897
Latest PR Test (AMD ROCm 10): ❌ Run #36233355046