FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell
Fedor Velikonivtsev, Mikhail Koltakov, Denis Kuznedelev
Yandex Research
Published November 2, 2026 · on the feed November 3, 2026 · ▲ 18 · arXiv 2611.03344
The 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.40times on B200 and 1.06--1.42times 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.
The Brief
2:00A short podcast summary of the paper. Written and voiced by AI.
Transcript
FLAMIA. FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell. By Fedor Velikonivtsev and two others, from Yandex Research. Published on the second of November.
Attention is still a large part of the compute bill in Transformer training, and the bill grows at long context lengths. NVIDIA Blackwell GPUs offer high-throughput, low-precision matrix multiplication, which looks like an obvious fit. Using it inside attention is not free, though. There is the overhead of quantizing in the first place. There is causal information leakage through shared scales, which matters for a causal model. And the forward and backward passes can disagree numerically.
FLAMIA is a family of attention kernels for Blackwell that takes on all three problems together. It has an MXFP8 training path and a mixed NVFP4 and FP8 path. The premise is that quantization, training numerics and document scheduling have to be designed jointly, rather than bolted on one at a time.
Several pieces fit together. Token-block scaling lets quantized operands be reused across matrix orientations. Targeted safeguards guard against leakage induced by scales. Forward and backward computations are kept consistent with each other. To cut overhead, quantization is fused and the kernels get architecture-specific optimizations. And for training data packed from many documents, document-local scaling and workload-aware scheduling handle variable-length documents directly in the packed layout, with no repacking.
The tests use causal packed-document workloads. The MXFP8 path speeds up the combined forward and backward attention call, with quantization counted in. On B200 the speedup ranges from one point one times to one point four times. On B300 it ranges from one point zero six times to one point four two times. Each comparison is against the fastest evaluated baseline for that workload. The authors also pretrain Llama-style models to check that training stays stable and quality is preserved.
Raw low-bit throughput on new hardware does not turn into faster training by itself. Here the gain comes from treating scaling, numerical consistency and packed documents as one design problem.
More with an account
Listen to every day’s papers
Daily playlists, deep dives, quests for the papers, and quizzes on their key terms.
Join free