Implementation:NVIDIA TransformerEngine Transpose C API

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

Overview

Declares the C API for transpose operations, including fused cast-transpose, cast-transpose with bias gradient reduction, and fused activation-backward + cast-transpose + dbias operations.

Description

transpose.h is critical for FP8 training performance. The backward pass of linear layers requires transposed weight/activation gradients in FP8 format plus bias gradients. Fusing these into single kernels eliminates multiple passes over large tensors.

Key functions:

• nvte_transpose: Plain transpose
• nvte_cast_transpose: Cast + produce both rowwise and columnwise data
• nvte_cast_transpose_dbias: Cast-transpose + column reduction for bias gradients
• nvte_fp8_transpose_dbias: FP8 transpose + dbias
• nvte_multi_cast_transpose: Batch cast-transpose for multiple tensors
• Fused dactivation variants: nvte_cast_transpose_dbias_dgelu/dsilu/drelu/dqgelu/dsrelu and their gated variants (dgeglu, dswiglu, dreglu, dqgeglu, dsreglu)

Usage

Use in the backward pass of linear and MLP layers to efficiently produce FP8-quantized transposed gradients along with bias gradients.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/include/transformer_engine/transpose.h

Lines

1--333

Signature

void nvte_transpose(const NVTETensor input, NVTETensor transposed_output,
                    cudaStream_t stream);
void nvte_cast_transpose(const NVTETensor input, NVTETensor output,
                         cudaStream_t stream);
void nvte_cast_transpose_dbias(const NVTETensor input, NVTETensor output,
                               NVTETensor dbias, NVTETensor workspace,
                               cudaStream_t stream);
void nvte_multi_cast_transpose(size_t num_tensors, const NVTETensor* input_list,
                               NVTETensor* output_list, cudaStream_t stream);
void nvte_cast_transpose_dbias_dgelu(const NVTETensor input,
                                     const NVTETensor act_input, NVTETensor output,
                                     NVTETensor dbias, NVTETensor workspace,
                                     cudaStream_t stream);

Import

#include "transformer_engine/transpose.h"

I/O Contract

Inputs

Name	Type	Required	Description
input	NVTETensor	Yes	Input tensor of shape [N, H]
stream	cudaStream_t	Yes	CUDA stream for the operation

Outputs

Name	Type	Description
output	NVTETensor	Rowwise data [N, H] and columnwise data [H, N]
dbias	NVTETensor	Bias gradient from column reduction [H]

Usage Examples

#include "transformer_engine/transpose.h"

// Fused cast-transpose with bias gradient
nvte_cast_transpose_dbias(input, fp8_output, dbias, workspace, stream);

// Fused dGELU + cast-transpose + dbias
nvte_cast_transpose_dbias_dgelu(grad_output, gelu_input,
                                fp8_output, dbias, workspace, stream);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements