FLAMIA is a set of attention kernels for NVIDIA's Blackwell GPUs, aimed at speeding up Transformer training at long context lengths. The paper says it does this while keeping training stable.
Precision paths: It offers an MXFP8 path and a mixed NVFP4 and FP8 path. The paper notes that going all-FP4 adds quantization overhead and slightly worsens loss.
Speed: The paper reports 1.10 to 1.40 times faster on B200 and 1.06 to 1.42 times faster on B300 for the combined forward and backward attention call, quantization included, against the fastest baseline for each workload.
Causality: Shared quantization scales can leak future-token information, so the paper uses local scales and other safeguards to mitigate this.
Packed documents: Variable-length documents are handled directly in packed layouts, without repacking.
Sign in to comment.
TL;DR of the paper, generated automatically.
FLAMIA is a set of attention kernels for NVIDIA's Blackwell GPUs, aimed at speeding up Transformer training at long context lengths. The paper says it does this while keeping training stable.