FLAMIA Brings Causally Safe Low-Precision Attention to Blackwell

Attention remains one of the largest costs in Transformer training, especially as context windows grow. Blackwell GPUs offer high-throughput matrix multiplication for FP8 and FP4 formats, but exploiting that hardware inside Attention requires more than converting tensors to a smaller representation. Quantization can add substantial overhead, shared scales can undermine causal semantics, and mismatched forward and backward paths can hurt training stability.

Yandex Research presents FLAMIA, or FLAsh MIcroscaling Attention, as a family of Blackwell-oriented kernels that integrates low-precision arithmetic with the broader Attention execution pipeline. The design includes an MXFP8 training path and a mixed path that uses NVFP4 for query-key score computation and its backward recomputation, while retaining FP8 for the remaining products. The authors avoid an all-FP4 design because its additional quantization cost and slightly worse loss behavior can offset the benefits of lower precision.

Key design choices

• Reusable quantized operands. FLAMIA uses power-of-two scales over blocks of 32 tokens and uses 32×32 blocks for score gradients. The resulting representations can serve different matrix orientations, reducing duplicate quantization, memory traffic, and intermediate movement.
• Causality-aware scaling. In causal Attention, a token must not receive information from future positions. Yet a shared scale calculated from a block containing future tokens can itself expose information, even when a causal mask is later applied. FLAMIA limits this risk with document-local scaling, saturation avoidance, and per-key scale folding. Exact causal options are also available for supported batched calls.
• Matched numerical paths. During backward recomputation, probabilities follow compatible quantization, tiling, scaling, and exponential conventions from the forward pass. This is important because a high single-step gradient similarity does not by itself guarantee stable long-running training.
• Hardware-aware execution. Quantization is fused into the kernels, while softmax work is overlapped with matrix multiplication. The low-precision backward path uses two cooperative thread arrays to improve GPU utilization. FLAMIA also handles variable-length documents directly in packed layouts through document-local scaling and workload-aware scheduling, avoiding a separate repacking step.

Results and implications

Across several causal packed-document workloads, the MXFP8 path accelerated the combined forward and backward Attention call, including quantization. Relative to the fastest evaluated baseline for each workload, the reported speedup was approximately 1.10–1.40× on B200 and 1.06–1.42× on B300. Pretraining experiments with Llama-style models were used to validate training stability and preservation of model quality.

The broader lesson is that low-bit acceleration is not simply a matter of selecting a smaller data type. Practical gains depend on coordinating quantization granularity, causal semantics, forward-backward numerical behavior, and GPU scheduling. FLAMIA illustrates how those concerns can be addressed together for Blackwell. Its results are most directly relevant to the tested packed-document workloads, however, and should not be interpreted as a fixed speedup for every model, sequence distribution, or training configuration.

Source: Hugging Face Daily Papers