[VikingDeng]

Description

First step of #3724, reworked after review: instead of a dedicated RMSNormGated op, this adds a small SiLUGate basic op, so gated RMSNorm is

te_ops.Sequential(te_ops.RMSNorm(hidden_size), te_ops.SiLUGate())

SiLUGate takes the gate as an extra input and computes y = x * silu(gate). The SiLU part is computed in FP32; y and dx are returned in the input dtype and dgate in the gate's dtype. Output quantization isn't supported, so it returns an unquantized tensor and the next op quantizes. A fused norm+gate kernel can come later as a fused op that only covers the cases the kernel supports.

Refs #3724, NVIDIA/Megatron-LM#6757

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• transformer_engine/pytorch/ops/basic/silu_gate.py: new SiLUGate basic op with one extra input (the gate).
• Export it from transformer_engine.pytorch.ops and add it to the API docs.
• tests/pytorch/test_fusible_ops.py:
  • test_silu_gate checks y, dx and dgate against a reference for each dtype and two shapes.
  • test_rmsnorm_silu_gate checks Sequential(RMSNorm, SiLUGate) end to end (y, dx, dweight, dgate), with zero-centered gamma on and off. The reference rounds the normalized output and its gradient to the input dtype like the op chain does, instead of loosening tolerances.

Tested on an RTX 3090 (SM86), torch 2.13.0+cu130:

python -m pytest -q tests/pytorch/test_fusible_ops.py -k "silu_gate or rmsnorm"   # 96 passed

I haven't run it on Hopper/Blackwell yet.

AI assistance: I used Claude Code to help write the code and tests. I reviewed the change, ran the validation above on my own GPU, and I'll handle the review.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

This seems like a rather large and bespoke operation to treat as a first-class citizen. How about we separate the RMSNorm and gate logic by just implementing a SiLUGate basic op? Then you can reproduce this function with Sequential(RMSNorm(...), SiLUGate()). You can easily change to a different norm or a different gating function (if we add more in the future).

Once you have a fused kernel, the right design is to expose it in a fused op. This has several advantages:

• Your fused kernel doesn't need to be fully general. Only enable the fusion for the cases that the kernel supports, and in other cases we fallback to the unfused impl.
• If you only fuse the forward pass or only fuse the backward pass, your implementation only needs to consider that case.

[VikingDeng]

Thanks, that's cleaner. I replaced RMSNormGated with a SiLUGate basic op (gate as an extra input), so the gated norm is now Sequential(RMSNorm(...), SiLUGate()). Tests cover the op alone and the RMSNorm + SiLUGate chain. A fused kernel would come later as a fused op, as you suggested.

[VikingDeng]

@timmoon10 when you have time, could you take another look? This is now just the SiLUGate basic op (gate as an extra input), with gated RMSNorm expressed as Sequential(RMSNorm(...), SiLUGate()), as you suggested.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge, with no outstanding findings.

Summary

Adds SiLUGate, which computes x * silu(gate) as a fusible PyTorch operation.

• PyTorch sequences can multiply a tensor by the SiLU of a separate gate.

_{Reviews (2) · Last reviewed commit: "[PyTorch] Cover negative gates in SiLUGa..." · Reviewed by Greptile}