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

AI summary

Overview

Research area: Efficient machine learning systems — specifically low-precision attention kernels for Transformer training on NVIDIA Blackwell GPUs, sitting at the intersection of numerics (microscaling quantization), kernel engineering, and training stability.

Technical level: Advanced. The paper assumes familiarity with FlashAttention-style tiling, online softmax, and block-scaled matrix multiplication formats (MXFP8, NVFP4). The prose is readable, but the core contributions are kernel-level.

Scope (one sentence): The paper designs and evaluates a family of attention kernels (MXFP8 and mixed NVFP4/FP8) that make low-bit attention training on Blackwell both fast and causally correct, including on variable-length packed documents, while preserving model quality in Llama-style pretraining.

What This Paper Is About

Attention is a major computational cost in Transformer training, especially at long context lengths, and Blackwell tensor cores can run FP8 at twice and FP4 at four times the BF16 rate — but simply switching to low precision breaks in three specific ways. First, Blackwell's block-scaled matrix multiply requires scaling groups along each operand's contraction axis, and attention's six logical operands each participate in two products with different contraction axes, so naive scaling doubles quantization work. Second, sharing a quantization scale across token positions can let a future token change an earlier token's output even under a causal mask. Third, the backward pass must reconstruct exactly what the forward pass computed, or training destabilizes. The goal is to exploit low-bit hardware throughput for attention training without sacrificing causal correctness, numerical consistency, or throughput on packed variable-length documents.

Key Contributions

1. Two training paths built on token-block scaling. An MXFP8 path and a mixed NVFP4/FP8 path, both using a scaling scheme that lets one 8-bit payload serve both matrix orientations, avoiding separately quantized transposed copies of Q, K, V, dO, P, and dS.

2. A diagnosis and mitigation of scale-induced causal leakage and forward–backward inconsistency. This includes per-key scales on diagonal tiles, a "K fold," exponent rounding-up so nothing clamps, an optional causal dQ scale, and consistent reconstruction of quantized intermediates across passes — plus experiments showing that single-step gradient similarity does not establish training stability.

3. Architecture-specific kernel engineering. Fused quantization (Hadamard rotation, block statistics, and code/scale generation in one launch), coordinated softmax and MMA scheduling, a low-precision two-CTA backward that needs no TMEM scale storage for dP, GPU-based load-balanced scheduling for packed documents without host synchronization, and two bit-level exponential approximations for inference.

4. An evaluation against six baseline families (FlashAttention-4 BF16/FP8, cuDNN BF16/FP8/MXFP8, FlashInfer, Low Precision FA4, HW-FP4 FA4, and PyTorch SDPA on nested tensors), plus training validation on Llama-style pretraining at 1B parameters over 100B tokens. Code is released at https://github.com/yandex-research/flamia-blackwell-attention.

Main Findings

• MXFP8 wins every causal packed-document training call. Across all nine causal workloads on both B200 and B300, the MXFP8 path has the lowest combined forward + backward attention latency including quantization, at 1.10–1.40× faster on B200 and 1.06–1.42× faster on B300 than the fastest evaluated baseline for each workload. Example: on B200 with 4×4096-token documents, MXFP8 runs the training call in 0.89 ms (1079 TFLOP/s) versus 1.06 ms for the fastest baseline (cuDNN FP8), a 1.18× speedup.

• The all-4-bit training path is not the default, and for good reasons. The 4-bit MMA instructions used here have no MN-major form, so each orientation needs its own copy; the six quantized copies make quantization the largest part of that path's backward time. It is slower than paths keeping most operands in 8 bits, reaches slightly higher loss, and 4-bit V adds per-key scaling work on diagonal tiles and needs separate representations for PV and dP.

• NVFP4 has the fastest backward; MXFP8 has the fastest forward. In kernel-only measurements, NVFP4 achieves the lowest causal backward latency on every workload, at 1.22–2.01× on B200 and 1.17–2.00× on B300. MXFP8 leads the forward on all nine B200 workloads and eight of nine B300 workloads, with ranges of 1.07–1.72× and 0.98–1.69×. NVFP4's forward is slower than MXFP8 because its extra softmax work outweighs savings in QKᵀ.

• Leakage and instability are easier to miss than to measure. A model trained with a leaking kernel obtains lower loss when evaluated through that kernel but performs worse under exact causal attention. Separately, mismatched tiling or exponential evaluation between forward and backward can destabilize training even when single-step cosine similarity is high. Validation loss measured through the training kernel alone is therefore insufficient.

• The dS quantization scale mattered most among numerical choices. In ablations, the fraction of small dS entries rounded to zero tracked training degradation more closely than the reported norm error, motivating a per-32×32-block scale derived from the tile already held in the kernel.

• Scale bounding matters for correctness, not just accuracy. Conditional rescaling lets unnormalized weights reach 2^(c+τ) after scaling; since E4M3's largest finite value is 448, the kernels enforce c + τ ≤ log₂448. Violating this can saturate weights and corrupt output without producing a NaN or runtime error.

• Quality is preserved in real pretraining. In 1B-model pretraining on 100B tokens, MXFP8 and mixed NVFP4/FP8 lead only to a marginal increase in training loss and a change in downstream metrics (the paper does not report per-metric deltas in the provided text).

• Non-causal results are more mixed on B300. MXFP8 leads all nine non-causal B200 cases (1.04–1.30×); on B300 there are five clear wins, two ties at displayed precision, and two small losses. Non-causal backward speedups are 1.17–1.49× (B200) and 1.01–1.43× (B300).

• B300 does not shrink the advantage uniformly. B300/B200 training-call latency ratios are 0.85–0.93 for cuDNN FP8, 0.92–0.95 for FA4 BF16, and 0.89–0.94 for MXFP8.

Methodology in Plain English

Start from the fastest existing design and fix what low precision breaks. The authors extend FlashAttention-4's warp-specialized kernel structure — separate warps for loading, MMA, softmax, correction, and epilogue, communicating through tensor memory — and then rework the parts that assume BF16.

Make one quantized copy do two jobs. Low-precision MMA needs each operand divided into scaling groups along the axis it gets contracted over. For attention, that axis differs between the forward and each gradient computation, which normally means quantizing every tensor twice. The authors instead share one power-of-two exponent across 32 consecutive tokens and the whole head dimension, so the representation survives transposition, and use 32×32 blocks for dS so it works for both dQ (contracted over keys) and dK (contracted over queries). Since 8-bit MMA accepts both K-major and transposed layouts, one payload serves both orientations. In the mixed path, only Q and K go to NVFP4 — using a scale per 16 head-dimension elements per token, identical for forward scores and backward score recomputation, which share a contraction axis.

Plug the leaks rather than check for them afterwards. A shared scale is a channel through which a masked future token can influence the past. In the diagonal tiles, where the visible/masked boundary cuts through a scale group, the authors give V a per-key scale folded into P; they do the same for V in dP and for K in dQ (the "K fold"). Only diagonal tiles pay this extra cost. They also round scale exponents up so the scaled maximum never exceeds the largest code, eliminating the clamping that otherwise lets a later token re-round earlier codes. An optional causal dQ scale removes the sequence-wide dependency introduced by the FP16 dQ accumulator's prescale. Explicit "exact" options remove the residual within-block effect on supported batched causal calls.

Make backward rebuild exactly what forward built. Score recomputation uses the same quantized Q and K; the backward reproduces the forward's tiling, diagonal decisions, exponential evaluation, and scaling relative to the row maximum; and D is computed from the same quantized dO used in dP.

Overlap everything that can be overlapped. Quantization of each input is fused into a single launch that also does the Hadamard rotation and block statistics. The MMA warp issues half of the next tile's QKᵀ before the current tile's PV and the other half after, keeping tensor cores busy while softmax works. The NVFP4 path quantizes probabilities in the log domain — computing each 32-key group's maximum before exponentiation and subtracting it from the exponential argument — removing a per-entry division and a scaling pass.

Handle packed documents in place. Quantizers read and write directly in the packed layout, keeping token blocks and statistics inside document boundaries and producing dense per-token scale tables. Since document starts don't align to the 128×4 scale-factor atoms the MMA consumes, the load warp assembles each tile's atom in shared memory from those tables, so no repacking or copying of Q/K/V/gradients is needed. Invalid query rows get an infinite log-sum-exp sentinel; out-of-document keys are masked.

Balance the schedule on the GPU. Work items of 256 consecutive queries are costed by how many key blocks they attend to, sorted by decreasing cost, and assigned to CTAs in alternating directions by a persistent grid; skewed workloads use Cluster Launch Control to acquire items dynamically — all built once per microbatch with no host synchronization.

Then measure honestly. Benchmarks time the full forward + backward training call including quantization on nine packed-document workloads (equal-length 4×4096, 16×4096, 16×1024, 64×1024; skewed 8192+32×256 and 8192+224×256; log-uniform 16k, 16k with GQA, and 64k), with 16 heads and d=128, on B200 and B300, against baselines using their supported ragged APIs.

Why This Matters

Impact on research. The paper's most transferable lesson is methodological: quantization for training cannot be validated by single-step gradient similarity or by loss measured through the quantized kernel itself. Scale-induced leakage and forward–backward mismatch produce structured errors that accumulate, and a leaking kernel can look better on its own metric while being worse under exact attention. This reframes how low-precision training work should be evaluated. The token-block scaling trick — choosing scale-group shapes so one payload serves both contraction axes — is also a reusable idea beyond attention.

• Long-context model pretraining: variable-length documents packed into single rows are the standard pretraining format, and attention cost grows with context; the packed-document, no-repacking path targets exactly that setting.
• Foundation-model training infrastructure: the 1B-parameter, 100B-token pretraining validation shows the speedups are usable at scale rather than only in microbenchmarks.
• Inference serving: the two bit-level exponential approximations (Mitchell-P for 4-bit probabilities, Schraudolph-style for 8-bit) target the softmax exponential, which the paper identifies as a major forward bottleneck — relevant to latency-sensitive serving.
• Retrieval and long-document workloads: the log-uniform 16k/64k and skewed 8192+short-document workloads mirror real mixtures of one long document among many short ones.

Industry relevance. Attention kernels sit directly on the critical path of training cost, and Blackwell's FP8/FP4 throughput advantages are only realizable if quantization overhead, causal correctness, and backward consistency are all solved together. The paper's claim that jointly designing quantization, training numerics, and document scheduling is what converts low-bit hardware throughput into practical speedup is directly actionable for teams running large-scale pretraining. The Apache-style public code release lowers the barrier to adoption.

Future Directions

• Closing the all-4-bit gap. The all-4-bit training path is currently slower and slightly less accurate than the 8-bit-dominant paths, largely because 4-bit MMA has no MN-major form and needs six separate quantized copies. A layout or instruction-level fix that removes this duplication would change the calculus.

• Pushing the 8-bit path's remaining dependencies to zero by default. The exact, leakage-free options exist only for supported batched causal calls and cost extra; making them cheap enough to be the default, and extending them to the fully general packed setting, is a natural next step.

• Reconciling inference approximations with training. The exponential approximations are inference-only because a backward using the same approximation lost more time than the forward gained. A backward formulation that tolerates the approximation would let those savings reach training.

• Generalizing beyond the evaluated configurations. Results are reported at 16 heads and head dimension 128, with MHA and 16/4 GQA; behavior at other head dimensions and GQA ratios, and the non-causal B300 regressions, remain open questions. The paper also notes bitwise gradient equality depends on accumulation order because the default dQ reduction uses atomics.

Target Audience

Kernel and systems engineers working on Transformer training infrastructure; researchers in low-precision / quantized training and microscaling formats; performance engineers targeting NVIDIA Blackwell (B200/B300); and ML practitioners running long-context or packed-document pretraining who need to know whether low-bit attention is safe to adopt and under what conditions. Readers without background in GPU kernel design or block-scaled numerics will find the performance results accessible but the mechanisms dense.

Authors’ 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\times$ on B200 and $1.06$--$1.42\times$ 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.