Implementation:NVIDIA TransformerEngine Ops Fused Forward Linear Bias 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, optional bias, and in-place residual addition, accumulating the GEMM output directly into the residual tensor.

Description

ForwardLinearBiasAdd is a FusedOperation that fuses the forward pass of BasicLinear, optional Bias, and in-place AddExtraInput. It passes the extra input (residual) tensor as the out parameter to BasicLinear._functional_forward with accumulate_into_out=True, eliminating a separate addition kernel. The fusion requires non-row tensor parallelism and an in-place AddExtraInput.

Usage

Automatically applied by the operation fuser when it detects a compatible forward pass pattern. Common in residual connection architectures.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py

Lines

1--193

Signature

class ForwardLinearBiasAdd(FusedOperation):
    def __init__(self, *, linear: BasicLinear, bias: Optional[Bias], 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_bias_add import ForwardLinearBiasAdd

I/O Contract

Inputs

Name	Type	Required	Description
linear	BasicLinear	Yes	The linear operation (not row TP)
bias	Optional[Bias]	No	Optional bias operation
add	AddExtraInput	Yes	In-place AddExtraInput for residual
input_	torch.Tensor	Yes	Input tensor
extra_input	torch.Tensor	Yes	Residual tensor (provided via fuser)

Outputs

Name	Type	Description
output	torch.Tensor	GEMM + bias output accumulated into the residual tensor

Usage Examples

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

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements