# FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell

**Authors:** Fedor Velikonivtsev, Mikhail Koltakov, Denis Kuznedelev  
**Lab:** Yandex Research · **Published:** 2 Nov 2026 · **Upvotes:** 17  
**Links:** [arXiv](https://arxiv.org/abs/2611.03344) · [PDF](https://arxiv.org/pdf/2611.03344) · [This page](https://omen-bit.github.io/ai-paper-daily/paper/2611.03344.html)

## 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× on B200 and $1.06$--1.42× 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.

---

Notes written by AI from the paper, numbers copied from the abstract. Check the paper before you rely on them. From AI Papers Daily.