Transformer Engine

Transformer Engine (TE) is a library developed by NVIDIA for accelerating Transformer models on NVIDIA GPUs. It enables the use of 8-bit floating point (FP8) and 4-bit floating point (NVFP4) precision on architectures like Hopper, Ada, and Blackwell, significantly improving performance and reducing memory utilization during both training and inference. TE provides highly optimized building blocks for common Transformer architectures and an automatic mixed-precision-like API that integrates seamlessly with PyTorch and JAX. The library has frequent releases, often aligned with updates to NVIDIA's deep learning software stack.