Implementation:NVIDIA TransformerEngine Ops Fused Backward Add RMSNorm

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch, Optimization
Last Updated	2026-02-07 14:00 GMT

Overview

Fused backward pass operation that combines RMSNorm gradient computation with residual gradient addition into a single kernel.

Description

BackwardAddRMSNorm is a FusedOperation that fuses the backward pass of MakeExtraOutput (residual branch) and RMSNorm into a single tex.rmsnorm_bwd_add kernel call. This avoids a separate kernel launch for adding the residual gradient to the RMSNorm input gradient. The fusion requires a non-in-place MakeExtraOutput followed by RMSNorm in the backward operation list.

Usage

Automatically applied by the operation fuser when it detects the compatible backward pass pattern. This fusion is common in pre-norm transformer architectures.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py

Lines

1--127

Signature

class BackwardAddRMSNorm(FusedOperation):
    def __init__(self, *, add: MakeExtraOutput, rmsnorm: RMSNorm): ...
    def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs) -> Tuple: ...

    @staticmethod
    def fuse_backward_ops(ops, **unused) -> list[FusibleOperation]: ...

Import

from transformer_engine.pytorch.ops.fused.backward_add_rmsnorm import BackwardAddRMSNorm

I/O Contract

Inputs

Name	Type	Required	Description
add	MakeExtraOutput	Yes	The MakeExtraOutput (residual) operation (must not be in-place)
rmsnorm	RMSNorm	Yes	The RMSNorm operation
grad_output	torch.Tensor	Yes	Upstream gradient
basic_op_grad_extra_outputs	list	Yes	Contains residual gradient from extra output

Outputs

Name	Type	Description
grad_input	torch.Tensor	Combined gradient (RMSNorm backward + residual)
grad_weight	torch.Tensor	Gradient w.r.t. RMSNorm weight parameter

Usage Examples

# Automatically fused by the operation fuser when detecting pattern:
# [MakeExtraOutput(in_place=False), RMSNorm] in the backward pass
# No manual usage required - the fuser handles fusion automatically

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements