FLAMIA: low-precision attention for Transformer training on Blackwell

FLAMIA is a new family of attention kernels designed to accelerate Transformer training on NVIDIA Blackwell GPUs. The paper shows that by jointly designing quantization, training numerics, and document scheduling, lower hardware precision can be turned into practical speedups for large language model training without harming model quality or stability (arxiv.org).

The problem

Attention mechanisms are a major computational bottleneck during Transformer training, especially when processing long text contexts (arxiv.org). While NVIDIA Blackwell GPUs offer high-throughput tensor cores for low-precision matrix multiplication—running 8-bit floating-point (FP8) operations at twice and 4-bit (FP4) operations at four times the rate of bfloat16 (BF16)—exploiting these rates in attention is difficult (arxiv.org). The authors identify three key challenges: quantization overhead, causal information leakage through shared scales, and numerical inconsistencies between the forward and backward training passes (arxiv.org). Conventional block scaling often requires multiple quantized copies of data along different contraction axes, creating memory bottlenecks. Furthermore, data-dependent scaling factors can inadvertently leak information about future masked tokens into earlier token outputs, and mismatches between forward and backward computations can destabilize training (arxiv.org).

The idea

The authors present FLAMIA, which includes both MXFP8 and mixed NVFP4/FP8 training paths (arxiv.org). To avoid duplicating quantized representations across matrix orientations, the design uses token-block scaling that reuses quantized operands (arxiv.org). In the mixed NVFP4/FP8 path, queries and keys use NVFP4 with one scale per sixteen consecutive head-dimension elements of each token, serving both forward score computation and backward score recomputation (arxiv.org). Other matrix multiplications use 8-bit E4M3 operands (arxiv.org). FLAMIA combines fused quantization, coordinated softmax and matrix-multiplication scheduling, a low-precision two-CTA backward pass, and GPU-based scheduling for packed documents (arxiv.org). To prevent scale-induced causal leakage, the scaling restricts shared exponents so that no scale group covers more than thirty-two consecutive positions, ensuring future tokens cannot alter past outputs outside their local block (arxiv.org).

The results

The authors evaluate FLAMIA against several baselines, including FlashAttention-4, cuDNN, FlashInfer, Low Precision FA4, and Hardware-Aware FP4 FlashAttention-4 (arxiv.org). Across various causal packed-document workloads, the MXFP8 path accelerates the combined forward and backward attention call—including quantization overhead—by factors of 1.10 to 1.40 times on the B200 GPU, and by 1.06 to 1.42 times on the B300 GPU, relative to the fastest evaluated baseline for each workload (arxiv.org). Additionally, in pretraining experiments on 1-billion-parameter Llama-style models over 100 billion tokens, both the MXFP8 and mixed NVFP4/FP8 training paths resulted in only a marginal increase in training loss and minimal changes to downstream metrics (arxiv.org).

Why it matters

Low-precision numerical formats offer massive theoretical performance gains on modern AI hardware, but exploiting them in complex operations like self-attention has historically introduced numerical instability, causal violations, or excessive overhead. FLAMIA demonstrates a concrete method for aligning hardware-level micro-scaling formats with the mathematical requirements of causal attention. By proving that low-bit training can match baseline stability while significantly reducing latency on Blackwell architecture, the work provides a practical blueprint for more efficient large language model pretraining (arxiv.org).

Caveats

The authors note that single-step accuracy metrics alone are insufficient to prove causal correctness or long-term training stability, requiring full pretraining validation (arxiv.org). They also evaluate their methods specifically on Llama-style model architectures and designated Blackwell GPU configurations (arxiv.org). Readers should note that this paper is a preprint and has not yet been formally peer-reviewed (arxiv.org).