arxiv:2611.03344

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

Published on Nov 2

· Submitted by

Denis Kuznedelev on Nov 3

Yandex Research

Upvote

Authors:

Fedor Velikonivtsev ,

Denis Kuznedelev

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.

Generated by thinkingmachines/Inkling-Small

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.

View arXiv page View PDF GitHub 5 Add to collection

Community

SpiridonSunRotator

Paper author Paper submitter about 15 hours ago

•

edited about 15 hours ago

FLAMIA (FLAsh MIcroscaling Attention) adapts FlashAttention to Blackwell’s FP8/FP4 hardware. Its main design choices are:

Selective low precision. Two main training paths: MXFP8, or NVFP4 for query–key score computation and its backward recomputation, with FP8 for the remaining products. All-FP4 adds quantization overhead and slightly worsens loss.
Reusable quantized representations. Power-of-two scales shared over blocks of 32 tokens across all head channels—and 32 × 32 blocks for score gradients—allow FP8 operands to serve both matrix orientations, reducing duplicate quantization and memory traffic.
Quantization-aware causality. Shared scales can leak future-token information despite causal masking. Local scales, saturation avoidance, and per-key scale folding mitigate leakage; exact causal options remove residual dependencies in supported batched calls.
Consistent forward and backward numerics. Probability recomputation follows matching quantization, tiling, scaling, and exponential conventions. High single-step gradient similarity alone does not guarantee stable training.
Hardware-aware kernels. Fused quantization, overlap between softmax and matrix multiplication, and a low-precision backward pass using two cooperative thread arrays reduce overhead and improve GPU utilization.
Direct packed-document support. Document-local quantization avoids repacking, while GPU scheduling balances work across variable document lengths.

FLAMIA achieves approximately 1.06–1.42× faster combined forward/backward attention, including quantization, across the evaluated causal packed-document workloads.