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/dsreluand 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);