Transformer Engine On ROCm and AMDGPU

This repository enables Transformer Engine (TE) on ROCm as a library to accelerate Transformer models on AMD GPUs, including using 8-bit floating point (FP8) precision on MI300 GPUs, to provide better performance with lower memory utilization in both training and inference. One of the missions is to provide an alternative to accelerate Transformer models that were previously run on NVIDIA GPUs like Hopper with best efforts to make the migration frictionless. Moreover, we add optimizations specific to AMD GPUs to get the best performance benefits out of AMD GPUs.

Feature Support Status

• Activation, cast, fused softmax, layernorm, rmsnorm, transpose, fused rope, fp8 recipe, HipRTC: fully supported
• GEMM: partially supported with following input/output types: (fp32/fp32), (fp16/fp16), (bf16/bf16), (fp8, bf8/fp16, bf16, fp32) Note: Support for rocBLAS as a GEMM backend has been removed; hipBLASLt is the only backend.
• Attention (Flash Attention, Fused Multihead Attention): partially supported: Fused Attention with AOTriton and CK backends, FlashAttention-2 without variable sequence length feature
• HipGraph, HipTX: partially supported
• Tensor Parallelism, Sequence Parallelism, Context Parallelism: supported

Installation

See docs/installation.rst for detailed installation instructions on ROCm and AMDGPU. For addtional build configuration parameters see Fused Attention Backends on ROCm section below.

CK JIT

TE uses the AITER library as fused attention backend on ROCm which includes Composable Kernels (CK) MHA kernels library. By default CK kernels are not compiled during the build of TE. Instead, they are installed to TE as sourcre files and compiled runtime with HIP compiler. This allows TE to support a wide range of configurations without the need to pre-compile kernels for all configurations, which would lead to a very long build time and large binary size. Installed TE has lib/ck_jit/ck_jit_prebuild.py script that allows manipulating JIT cache.

#In TE installation directory
python3 lib/ck_jit/ck_jit_prebuild.py --help

Like with any JIT it is important to perform warmup runs before benchmarking or using TE in production to ensure the JIT cache is populated and the kernels are compiled. After that list of built kernels can be obtained with ck_jit_prebuilt.py or by direct read of cache directory, and then ck_jit_prebuilt.py can be used on the same on another system for quick kernels rebuild if needed fully utilizing parallel multi-CPU compilation

AITER rebuilding

If CK JIT is disabled, the AITER library is built ahead of time via QoLA in the CMake build directory and rebuilt incrementally. To force a full AITER rebuild, delete the build directory and rebuild TE.

Test

Framework Agnostic C++ library unittests

After a successful Transformer Engine installation via pip install, execute the following commands to build and test the framework agnostic C++ library:

cd tests/cpp
cmake -GNinja -Bbuild . && cmake --build build
# To run util tests
./build/util/test_util
# To run operator tests using 64 threads
OMP_NUM_THREADS=64 ./build/operator/test_operator

Pytorch framework integration tests

Pytorch integration pytests under tests/pytorch/ and tests/pytorch/fused_attn/ are supported.

Jax framework integration tests

All JAX pytests are supported.

Examples

Pytorch

MNIST with optional FP8

cd examples/pytorch/mnist
python main.py
python main.py --use-te   # Linear layers from TransformerEngine
python main.py --use-fp8  # FP8 + TransformerEngine for Linear layers

Sort with minGPT

cd examples/pytorch/minGPT
python gptSort.py --use-te # Linear and layernorm from TransformerEngine
python gptSort.py --use-te --ln-mlp # In addition, use LayernormMLP from transformer engine
python gptSort.py --use-te --ln-mlp --use-fp8 # In addition, use fp8

Jax

Flax

import flax
import jax
import jax.numpy as jnp
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
from transformer_engine.common import recipe

BATCH = 32
SEQLEN = 128
HIDDEN = 1024

# Initialize RNG and inputs.
rng = jax.random.PRNGKey(0)
init_rng, data_rng = jax.random.split(rng)
inp = jax.random.normal(data_rng, [BATCH, SEQLEN, HIDDEN], jnp.float32)

# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.HYBRID)

# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    model = te_flax.DenseGeneral(features=HIDDEN)

    def loss_fn(params, other_vars, inp):
      out = model.apply({'params':params, **other_vars}, inp)
      return jnp.mean(out)

    # Initialize models.
    variables = model.init(init_rng, inp)
    other_variables, params = flax.core.pop(variables, 'params')

    # Construct the forward and backward function
    fwd_bwd_fn = jax.value_and_grad(loss_fn, argnums=(0, 1))

    for _ in range(10):
      loss, (param_grads, other_grads) = fwd_bwd_fn(params, other_variables, inp)
      # Update FP8 metas
      other_variables = te.update_fp8_metas(other_grads)

MNIST

cd examples/jax/mnist
python test_single_gpu_mnist.py # Use Flax to train MNIST with BF16 as usual
python test_single_gpu_mnist.py --use-te # Use `te.DenseGeneral` provided by Transformer Engine to train MNIST with BF16
python test_single_gpu_mnist.py --use-fp8 # Use `te.DenseGeneral` provided by Transformer Engine to train MNIST and enable FP8 training and evaluation.

Encoder

cd examples/jax/encoder
python test_single_gpu_encoder.py
python test_single_gpu_encoder.py --use-fp8

Features on ROCm Platform

GEMM tuning with hipBlasLt

TE provides an ability to manually or automatically select a GPU algorithm to use from a list generated by hipBlasLt. Selected algorithms info can be stored to file and read on further applications run. This ability is controlled by environment variables when calling GEMM operations with a specific config for the first time.

• TE_HIPBLASLT_ALGO_SELECTION - algorithm index to use in the list returned by hipBlasLt for the config or the first algorithm to select from if auto-selection is enabled; default=0.
• TE_HIPBLASLT_TUNING_RUN_COUNT - number of profiling loops for algorithm auto-selection; default=0 which means no auto-selection. For small tasks where run-to-run time variation is relatively high, using higher number of loops may give better auto-selection results.
• TE_HIPBLASLT_TUNING_ALGO_COUNT - maximal number of algorithms to check when auto-selection is enabled; default=16.
• TE_HIPBLASLT_ALGO_LOAD - filename of algorithm selection data saved by previous GEMM operation runs; if file does not exist, algorithm selection logic proceeds as if no filename were specified
• TE_HIPBLASLT_ALGO_SAVE - filename to save algorithm selection data to; can be the same as a filename to load in which case the file will be read first and then overwritten with updated results; filename may contain %i, that is replaced with the process ID. For example auto_tune_%i.csv.

It is not guaranteed that the algorithm selection data file created with one version of TE or hipBlasLt will work with other versions. Even if it works, it is highly recommended to perform algorithm selection tuning again when switching to new libraries versions because newer hipBLASLt versions may have optimized algorithms.

Typical usage is as follows:

1. Run single iteration of training enabling algorithm selection autotuning and saving:

export TE_HIPBLASLT_TUNING_RUN_COUNT=20
export TE_HIPBLASLT_TUNING_ALGO_COUNT=400
export TE_HIPBLASLT_ALGO_SAVE=algo_tune.csv
some_training_app

2. Use resulting algo_tune.csv for further training runs

unset TE_HIPBLASLT_TUNING_RUN_COUNT TE_HIPBLASLT_TUNING_ALGO_COUNT TE_HIPBLASLT_ALGO_SAVE #these variables are not needed anymore
export TE_HIPBLASLT_ALGO_LOAD=algo_tune.csv
some_training_app

If you want to check that only previously tuned algorithms are used by your application, it can be done by keeping selection data saving enabled.

export TE_HIPBLASLT_ALGO_SAVE=algo_tune_check.csv
export TE_HIPBLASLT_ALGO_LOAD=algo_tune.csv
some_training_app
#If the files are different, some not previously cached algorithms are probably used
diff algo_tune.csv algo_tune_check.csv

Single-Launch Grouped GEMM Backends on ROCm

For Mixture of Experts workloads, ROCm TE supports single-launch grouped GEMMs. Two backends are available:

• HipKittens -- MXFP8 grouped GEMM only, requires expert dimensions to be 256-aligned. Default when the grouped GEMM path is enabled on gfx950.
• CK (Composable Kernel) -- Grouped GEMM used as a fallback when HipKittens constraints are not met, or when explicitly selected.

These backends are not enabled by default. The following environment variables control backend selection:

• NVTE_USE_CUTLASS_GROUPED_GEMM=1 -- Enable the single-launch grouped GEMM path. On gfx950, defaults to HipKittens with CK fallback.
• NVTE_USE_HIPKITTENS_GROUPED_GEMM=1 -- Explicitly enable HipKittens grouped GEMM. Takes priority over CK.
• NVTE_USE_CK_GROUPED_GEMM=1 -- Explicitly select CK, bypassing HipKittens.

When none are set, TE uses multi-stream dispatch (one hipBLASLt GEMM per expert).

FlyDSL GEMM Backend on ROCm

ROCm TE provides an optional FlyDSL GEMM backend for dense (non-grouped) GEMMs, covering BF16, FP16, FP32, tensor-wise FP8, and MXFP8.

Support matrix:

• Architecture -- gfx950 (CDNA4) is the targeted architecture for now. On any other architecture the backend is not selected and TE uses its default GEMM path. gfx942 support is planned for a future release.
• FlyDSL package -- requires flydsl >= 0.3.0. FlyDSL is not a declared dependency of TE (similar to flash-attention); install it yourself with pip install flydsl.

The backend is off by default and enabled via an environment variable:

• NVTE_GEMM_BACKEND=FLYDSL -- dispatch dense GEMMs through FlyDSL when running on gfx950. Requires flydsl to be installed. (Leave NVTE_GEMM_BACKEND unset for the default C++/hipBLASLt backend, or set it to TRITON for the Triton GEMM backend.)
• NVTE_FLYDSL_GEMM_WARN_FALLBACK=1 -- emit a warning whenever a GEMM that FlyDSL cannot serve (unsupported shape/config) falls back to the default backend. Off by default.

If NVTE_GEMM_BACKEND=FLYDSL is set but flydsl is missing or older than 0.3.0, TE warns once and falls back to the default GEMM backend. Configurations FlyDSL does not support (e.g. shapes that are not tile-aligned) also fall back transparently.

Fused Attention Backends on ROCm

Currently ROCm TE supports two backends, AOTriton and CK, for fused attention. To enable specific backends in compilation and/or in runtime, the following environment variables can be used:

• NVTE_FUSED_ATTN - enable the fused attention, default = 1;
• NVTE_FUSED_ATTN_CK - enable the AITER and CK backends, default = 1;
• NVTE_FUSED_ATTN_AOTRITON - enable the AOTriton backend, default = 1.

Setting env NVTE_FUSED_ATTN_<BACKEND>=0 in compilation will skip the build of the specific backend, which saves the overall building time. Setting env NVTE_FUSED_ATTN_<BACKEND>=0 in runtime provides the option to choose specific backends in runtime. Note that one backend can be enabled in compilation but disabled in runtime. However, if one backend is disabled in compilation, the same env NVTE_FUSED_ATTN_<BACKEND>=0 is required during runtime. Otherwise TE will error out that the specific backend is not compiled.

NVTE_FUSED_ATTN has higher priority than NVTE_FUSED_ATTN_CK and NVTE_FUSED_ATTN_AOTRITON. NVTE_FUSED_ATTN=0 will use the TE unfused attention even if NVTE_FUSED_ATTN_CK or NVTE_FUSED_ATTN_AOTRITON is set. Fused attention backends are chosen according to the match results between the actual problem config and the support matrix of the specific backend. For the scenario that both backends are enabled and match the problem configuration, the CK backend will be chosen with higher priority.

Note that when using THD format tensors with CK Fused Attention, one should pass None for cu_seqlens_q_padded, cu_seqlens_kv_padded to indicate that there is no padding between sequences. Otherwise, passing proper tensors will indicate padding between sequences. This is the case for both the FusedAttention and DotProductAttention modules.

Certain settings can be enabled to potentially optimize workloads depending on the nature of the inputs and expected outputs:

• NVTE_CK_RUNTIME_NUM_SEGMENTS - by default 0, if set to 1 then the JAX integration will calculate the number of segments at runtime. Enabling this requires also disabling the GPU graph by setting XLA_FLAGS="--xla_gpu_graph_level=0".
• NVTE_CK_RUNTIME_MAX_SEQLEN - by default 0, if set to 1 then the max sequence length will be calculated at runtime. This can result in speedups in cases where there are many zero-length sequences. Enabling this while using the JAX integration requires also disabling the GPU graph by setting XLA_FLAGS="--xla_gpu_graph_level=0".
• NVTE_CK_ZERO_OUT_PAD - by default 1, if set to 0 then the output of the FA forward pass will not be initialized to zero, meaning invalid regions (representing padding) may take nonzero values. Only used if input has padding.

AITER FA v3 Kernels

ROCm TE supports flash-attention v3 fwd/bwd kernels on gfx942 and gfx950 using AITER backend. This functionality can be controlled by the following environment variables:

• NVTE_CK_USES_FWD_V3 - by default 1, if set to 0, v3 kernels will not be used for fwd pass;
• NVTE_CK_USES_BWD_V3 - by default 1, if set to 0, v3 kernels will not be used for bwd pass;
• NVTE_CK_IS_V3_ATOMIC_FP32 - by default 1, if set to 0 will use atomic fp16/bf16(w/o convert_dq kernel) in bwd pass when v3 is enabled;
• NVTE_CK_HOW_V3_BF16_CVT - by default 1, float to bf16 convert type when v3 is enabled, 0:RTNE; 1:RTNA; 2:RTZ, only applicable to the gfx942 architecture.

Float to BFloat16 Conversion in CK Backend (gfx942 only)

How fp32 converts to bf16 affects both the performance and accuracy in ck fused attn. ROCm TE provides the compile-time env NVTE_CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT with the following values available to choose from:

• 0 - standard;
• 1 - truncate with nan;
• 2 - truncate;
• 3 - standard asm, default;
• 4 - rta_asm.

Small-Sequence Attention (gfx942/gfx950)

For workloads with very short sequences (up to 17 tokens), ROCm TE provides dedicated HIP MFMA attention kernels that are more efficient than the general fused-attention path for these shapes.

This path is opt-in at runtime on the CK fused-attention backend (it replaces the usual CK/AITER kernel for eligible problems; otherwise, TE transparently falls back to the regular CK/AITER fused-attention path):

• NVTE_FUSED_ATTN_CK_SMALLSEQ - by default 0 (disabled); set to 1 to route eligible problems through the small-seq kernels.

It requires the CK backend to be enabled (NVTE_FUSED_ATTN_CK=1, the default). When enabled, a problem is routed to the small-seq kernels only when all of the following hold; otherwise TE transparently falls back to the regular CK/AITER fused-attention path:

• GPU architecture is gfx942 or gfx950;
• data type is BF16.
• head dimension is 128 or 256, with matching Q/K and V head dimensions;
• number of attention heads is 16 or 32, with no GQA/MQA (num_heads == num_gqa_groups);
• no attention bias and no dropout;
• Q, K and V in separate tensors (no packed QKV or KV layouts);
• mask type is no_mask or padding (no causal masks), without a sliding window;

Both THD (variable-length, e.g. cross-attention) and BSHD (dense self-attention with s_q == s_kv) layouts are supported, with different sequence-length ranges:

• THD — eligibility uses the runtime maximum sequence length per batch (from cu_seqlens). 1 <= s_q <= 17 and 2 <= s_kv <= 17 (s_q = 1 is supported for cross-attention).
• BSHD — eligibility uses the static sequence length. Requires no mask and s_q == s_kv with 2 <= s_q <= 17.

When using the JAX integration, the small-seq path requires XLA GPU graph capture (command buffers) to be disabled, and XLA_FLAGS must be set before the process starts, for example:

XLA_FLAGS='--xla_gpu_enable_command_buffer=' NVTE_FUSED_ATTN_CK_SMALLSEQ=1 python your_script.py

Experimental Triton Kernels on ROCm

Most CUDA kernels in Transformer Engine are hipified to run on ROCm. While the hipifiled CUDA kernels are functional, they are not necessarily optimal on ROCm. We added some Triton kernels to TE ROCm to improve the performance over the hipified kernels. Currently, we have integrated Triton kernels for cast_transpose and cast_transpose_bgrad, which are commonly used in fp8 training, and also rmsnorm kernels. This feature is still experimental as it requires relatievely newer version of Pytorch (with version >= 2.4) and Triton. Also, it only works on Pytorch extension as JAX extension does not use it.

At runtime, you can enable specific triton kernels using the specific environment variables:

• NVTE_USE_CAST_TRANSPOSE_TRITON=1 can be used to enable cast transpose (bgrad) triton kernels;
• NVTE_USE_LAYERNORM_TRITON=1 can be used to enable layernorm triton kernels.
• NVTE_USE_RMSNORM_TRITON=1 can be used to enable rmsnorm triton kernels.

MXFP8 support on ROCm (gfx95x only)

MXFP8 is supported on gfx95x GPUs for limited number of configurations and its usage is disabled by default. To enable MXFP8 support, use NVTE_ROCM_ENABLE_MXFP8 environment variable which can take the following values: * 0 - disable MXFP8 support (default); * 1 - enable MXFP8 support in fp8; * 2 - make MXFP8 a default fp8 recipe.

MXFP4 GEMM support on ROCm (gfx95x only)

MXFP4 GEMM is supported on gfx95x GPUs for a limited number of configurations. A native hipBLASLt MXFP4 (F4F4) path is available when built against hipBLASLt >= 1.3, alongside the default AITER a4w4 backend. To select the GEMM backend, use NVTE_ROCM_USE_HIPBLASLT_MXFP4 environment variable which can take the following values: * 0 - use the AITER a4w4 backend (default); * 1 - use the hipBLASLt MXFP4 GEMM backend.

Blockwise FP8 GEMM support on ROCm (gfx942 and gfx950)

Blockwise FP8 GEMM is supported on gfx942 and gfx950 GPUs through the HipKittens backend and is enabled by default. It supports 1Dx2D and 1Dx1D scaling, e4m3/e5m2 inputs in any combination except e5m2 x e5m2, bf16/fp32/fp16 output, and bias/beta/DGELU epilogues. The K and N dimensions must be multiples of 16. The backend shares the existing USE_HIPKITTENS_GEMM build option with MXFP8:

• OFF - disable the HipKittens blockwise FP8 and MXFP8 GEMM backend;
• ON - enable the HipKittens GEMM backend (default).

On gfx950 there are two blockwise FP8 GEMM kernels, selected at runtime by the NVTE_FP8_BLOCK_SCALING_FP32_SCALES environment variable, which also controls the scale produced by the Float8BlockScaling recipe:

• 0 - cast the incoming FP32 scales to E8M0 (power-of-2) and compute the GEMM with those scales (default);
• 1 - compute the GEMM with the FP32 scales directly.

Two-stage amax Kernel

Transformer Engine uses an optimized two-stage amax kernel by default. The first stage computes a per-block amax, and the second stage performs a small reduction over these block-wise maxima. This implementation reduces contention from atomic operations and provides higher performance in many cases.

If needed (e.g., for debugging or performance comparisons), you can restore the legacy single-stage atomic kernel by setting:

NVTE_USE_ATOMIC_AMAX=1

Grouped GEMM using CK_Tile

Transformer Engine provides a CK_Tile–based implementation of grouped GEMM as an alternative to the hipBlasLt-based default grouped GEMM implementation. This will provide performance improvements in most supported cases.

You can enable the CK_Tile-based backend using the same environment variables as in the upstream CUTLASS implementation:

NVTE_USE_CUTLASS_GROUPED_GEMM=1 # Enable CK_Tile-based grouped GEMM NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK=1 # Print a warning if falling back to hipBlasLt backend (e.g., due to an unsupported config)

Transformer Engine

Quick start | User guide | PyTorch API | JAX API | Examples | Releases

What is Transformer Engine?

Transformer Engine (TE) is an NVIDIA library for accelerating Transformer model training on NVIDIA GPUs. It combines optimized building blocks and fused kernels with automatic mixed-precision-style APIs for PyTorch and JAX, so low-precision training can be adopted without rewriting a training stack.

Transformer Engine manages the scaling factors, amax histories, and quantization metadata required by low-precision recipes. Its modules cover attention, linear layers, normalization, Mixture-of-Experts (MoE), and communication operations used in large-scale distributed training.

Highlights

• FP8 training on NVIDIA Hopper, Ada, Blackwell, and Rubin GPUs.
• MXFP8 and NVFP4 training on NVIDIA Blackwell GPUs.
• Optimized attention, GEMM, normalization, quantization, and fused Transformer and MoE modules.
• PyTorch and JAX APIs with autocast-style contexts and configurable low-precision recipes.
• Support for tensor, sequence, context, and EP, including communication overlap.
• FP16 and BF16 optimizations on NVIDIA Ampere architecture GPUs and later.

News

• [09/2026] Transformer Engine v2.19 adds Rubin support, hybrid quantization, MXFP8 EP communication, and expanded FP8 attention support.
• [09/2026] Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine describes optimized JAX MoE training on GB200 and GB300 systems.
• [08/2026] Transformer Engine v2.18 adds FP8 block scaling in PyTorch, zero-copy EP, and CUDA Graph support for THD attention.
• [07/2026] Transformer Engine v2.17 introduces NCCL-backed EP for PyTorch and JAX, faster MoE routing, and a JAX Flax MoE block.
• [06/2026] Boosting MoE Training Throughput with Advanced Fusion Kernels shows how fused Transformer Engine operations accelerate MoE training.

See the project updates archive for earlier news.

Quick start

Install the latest stable release for your framework:

# PyTorch
pip install --no-build-isolation "transformer_engine[pytorch]"

# JAX
pip install --no-build-isolation "transformer_engine[jax]"

For a ready-to-run environment, use an NVIDIA NGC framework container. Replace <YY.MM> with a container release listed in the NVIDIA Deep Learning Frameworks Support Matrix.

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:<YY.MM>-py3
docker run --gpus all -it --rm nvcr.io/nvidia/jax:<YY.MM>-py3

Continue with the PyTorch and JAX getting started guide. For prerequisites, source builds, environment variables, and troubleshooting, see the installation guide.

Integrations

Transformer Engine has been integrated with popular LLM frameworks such as:

• Hugging Face Accelerate
• Lightning
• NVIDIA BioNeMo Recipes
• NVIDIA JAX Toolbox
• NVIDIA Megatron-LM
• NVIDIA NeMo Megatron Bridge
• Amazon SageMaker Model Parallel Library

See Ecosystem and historical integrations for additional community integrations and projects that have worked with Transformer Engine.

Contributing

We welcome contributions to Transformer Engine! To contribute to Transformer Engine and make pull requests, follow the guidelines outlined in the CONTRIBUTING.rst guide.

Technical deep dives

• Low precision training guide — FP8, MXFP8, NVFP4, scaling recipes, and performance considerations.
• Using FP8 and FP4 with Transformer Engine — a practical introduction with code examples.
• FP8 Formats for Deep Learning — the foundational paper describing the FP8 formats used for deep learning.
• Stable and Scalable FP8 Deep Learning Training on Blackwell — a technical GTC 2025 session on training numerics and scale.

See Resources for the complete collection of papers and recorded talks.