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

Hugging Face Daily Papers Papers
attention-kernels quantization blackwell-gpu mxfp8 training-efficiency low-precision

Summary

FLAMIA 提出了一族面向 NVIDIA Blackwell GPU 的混合精度注意力内核，利用 MXFP8 与 NVFP4/FP8 路径，在避免因果信息泄漏的同时加速 Transformer 训练，在 B200/B300 上实现约 1.06–1.42 倍的前向与反向注意力加速。

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.
Original Article
View Cached Full Text

Cached at: 11/03/26, 11:37 AM

Paper page - FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell

Source: https://huggingface.co/papers/2611.03344

Abstract

FLAMIA accelerates Transformer attention on NVIDIA Blackwell GPUs via MXFP8 and mixed-precision kernels that reduce quantization overhead, prevent causal leakage, and support variable-length packed documents.

Attention remains a substantial computational cost in Transformer training, particularly at long context lengths. NVIDIABlackwell GPUsprovide high-throughput low-precision matrix multiplication, but exploiting it in attention requires addressing quantization overhead,causal information leakagethrough shared scales, and numerical inconsistencies between theforward and backward passes. We present FLAMIA, a family ofattention kernelsfor Blackwell withMXFP8and mixedNVFP4/FP8training paths. Our design combinestoken-block scalingthat reuses quantized operands across matrix orientations, targeted safeguards against scale-induced leakage, and consistent forward and backward computations.Fused quantizationand architecture-specific kernel optimizations reduce overhead, whiledocument-local scalingand workload-aware scheduling support variable-length documents directly inpacked layoutswithout repacking. Across various causal packed-document workloads, theMXFP8path 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.

View arXiv pageView PDFGitHub0Add to collection

Get this paper in your agent:

hf papers read 2611\.03344

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2611.03344 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2611.03344 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2611.03344 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

Format-Aware Fusion for Fast FP4 Pretraining

arXiv cs.LG

提出了 format-aware fusion 方法，将 FP4 量化中的 scale 计算、操作数打包与反向状态保存与 GEMM 融合，在 Llama-3 8B 预训练中达到最高 37.9K tokens/s/GPU，MXFP4 路线在略高于 bfloat16 训练损失终点的情况下实现 37.2K tokens/s/GPU（86.3% MFU），显著超过 Transformer Engine 的 27.6K。

@charles_irl: dflash go brr

X AI KOLs Timeline

NVIDIA announces DFlash, an open source block diffusion model for speculative decoding that achieves up to 15x higher inference throughput on Blackwell GPUs while maintaining interactivity.