Implementation:NVIDIA TransformerEngine Ops Fused Backward Linear Scale

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

Overview

Fused backward pass operation that combines dgrad GEMM with constant scaling, applying the scale factor directly during the GEMM computation.

Description

BackwardLinearScale is a FusedOperation that fuses the backward pass of BasicLinear and ConstantScale by passing the scale factor as grad_input_alpha and grad_weight_alpha to the BasicLinear._functional_backward call. This avoids a separate scaling kernel. The fusion requires a BasicLinear followed by ConstantScale in the backward operation list, and column tensor parallelism is not supported.

Usage

Automatically applied by the operation fuser in backward pass pattern matching. Common in architectures that use residual scaling.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/fused/backward_linear_scale.py

Lines

1--165

Signature

class BackwardLinearScale(FusedOperation):
    def __init__(self, *, scale: ConstantScale, linear: BasicLinear): ...
    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_linear_scale import BackwardLinearScale

I/O Contract

Inputs

Name	Type	Required	Description
scale	ConstantScale	Yes	The constant scaling operation
linear	BasicLinear	Yes	BasicLinear operation (not column TP)
grad_output	torch.Tensor	Yes	Upstream gradient

Outputs

Name	Type	Description
grad_input	torch.Tensor	Scaled dgrad
grad_weight	torch.Tensor	Scaled weight gradient

Usage Examples

# Automatically fused by the operation fuser when detecting pattern:
# [BasicLinear, ConstantScale] in the backward pass
# No manual usage required

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements