Implementation:NVIDIA TransformerEngine Ops Fused Forward Linear Scale Add

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

Overview

Fused forward pass operation that combines GEMM with constant scaling and in-place residual addition into a single operation.

Description

ForwardLinearScaleAdd is a FusedOperation that fuses the forward pass of BasicLinear, ConstantScale, and in-place AddExtraInput. It passes the scale factor as the alpha parameter to BasicLinear._functional_forward and uses the extra input as the output buffer with accumulate_into_out=True. This computes output = residual + alpha * (input @ weight^T) in a single fused operation. The fusion requires non-row tensor parallelism and an in-place AddExtraInput.

Usage

Automatically applied by the operation fuser when it detects the 3-operation pattern in the forward pass. Common in scaled residual architectures.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py

Lines

1--175

Signature

class ForwardLinearScaleAdd(FusedOperation):
    def __init__(self, *, linear: BasicLinear, scale: ConstantScale, add: AddExtraInput): ...
    def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, prev_op_grad_output_quantizer, next_op_input_quantizer, basic_op_kwargs) -> Tuple: ...

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

Import

from transformer_engine.pytorch.ops.fused.forward_linear_scale_add import ForwardLinearScaleAdd

I/O Contract

Inputs

Name	Type	Required	Description
linear	BasicLinear	Yes	The linear operation (not row TP)
scale	ConstantScale	Yes	Constant scaling factor
add	AddExtraInput	Yes	In-place AddExtraInput for residual
input_	torch.Tensor	Yes	Input tensor
extra_input	torch.Tensor	Yes	Residual tensor

Outputs

Name	Type	Description
output	torch.Tensor	residual + scale * (input @ weight^T)

Usage Examples

# Automatically fused by the operation fuser when detecting pattern:
# [BasicLinear, ConstantScale, AddExtraInput(in_place=True)] in forward pass
# No manual usage required

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements