[MCore] Add SiTU-GLU through Transformer Engine

[harryzhou2000]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add Kimi K3 SiTU-GLU as a global FFN activation on main.

Paired dev-target PR: #6673

Activation

For gate projection G, up projection U, and K3 defaults beta1=4, beta2=25:

T_g = beta1 * tanh(G / beta1) * sigmoid(G)
T_u = beta2 * tanh(U / beta2)
Y   = T_g * T_u

Reference: Kimi K3 technical report, Figure 4 and Equation 12.

Design

• Add --situ-glu, with --moe-use-situ-glu as an alias, and apply it consistently to dense, routed-expert, and shared-expert FFNs.
• Add situ_glu_beta1=4 and situ_glu_beta2=25 to TransformerConfig, with validation that SiTU-GLU is gated, unclamped, uses zero GLU linear offset, and has finite positive beta values.
• Add situlu as the correct PyTorch pointwise reference and configuration marker. It serves non-TE/custom module paths until PyTorch provides a dedicated torch.nn.functional.situlu-style operation; unary F.silu is not equivalent.
• Use transformer_engine.pytorch.ops.SiTUGLU for ordinary dense, sequential-expert, and ordinary shared-expert paths.
• Use transformer_engine.pytorch.ops.ScaledSiTUGLU for fused grouped dense, routed-expert, and shared-expert paths.
• Fail clearly when a requested TE path does not expose the required SiTU-GLU operator. cuDNN frontend capability selection remains inside TE.
• Extend checkpoint conversion, default FFN sizing, FLOP accounting, and memory estimation to treat SiTU-GLU as a gated GLU.

The required TE interface is available in NVIDIA/TransformerEngine#3402.

With that TE interface, an older cuDNN frontend without SiTU parameters uses TE's unfused GroupedLinear -> ScaledSiTUGLU -> GroupedLinear operation sequence. A SiTU-capable cuDNN frontend enables fused grouped GEMM + SiTU-GLU without an MCore-side backend check.

Test results

Validated on NVIDIA B300 with NVIDIA/TransformerEngine#3402 and cuDNN Frontend develop through NVIDIA/cudnn-frontend#670, using CuTe DSL 4.6.2:

• final focused SiTU-GLU and grouped-MLP suite on the review-fixed tip: 60 passed;
• broader selected main suites: 126 passed, 11 skipped, 4 deselected; the deselected cases are existing multi-rank shared-expert parameterizations that cannot run on one GPU;
• BF16 real MCore MoE routed forward/backward through ScaledSiTUGLU: passed;
• MXFP8 real MCore MoE forward/backward: routed and fused-shared ScaledSiTUGLU, with the fused cuDNN/CuTe grouped-MLP operation asserted: passed;
• NVFP4 routed-expert forward/backward through ScaledSiTUGLU, with the fused Hadamard forward operation asserted: passed;
• git diff --check, Python compilation, and Ruff on all changed source and test files: passed;
• SSH commit signature and DCO: passed.
• --use-checkpoint-args restoration passed two parametrized SwiGLU/SiTU-GLU cases, including the SiTU selector and non-default beta values.
• The hybrid-model golden-config comparator passed on the current rebased tip.

Fused-shared NVFP4 currently reaches an existing TE FC1 weight-scale reshape failure that reproduces with an equivalent SwiGLU control, so this PR does not claim that combination. MXFP8 dSiTU validation used the cuDNN frontend correction included in NVIDIA/cudnn-frontend#670; the published cuDNN frontend 1.27.0 package tested here does not yet contain that correction.

Scope and dependencies

This PR depends on NVIDIA/TransformerEngine#3402 for TE-backed SiTU-GLU functionality. The cuDNN frontend support is optional for correctness and required only for fused block-scaled grouped execution. A functional test is deferred until that TE interface is available in Megatron-LM CI; the focused MCore unit and GPU suites cover the current integration boundary.

This PR is independent of Quantile Balancing and #6638.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue: TODO — link a Megatron-LM feature-request issue before marking this draft ready for review.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[harryzhou2000]

Mirrored the SiTU-GLU review fixes from the paired dev PR #6673 in 89b265c25: narrow SiTU-only recompute gating, preserved SwiGLU/quick-GEGLU behavior, checkpoint-converter gated-layout metadata, fixed-value saturation coverage, and exact formatter cleanup. The final focused NVIDIA B300 suite is 60 passed.