FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell
Fedor Velikonivtsev, Mikhail Koltakov, Denis Kuznedelev
cs.LG
2026-11-03
MXFP8 and mixed NVFP4/FP8 attention training kernels for Blackwell: 1.06-1.42x faster fwd+bwd with quantization, near-identical 1B pretraining loss, no causal leakage from scales.
Blackwell tensor cores run FP8 at twice and FP4 at four times the BF16 peak, yet FlashAttention-4 ships in BF16. Three things block low-precision attention training. Quantization overhead: each of attention's six operands enters two products with different contraction axes, so conventional block scaling stores two quantized copies per operand, and quantization is memory-bound. Causal leakage: a causal mask does not stop data-dependent scale sharing, so a future token that raises a shared scale maximum changes how earlier tokens are reconstructed. Forward-backward inconsistency: mismatches in quantization, tiling, or exponential evaluation destabilize training even when single-step gradient similarity stays high. Existing low-bit kernels each miss something: FlashInfer offers only a packed forward, cuDNN FP8 uses tensor-wide scales, Low Precision FA4 lacks causal masking.
FLAMIA ships two training paths, MXFP8 and mixed NVFP4/FP8, built on FlashAttention-4's warp-specialized design.
The central idea is token-block scaling: 32 consecutive tokens share one power-of-two exponent across the full head dimension, so a block's codes and scales stay valid under transposition. Since 8-bit MMA accepts both K-major and MN-major layouts, one E4M3 payload serves both orientations and the duplicate copies disappear. NVFP4 covers only QK^T and its backward recomputation, which share a contraction axis; everything else runs in 8 bits.
Causal safety is surgical. No scale group spans more than 32 consecutive positions. Scale exponents are rounded up so no element saturates. A 4-bit V switches to per-key scales only on diagonal tiles that straddle the causal boundary, folding the scale into P̂; at 4096 tokens this costs 1.17× on the forward versus 2.33× if applied to every tile. The FP16 dQ accumulator defaults to sequence-wide statistics; a causal dQ scale option sizes each row's factor from visible keys only.
Forward-backward consistency comes from replication: the backward recomputes probabilities with the same quantized representations, tiling, diagonal decisions, and exponential evaluation as the forward. dS is quantized to E4M3 with one scale per 32×32 block, so one payload serves both dQ and dK.
Packed documents are processed in place. Quantization statistics never cross document boundaries, and scheduling is sorted by tile cost on the GPU: building the work list for 131,072 one-token documents takes 0.26 ms versus 50.5 ms for the reference construction. Peak memory at T=131072 is 2.7 GB versus 16.8 GB with aligned padding. An all-4-bit path exists but is not the default: 4-bit MMA has no MN-major form, six quantized copies take 45-52% of backward time, and loss is slightly worse. Two bit-level exponential approximations (Mitchell-P, Schraudolph) are inference-only, because the backward must rebuild forward probabilities exactly.
| Setting | Metric | Result |
| MXFP8, B200, 9 causal packed workloads | training-call speedup vs fastest baseline per workload | 1.10-1.40x |
| MXFP8, B300 | same | 1.06-1.42x |
| NVFP4 kernel-only causal backward | speedup vs fastest baseline | 1.22-2.01x (B200) / 1.17-2.00x (B300) |
| MXFP8 numerics | O/dQ cosine vs FP32 | 0.9988 / 0.9979 (cuDNN MXFP8: 0.9988 / 0.9978) |
| 1B Llama-3, 100B FineWeb-Edu tokens | training loss | BF16 2.392 / MXFP8 2.397 / NVFP4 2.406 |
| same | mean of five downstream tasks | 0.5832 / 0.5737 / 0.5814 |
Baselines include BF16 and FP8 FlashAttention-4, cuDNN in three precisions, FlashInfer, Low Precision FA4, HW-FP4 FA4, and PyTorch SDPA; support differs across them, so speedups are against the fastest available one per workload.
The leakage experiment is the sharpest result. On enwik8 with a 6-layer GPT, a leaking 4-bit configuration reaches validation loss 0.811-0.827 through its own kernel, lower than BF16's 0.852; evaluated with exact attention, the same weights score 0.888-0.940, against 0.862 for a model trained without the leak. Self-evaluation hides the damage entirely. Forward-backward mismatches are just as dangerous: handling diagonals over 256 queries in forward and 128 in backward keeps single-step dQ cosine at 0.945 yet blows the gradient norm from 0.5 to 2×10^5; an oversized dS scale pushes the fraction of entries rounded to zero from 0.4% to 29% and widens the loss gap from 0.0008 to 0.049. Probing existing implementations, default cuDNN MXFP8 and Low Precision FA4 leak at the 1e-6 level, HW-FP4 FA4's causal path at 1.7e-1, FlashInfer's NVFP4 KV at 2.6e-1; only FLAMIA's exact options measure 0.
Teams pretraining on Blackwell can swap the kernel in directly: the code is open with a TorchTitan integration, and 1.1-1.4x on attention changes no hyperparameters. In-place packed processing pays off more at long context: full 1B-model training reaches 96.8k tokens/s packed versus 84.3k with padding at 64k row length, and padding runs out of memory at 128k.
The measurement lesson outlives the kernel. Single-step cosine and self-evaluated loss cannot certify a low-bit attention implementation; causal-dependency probes and actual training can. The paper applies that probe to mainstream implementations and finds nonzero leakage in several defaults.
Defaults are not leak-free: a within-block dependency across 32 tokens and a sequence-wide backward accumulator dependency remain. The exact options that remove them cover only supported batched causal calls, and the causal dQ scale costs 6-9% of the training call.
Training validation is small: 1B parameters, 100B tokens, 2k context. The motivating scenario is long-context attention, yet no experiment goes beyond 2k. Downstream scores wobble: ARC Challenge falls from 0.3925 to 0.3618 (MXFP8) and 0.3677 (NVFP4) even though the mean stays close.
Speedups are relative to the fastest baseline per workload, and baseline support is uneven, so this is not a strict apples-to-apples comparison; on B300 non-causal workloads FLAMIA loses two of nine. The all-4-bit path is slower with worse loss, which says 4-bit peak throughput does not convert to attention speed while the backward is not tensor-core-bound.