PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Fedor Velikonivtsev, Mikhail Koltakov, Denis Kuznedelev

arXiv:2611.03344v1cs.LG

TL;DR

Long-context attention remains costly, and Blackwell’s higher low-precision throughput is difficult to exploit without addressing quantization overhead, causal leakage, and forward–backward numerical consistency. FLAMIA combines MXFP8 and mixed NVFP4/FP8 kernels with token-block scaling, leakage safeguards, and packed-document scheduling. Its MXFP8 path achieves the lowest measured combined forward and backward latency, including quantization, across the evaluated causal packed-document workloads.

  • Problem

    Attention remains a substantial Transformer training cost at long contexts, while exploiting Blackwell’s low-precision throughput requires addressing quantization overhead, causal leakage, and forward–backward numerical consistency.

    Source: §1 · p. 1

  • Method

    FLAMIA provides MXFP8 and mixed NVFP4/FP8 attention kernels that reuse token-block-scaled representations, mitigate scale-induced leakage, and compute consistently across forward and backward passes.

    Source: §1 · p. 2

  • Results

    1.10–1.40× on B200 and 1.06–1.42× on B300: MXFP8 speeds combined forward and backward attention latency, including quantization, over the fastest evaluated baseline per workload.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across evaluated causal packed-document workloads, MXFP8 achieves the lowest measured combined forward and backward attention latency, including quantization.

    Source: §5 · p. 15

    The evaluated all-4-bit training path is slower than paths retaining most operands in 8 bits and reaches slightly higher loss.

    Source: §3.4 · p. 6

Abstract

from arXiv · show

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.

1. Introduction

FLAMIA targets low-precision attention on Blackwell by addressing quantization overhead, scale-induced causal leakage, forward–backward numerical consistency, and packed-document support. Its MXFP8 path leads the evaluated causal packed-document workloads in combined forward and backward attention latency, including quantization.

1. Introduction

  • Blackwell’s higher-throughput FP8 and FP4 matrix multiplication makes attention a low-precision opportunity, but quantization overhead, causal leakage, numerical consistency, and packed documents remain design constraints.Some evaluated implementations provide only a forward path for packed inputs, use tensor-wide FP8 scaling, or lack causal-mask support.Source: §1 · pp. 1–2
  • FLAMIA provides MXFP8 and mixed NVFP4/FP8 training paths, using token-block scaling to reuse 8-bit representations across matrix orientations.The mixed path uses NVFP4 for QK⊤ and its recomputation, and 8-bit operands for the remaining matrix multiplications.Source: §1 · p. 2
  • FLAMIA combines targeted scale-leakage safeguards with consistent forward and backward computations; single-step accuracy alone does not establish causal correctness or training stability.Default scaling retains within-block dependencies and a sequence-wide backward accumulator dependency; exact options remove them in supported batched causal calls.Source: §1 · pp. 1–2
  • Fused quantization, coordinated softmax and matrix-multiplication scheduling, a low-precision two-CTA backward pass, and GPU-based scheduling support packed documents without repacking.The kernels also include two bit-level exponential approximations for inference.Source: §1 · p. 2
  • 1.10–1.40× on B200 and 1.06–1.42× on B300: MXFP8 has the lowest combined forward and backward attention latency, including quantization, versus the fastest evaluated baseline per workload.In 1B-model pretraining on 100B tokens, MXFP8 and mixed NVFP4/FP8 lead only to marginal increase of training loss and change of downstream metrics.Source: §1 · p. 2

2. Background and Related Work

The background traces attention optimization from tiled computation and online softmax to GPU-specific kernels, then situates FLAMIA among microscaling formats and low-precision attention work. Prior work spans inference and training methods, including mixed-precision paths and efforts to improve backward consistency.

Background and Related Work

  • FlashAttention [2] uses tiled computation and online softmax to keep intermediate tiles on-chip, reducing memory traffic; later versions improve work partitioning and target Hopper and Blackwell GPUs.FlashAttention-2 [10] reduces non-matmul operations; FlashAttention-3 [11] and FlashAttention-4 [3] introduce architecture-tailored optimizations.Source: §2 · pp. 2–3
  • Microscaling represents each block with low-bit codes and a shared scale; MXFP formats use 32-element blocks, while NVFP4 uses 16-element blocks plus tensor-level scaling.Blackwell’s block-scaled matrix multiply consumes these operands directly and applies their scales during accumulation.Source: §2 · p. 3
  • Low-precision attention work covers quantized inference and training, including mixed FP4/FP8 training and MXFP8 forward and backward kernels.Attn-QAT [25] emphasizes matching backward score recomputation to the quantized forward pass and correcting precision assumptions in gradient calculations.Source: §2 · p. 3

4. Experiments

The experiments evaluate FLAMIA’s latency, numerical behavior, optimizations, inference approximations, and training outcomes. MXFP8 leads the combined forward-and-backward comparison on causal packed workloads, while additional experiments examine accuracy, causality, and training stability.

4. Experiments

  • The evaluation benchmarks B200 and B300 against attention implementations with supported ragged APIs, while some baselines lack causal masking, GQA, or packed-document backward support.Training-call timing includes forward, backward, and quantization; kernel-only timing excludes quantization and uses a different layout, so subtraction cannot isolate quantization overhead.Source: §4 · p. 10, §4.1 · p. 11

4.1 Performance on Packed Documents

  • 1.10–1.40× on B200 and 1.06–1.42× on B300: MXFP8 has the lowest training-call latency across nine causal packed workloads, including forward, backward, and quantization.Packed inputs use 16 query heads and d = 128; at T = 131072, peak attention-call memory is 2.7 GB for packed inputs versus 16.8 GB for aligned, padded buffers.Source: Table 1 · p. 10, §4.1 · pp. 11–12

4.2 Numerical Accuracy

  • Across tensor-wide, per-32-element, and token-block scaling, evaluated 8-bit configurations achieve output cosine similarity of approximately 0.998–0.999 against an FP32 reference.Within FLAMIA, mixed NVFP4/FP8 has higher O/dQ cosines than all-4-bit in both causal and non-causal settings; these metrics do not establish training stability.Source: §4.2 · p. 12, Table 4 · p. 12

4.3 Ablation Study

  • Fused quantization reduces preprocessing from 789 to 261 µs for MXFP8 and from 0.64 to 0.11 ms for NVFP4 at 4×4096 tokens.The associated launch counts fall from 17 to one and 11 to one, respectively; other ablations report gains from log-domain probability quantization, splitting QK⊤, and folding score scales.Source: §4.3 · p. 13
  • Removing longest-first ordering incurs a 10–65% causal-forward penalty; just-in-time CLC acquisition improves performance in 13 of 16 evaluated causal configurations.On non-causal workloads, CLC helps on the two skewed cases but loses on equal-length cases; removing document-aware backward ordering incurs up to 9.7% overhead.Source: §4.3 · p. 13

4.4 Bit-Level Exponentials at Inference

  • Schraudolph’s approximation gives MXFP8 inference speedups of approximately 1.03–1.17× with output-cosine change below 10^-4, while Mitchell-P has different causal and non-causal trade-offs.Mitchell-P achieves 1.05–1.23× speedups over NVFP4 on non-causal workloads, but causal latency is approximately 1.25–1.77× NVFP4; its comparison also changes probability precision.Source: §4.4 · pp. 13–14

4.5 Causality and Training Stability

  • Scale-induced leakage appears in several evaluated quantization routes, while only exact options with the causal dQ scale leave earlier outputs unchanged; forward–backward mismatches can destabilize training.The causal dQ scale has 8–11% backward overhead; mismatched dO representations, probability precision, diagonal tiling, or exponential evaluation produced reported training failures.Source: §4.5 · p. 14

E. Benchmarking Pitfalls

The section identifies several benchmark pitfalls, including incorrect causal FLOP accounting, gradient and layout errors, and toolchain-sensitive timings. It also describes runtime integration that supports checkpointing, compilation, and CUDA graph capture.

E. Benchmarking Pitfalls

  • The benchmark convention omits causal FLOP discounting, making causal throughput appear twice the true value and yielding reported rates above GPU peak.The authors say they controlled for several effects that had affected their own measurements.Source: E. Benchmarking Pitfalls · p. 25
  • cuDNN’s BF16 ragged backward returns incorrect gradients when softmax statistics are padded to batch shape rather than kept ragged.The reported dQ cosine is 0.62 for causal and 0.93 for non-causal attention.Source: E. Benchmarking Pitfalls · p. 25
  • cuDNN’s MXFP8 graph uses one V input layout for both passes, but forward and backward require V blocks along different axes.Using the forward layout in backward gives dQ and dK cosines of 0.947–0.962, while dV remains 0.9997.Source: E. Benchmarking Pitfalls · p. 25
  • The same kernel source ran 18.6–20.6% faster or slower across CuTe DSL versions despite bitwise-identical output, so timings are toolchain-specific.The authors gate speed changes bitwise per route rather than per kernel family.Source: E. Benchmarking Pitfalls · p. 25

F. System Integration

  • Both entry points are custom operators that support selective activation checkpointing and torch.compile tracing through a fake implementation.Checkpointing can save attention outputs instead of rerunning the forward and quantizers during backward.Source: F. System Integration · p. 25
  • Calls keep scales, work lists, and stochastic-rounding seeds on the GPU, enabling CUDA graph capture without host synchronization.Graph replay is bitwise equal to eager execution.Source: F. System Integration · p. 25

G. Training Details

The appendix points to shared training hyperparameters and specifies the software and runtime used for the training experiments.

G. Training Details

  • Training experiments use the shared hyperparameters listed in Table 8 and run with PyTorch 2.15.0dev20260928 and CUDA 13.1.The supplied passage does not enumerate the hyperparameter values.Source: G. Training Details · p. 25

H. Non-Causal Packed Documents

The appendix repeats packed-document measurements for non-causal workloads and reports separate B300 kernel measurements.

H. Non-Causal Packed Documents

  • Tables 9–11 repeat Section 4.1 measurements on non-causal versions of every workload, where Low Precision FA4 also runs.The supplied passage does not report the measurement values.Source: H. Non-Causal Packed Documents · p. 25
  • Table 9 covers training calls on packed, non-causal documents with 16 heads and d = 128, reporting time, throughput, and speedup over the fastest baseline.The caption defines speedup as the fastest baseline’s time over FLAMIA’s time.Source: Table 9 · p. 26

I. Kernels on B300

  • Tables 12 and 13 repeat the Section 4.1 kernel measurements on B300.The supplied passages do not state the measurement outcomes.Source: §I · p. 26

J. Batched Attention

Batched evaluations report training-call and inference-call comparisons, alongside results for isolated forward and backward kernels. The training-call advantage varies with cuDNN FP8 scale computation, while isolated-kernel leadership differs by direction and GPU.

J. Batched Attention

  • For batched-sequence inference, Schraudolph’s exponential makes the MXFP8 route 7–26% faster than cuDNN FP8 on every shape; at 1×32768 non-causal, Mitchell-P takes 4.64 ms versus 6.17 ms for NVFP4 and 4.68 ms for the fastest 8-bit baseline.Source: J. Batched Attention · p. 26
  • The backward kernel is fastest on all twelve B200 rows and six of twelve B300 rows, while the forward leads on three B200 rows and none on B300.On the other six B300 backward rows, cuDNN FP8 leads by 3–8%; operands are quantized outside the timed region.Source: J. Batched Attention · p. 26, Table 15 · p. 30, Table 16 · p. 31
  • Without a causal mask, HW-FP4 FA4’s 4-bit Direct-P forward is fastest on every MHA row with d = 128.This finding is for the non-causal setting.Source: J. Batched Attention · p. 28
  • 7.4%, 10.3%, and 112% are the median training-call speedups over cuDNN FP8 when all scales are pre-calibrated, input scales are computed in-call, and all scales are computed in-call, respectively.Computing every cuDNN scale in-call requires a second forward and backward pass.Source: J. Batched Attention · p. 26, Table 14 · p. 29

Related papers

Loading 2611.03344v1…