Skip to content

[MCore] Add SiTU-GLU activation fallbacks - #6636

Closed
harryzhou2000 wants to merge 3 commits into
NVIDIA:mainfrom
harryzhou2000:hhanyu/situ-glu-fallback-main
Closed

[MCore] Add SiTU-GLU activation fallbacks#6636
harryzhou2000 wants to merge 3 commits into
NVIDIA:mainfrom
harryzhou2000:hhanyu/situ-glu-fallback-main

Conversation

@harryzhou2000

@harryzhou2000 harryzhou2000 commented Aug 18, 2026

Copy link
Copy Markdown
Member
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add SiTU-GLU as a global FFN activation on main. MCore uses the public interface proposed by NVIDIA/TransformerEngine#3402 when the complete native interface is installed and otherwise falls back atomically to MCore-local CuTe DSL operations.

Paired dev-target PR: #6635

Activation

The activation follows Kimi K3 Stable-LatentMoE:

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

K3 uses beta1=4 and beta2=25. For beta1=4, the local implementation reuses a = tanh(G / 4) and evaluates sigmoid(G) exactly as 0.5 + a / (1 + a^2), avoiding a separate exponential.

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

Implementation

  • Add --situ-glu, with --moe-use-situ-glu as a compatibility alias, and apply it to gated dense, routed-expert, and shared-expert FFNs.
  • Import SiTUGLU and ScaledSiTUGLU together from the public interface in [Common][PyTorch] Add SiTU-GLU activation TransformerEngine#3402:
    • if both imports succeed, construct the native operations with the same beta, interleave, cached-input, and activation-recompute arguments;
    • if either import is unavailable, use the matching MCore-local standalone and scaled CuTe DSL operations.
  • Preserve the existing SwiGLU and ScaledSwiGLU selection paths unchanged.
  • Keep the local scaled fallback as a distinct TE-fuser BasicOperation; it is not a ScaledSwiGLU subclass and therefore cannot be mistaken for SwiGLU by TE's grouped-MLP matcher.
  • Remove the former global cuDNN frontend class replacement and all CuTe DSL compatibility monkeypatches.
  • Preserve SwiGLU-compatible checkpoint conversion, FLOP accounting, and memory estimation.

Native block-scaled grouped SiTU-GLU fusion is supplied by NVIDIA/cudnn-frontend#645. MCore does not duplicate or patch those kernels.

Validation

Baseline B200 environment (TE 2.20.0.dev0, cuDNN frontend 1.26.0, CuTe DSL 4.5.2):

  • All 23 focused tests pass on the main-target worktree.
  • Real MCore MoE forward/backward passes for BF16, MXFP8, and NVFP4 with the local SiTU-GLU fallback.
  • Real MCore fused grouped-MLP forward/backward passes for the pre-existing SwiGLU path under MXFP8 and NVFP4; the realized fuser plan contains GroupedMLP_CuTeGEMMGLU.
  • Tests verify complete-native-interface preference, atomic fallback for an incomplete TE interface, activation-recompute contract handling, and preservation of the ordinary SwiGLU path.

Native-interface environment (NVIDIA/TransformerEngine#3402, NVIDIA/cudnn-frontend#645, CuTe DSL 4.6.2):

  • BF16 native SiTU-GLU and SwiGLU forward/backward pass.
  • NVFP4 fused SiTU-GLU and SwiGLU forward/backward pass on the main-target worktree, and both select GroupedMLP_CuTeGEMMGLU.
  • MXFP8 fused SwiGLU forward/backward passes on the equivalent dev-target implementation. Native SiTU-GLU reaches the fused forward path; its dSiTU backward kernel currently hits a vector.from_elements BF16-to-FP32x2 verifier error inside the dependency implementation. No MCore or cuDNN frontend workaround is included here.

The commits are SSH-signed and signed off. git diff --check and the repository tools/autoformat.sh pass; pylint rates the changed files 10/10 and Ruff reports no errors.

Scope and dependencies

This PR owns the MCore configuration, selection logic, and local fallback. It does not change grouped-GEMM fusion policy or patch Transformer Engine/cuDNN frontend internals.

This PR does not depend on Quantile Balancing or #6638. The paired integration checkout was used only to verify that the independent features coexist.

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.

Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 18, 2026

Copy link
Copy Markdown

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.

Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
@harryzhou2000

Copy link
Copy Markdown
Member Author

Superseded by the smaller TE-only main-target PR: #6674. The previous implementation is preserved at branch hhanyu/situ-glu-fallback-main-bkp.

@harryzhou2000
harryzhou2000 deleted the hhanyu/situ-glu-fallback-main branch August 19, 2026 19:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant