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

FLAMIA is a family of Blackwell attention kernels with MXFP8 and mixed NVFP4/FP8 training paths that use token-block scaling, causal-leakage safeguards, forward-backward numerical consistency, and GPU-side scheduling of packed documents, achieving 1.10–1.40× (B200) and 1.06–1.42× (B300) speedups on combined forward and backward attention calls.

Attention dominates Transformer training cost at long context, and although Blackwell tensor cores run FP8 and FP4 matmuls at twice and four times the BF16 rate, exploiting them requires solving quantization overhead, causal information leakage through shared scales, and numerical mismatch between forward and backward. Hardware block scaling follows the contraction axis, while attention's six logical operands participate in products with different contraction axes, so naive scaling duplicates quantized representations and makes memory-bound quantization a large share of training time. Shared scales can also let a future token change earlier outputs despite the causal mask, and the backward pass must reconstruct the forward's quantized intermediates consistently or training becomes unstable. Packed variable-length documents additionally require document-local scaling and load balancing without repacking.

FLAMIA builds on the warp-specialized FlashAttention-4 design and provides two default training paths: MXFP8 (E4M3 with token-block scales) and a mixed path using NVFP4 for QKᵀ and its recomputation with one scale per 16 head-dimension elements, plus E4M3 for the remaining products. Token-block scaling shares one power-of-two exponent across 32 consecutive tokens and the full head dimension so that a single 8-bit payload serves both K-major and MN-major orientations, while dS uses one exponent per 32×32 block. Causal leakage is confined by never letting an operand scale span more than 32 positions, rounding scale exponents up so no element clamps, and giving V per-key scales and a K fold on diagonal tiles, plus an optional causal dQ scale for the FP16 accumulator. Numerical consistency is enforced by recomputing scores from the same quantized Q and K, reproducing tiling, diagonal decisions and exponential evaluation, and computing D from the same quantized dO used in dP. Kernel-level work fuses Hadamard rotation, block statistics and code/scale generation into one launch, quantizes probabilities in the log domain for NVFP4, splits QKᵀ around P̂V to keep tensor cores busy, uses a low-precision two-CTA backward, and processes packed documents in place with a GPU-built longest-first work list plus Cluster Launch Control acquisition.

• MXFP8 achieves the lowest measured attention training-call latency, including quantization, on all nine causal packed-document workloads on both GPUs. §4.1 Performance on Packed Documents
• The 4-bit NVFP4 kernel achieves the lowest causal backward latency on every evaluated workload. §4.1 Performance on Packed Documents
• Leakage probes show that per-tensor scales, grouped 4-bit V along keys, HW-FP4 FA4 causal and FlashInfer NVFP4 KV change earlier outputs, while FLAMIA's exact options leave them unchanged; cuDNN MXFP8 and Low Precision FA4 change them slightly. §4.5 Causality and Training Stability9.8e-2 (not found in the paper)7.9e-2 (not found in the paper)2.8e-1 (not found in the paper)6.9e-6 (not found in the paper)1.7e-1 (not found in the paper)2.6e-1 (not found in the paper)1.1e-1 (not found in the paper)8.3e-2 (not found in the paper)
• Cosine similarity to an FP32 reference shows evaluated 8-bit configurations near 0.998–0.999, mixed NVFP4/FP8 at 0.9920/0.9906 for causal O/dQ, and all-4-bit at 0.9810/0.9456. §4.2 Numerical Accuracy
• In 1B-model pretraining on 100B tokens, MXFP8 and mixed NVFP4/FP8 lead to only a marginal increase in training loss and downstream scores close to the BF16 baseline. §4.6 Training validation
• Inference-time bit-level exponentials: Schraudolph's approximation gives MXFP8 speedups with negligible cosine change, and Mitchell-P reproduces exact-rounding codes far more often than Direct-P. §4.4 Bit-Level Exponentials at Inference

• Default scaling retains residual dependencies, a within-block residue of at most 31 positions and a sequence-wide dQ accumulator dependency; the exact options that remove them currently cover supported batched causal configurations, while packed documents, windows and the all-4-bit backward use the default scale (§5. Discussion and Conclusion; §4.5 Causality and Training Stability).
• The causal dQ scale that removes the accumulator dependency makes the backward 8–11 % slower at 32768 tokens and is therefore provided as an option (§4.5 Causality and Training Stability).
• The all-4-bit training path is not the default: it is slower than paths retaining most operands in 8 bits and reaches a slightly higher loss, with quantizing six copies accounting for 45–52 % of its backward time (§3.4 MXFP8 and Mixed NVFP4/FP8 Paths; §4.3 Ablation Study).
• The bit-level exponentials are not used in training because the backward must rebuild the forward's P̂ exactly, and on causal workloads Mitchell-P latency is approximately 1.25–1.77× that of NVFP4 (§3.7 Bit-Level Exponentials for Inference; §4.4 Bit-Level Exponentials at Inference).

• Token-block scaling shares one power-of-two exponent across 32 consecutive tokens and the full head dimension; because the exponent is common within a 32×d block, codes and scales stay valid under transposition and one E4M3 payload serves both orientations. §3.1 Token-Block Scaling
• Causal masking alone does not stop leakage through data-dependent scales; the design confines any operand scale to at most 32 positions, uses per-token NVFP4 scales for Q and K, and rounds scale exponents up so the scaled maximum is never clamped. §3.2 Causally Safe Scales
• For forward-backward consistency the backward recomputes scores from the same quantized Q and K, reproduces tiling, diagonal decisions, exponential evaluation and row-maximum scaling, and computes D from the same quantized dO used in dP. §3.3 Forward–Backward Numerical Consistency
• The fraction of dS entries rounded to zero tracked training degradation more closely than reported norm error in ablations, motivating scales derived from each 32×32 block's maximum absolute value. §3.3 Forward–Backward Numerical Consistency
• Fusing Hadamard rotation, block statistics and code/scale generation into one launch cuts quantization preprocessing from 789 to 261 µs for MXFP8 and from 0.64 to 0.11 ms for NVFP4 at 4×4096 tokens, reducing kernel launch counts from 17 to one and 11 to one. §4.3 Ablation Study
• Quantizing probability operands in the log domain for NVFP4 removes a per-entry division and a scaling pass, reported as the largest individual NVFP4 forward improvement at 14–19 %. §3.5.2 Softmax Throughput and MMA Scheduling
• Packed documents are quantized in place with per-token scale metadata, the load warp assembles each tile's 128×4 scale-factor atom in shared memory, and the schedule is built on the GPU once per microbatch with CLC used for skewed or causal calls. §3.6.2 Load-Balanced Scheduling
• Neither bit-level exponential is used in training, since the backward must rebuild the forward's P̂ exactly and a consistent Mitchell-P backward gave 1.60–1.64 ms versus 1.00 ms for the default NVFP4 path at 4×4096 tokens. §3.7 Bit-Level Exponentials for Inference