Flash attention method speeds up transformer training on Nvidia Blackwell GPUs
FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell
Machine Learning
Summary
Training transformers, a type of AI model, needs a lot of computing power because of a step called attention. The authors found a way to make this step faster on NVIDIA Blackwell GPUs using clever tricks with low-precision math and careful scheduling. Their method, called FLAMIA, keeps training stable and accurate while reducing the time needed. This helps AI models learn more quickly without losing quality.
What this means in practice
- •For machine learning engineers: Train large language models faster by running attention computations efficiently on Nvidia Blackwell GPUs using low-precision quantization with FLAMIA.
- •For gpu kernel developers: Implement attention kernels that avoid causal information leaks while exploiting hardware-specific quantization and fused operations for speed and accuracy.
Authors
Fedor Velikonivtsev, Mikhail Koltakov, Denis Kuznedelev
Abstract
Attention remains a substantial computational cost in Transformer training, particularly at long context lengths. NVIDIA Blackwell GPUs provide high-throughput low-precision matrix multiplication, but exploiting it in attention requires addressing quantization overhead, causal information leakage through shared scales, and numerical inconsistencies between the forward and backward passes. We present FLAMIA, a family of attention kernels for Blackwell with MXFP8 and mixed NVFP4/FP8 training paths. Our design combines token-block scaling that reuses quantized operands across matrix orientations, targeted safeguards against scale-induced leakage, and consistent forward and backward computations. Fused quantization and architecture-specific kernel optimizations reduce overhead, while document-local scaling and workload-aware scheduling support variable-length documents directly in packed layouts without repacking. Across various causal packed-document workloads, the MXFP8 path accelerates the combined forward and backward attention call, including quantization, by factors of $1.10$--$1.40\times$ on B200 and $1.06$--$1.42\times$ on B300 relative to the fastest evaluated baseline for each workload. We validate training stability and quality preservation on pretraining Llama-style models. These results show how jointly designing quantization, training numerics, and document scheduling translates low-bit hardware throughput into practical attention acceleration.