Latest / Embodied AI 101 / FlashAttention-3: Fast & Accurate Attention with Asynchrony & Low-Precision
Transcript
- FlashAttention-3: Fast & Accurate Attention with Asynchrony & Low-Precision. Transformer self-attention is famously a big computational sink: its runtime and memory both grow quadratically with sequence length. For modern large language models (LLMs) pushing very long contexts (GPT-4 uses 128K tokens, and newer models like Llama 3 even 1M tokens), attention can easily become the bottleneck. In 2022–2023, the FlashAttention series showed how to dramatically speed up attention by clever blocking and memory-saving tricks. FlashAttention (v1) reorders the QKT computation into block-wise tiles, keeping only a small softmax state on-chip, which avoids writing huge attention matrices to slow GPU memory and yields ~2–4× speedups over naïve attention. FlashAttention-2 then tuned the tiling and parallelization further for GPUs, yielding even more speedups in practice (up to roughly 5–9× on Ampere GPUs in some regimes). These methods have already become standard in machine-learning libraries and dramatically enabled longer context windows in LLMs. However, new GPU architectures have created fresh opportunities (and challenges) that FlashAttention-2 did not exploit. Nvidia’s Hopper GPUs (H100 and beyond) introduced asynchronous tensor cores (WGMMA), a Tensor Memory Accelerator (TMA), and hardware FP8 support, which can double raw throughput for matrix-multiply at the cost of precision. Without leveraging these, FlashAttention-2 was still only sitting at ~35% of H100’s theoretical peak FLOPS. FlashAttention-3 (Shah et al., 2024) is a major overhaul designed to unlock those new capabilities. The key ideas are: overlap (“pipeline”) the GEMM and softmax work to utilize idle units, and enable accurate low-precision (FP8) attention via a randomized transform. Taken together, these achieve roughly 1.5–2.0× throughput over FlashAttention-2 in practice, pushing GPU utilization up dramatically (to ~75% of peak on H100). Below we break down the core innovations of FlashAttention-3, focusing on the algorithmic insights and why they matter for accelerating Transformers on modern hardware. New GPU Features on Hopper: WGMMA, TMA, FP8. FlashAttention-3 starts by embracing Hopper’s innovative hardware units. The white paper for NVIDIA’s H100 describes two critical new features:. WGMMA (Warpgroup Matrix Multiply-Accumulate) – This allows the new Hopper tensor cores to perform matrix multiplies with much higher throughput than the old instruction on Ampere. In fact, WGMMA can reach beyond the Hopper peak (since Hopper doubled tensor throughput). In practice, rewriting FlashAttention to use WGMMA lifts a matmul performance to roughly 70–75% of Hopper’s peak (up from ~2/3 on Ampere). In simple terms, one large-block QKT or PV product can now compute almost immediately at up to ~989 TFLOPS on an H100 (FP16), thanks to WGMMA, versus only ~2/3 of that with the old path. TMA (Tensor Memory Accelerator) – This is a special DMA engine for copying blocks between device memory (HBM) and each SM’s shared memory. Normally index calculations and predicate checks (for bounds) eat up registers. The TMA offloads all that indexing work, freeing up registers and reducing overhead. In practice this means FlashAttention can load or store entire tiles of Q/K/V blocks with minimal fuss, allowing bigger tile sizes and higher occupancy. The net effect is that attention’s data movement (reading Q/K/V and writing O) becomes much cheaper and more parallel. FP8 Support – Hopper doubles the tensor core throughput (e.g. from 989 TFLOPS FP16 to 1978 TFLOPS FP8) by using an 8-bit floating format. The hardware can natively do FP8 matrix multiplies at double the rate of FP16, at the cost of precision. (Since softmax/exponentials are still done in higher precision, FP8 mostly affects the QKT and PV multiplies.) Naively using FP8 can introduce quantization error, but FlashAttention-3 adds a fix to keep accuracy high (see below). The FlashAttention-3 implementation fully exploits all of these Hopper features via CUDA (using NVIDIA’s CUTLASS library to manage the new instructions). If you rewrite the standard FlashAttention kernel to just use WGMMA/TMA/FP8, the throughput jumps from ~350 TFLOPS (H100) to about 540–570 TFLOPS on a forward pass. In other words, hardware alone lifts a big chunk of speed (simply by using bigger, faster math tiles and faster memory moves). But the real secret sauce is pipelining: because WGMMA and TMA are asynchronous (they can run in parallel with each other and with other warps), the FlashAttention-3 authors introduce new overlap strategies to squeeze out even more speed. In particular, rather than compute each block’s matmuls then softmax sequentially, they interleave and overlap those operations across warps and blocks so that the GPU’s math units and special-function units stay busy at once. Overlapping GEMM and Softmax (Asynchrony). An attention layer consists of two big steps: (a) matrix multiplies (GEMMs) for Q×K and then P×V, and (b) softmax (which is mostly exponentials). It is tempting to think "well, the GEMMs are most of the FLOPs, so focus on them". But in fact on modern GPUs the reverse is a bottleneck: the non-matrix operations (especially the exponentials for softmax) are far slower than the rich tensor-core throughput. For example, on an H100 the FP16 matrix-multiply throughput is about 989 TFLOPS (using WGMMA) whereas special functions like exp are done by the CUDA multi-function unit at only 3.9 TFLOPS — roughly a 256× gap. In a single attention head of dimension 128, that means there are ~512× more matmul FLOPs than exponential FLOPs. Concretely, with FP16 the softmax exp could consume on the order of half the kernel time compared to the matmul. And if you go to FP8, the matmul rate doubles (nearly 1978 TFLOPS) while the exp rate stays the same, making the imbalance even worse. This observation suggests a new trick: run softmax in parallel with the matmul. Ideally, while some warps are grinding out a QK multiplication, other warps should be using the special-function units to start processing the previous block’s softmax. FlashAttention-3 achieves this with a two-tier pipeline:. Inter-warpgroup overlap (“ping-pong” scheduling): They divide the threadblock’s warps into two warp-groups. In the first group, they perform the GEMMs for block i (for example) and turn in the partial softmax results. Meanwhile, the second warp-group waits idle. Then a synchronization barrier flips: warpgroup 1 switches to computing softmax (starting with the block just finished), and warpgroup 2 takes over the GEMMs for the next block (i+1). In effect, one group is always doing matmul while the other is doing softmax in the background. A diagram (not shown here) resembles a ping-pong: while block A’s matmul and softmax are pipelined across the groups, block B’s matmul overlaps block A’s softmax. This simple ping-pong schedule already boosts performance noticeably: in their tests, it raises the FP16 forward throughput from about 570 TFLOPS to around 620 TFLOPS on H100 (head-dim 128, seqlen 8K). Intra-warpgroup overlap: They go further by pipelining within each warp-group. Even while a warp-group is busy on some GEMM, it can allow the first few warps to start computing part of the next softmax (e.g. computing the exponentials or partial reduction on previously loaded data). Again picture multiple interleaved blocks: as soon as a piece of QK is done, some warps start the softmax even while others finish the outlying multiplications. This “inside-out” pipelining yields another bump: throughput climbs from ~620 to about 640–660 TFLOPS for the FP16 forward pass. The cost is a bit more register usage, since now each warp must hold both some GEMM accumulators and parts of the softmax state, but it’s a favorable tradeoff. Together, these overlapped pipelines keep the tensor cores and special-function units mostly busy. The GPU isn’t sitting idle waiting for exps; instead, exponentials run “in the shadow” (background) of the multiply-heavy operations. In short, by explicit warp scheduling and pipelining, FlashAttention-3 interleaves matmul and softmax to utilize the Hopper’s asynchrony. (As a footnote: modern GPUs’ hardware warp schedulers already do some of this scheduling by default, but FlashAttention-3 does even more via manual barriers and specialized warp roles.). Incoherent Processing for FP8 Accuracy. Using FP8 arithmetic yields huge speedups — but only if the model can tolerate the low precision. Naively quantizing Q/K/V into FP8 can cause large errors, especially when some elements (outliers) are much larger than others. FlashAttention-3 combats this with a technique borrowed from quantization research called “incoherent processing.” The idea is to randomly rotate or “spread out” the values before quantizing, so that extreme spikes (outliers) don’t all land in the same coordinate corner of FP8. Concretely, FlashAttention-3 multiplies the queries and keys by a random orthogonal matrix per head before computing attention. They implement this efficiently using a Hadamard transform (with random signs), which is an O(d log d) operation per head (rather than O(d²)). Since Rotary embeddings and similar steps are already memory-bound, they fuse this transform seamlessly (at essentially no extra cost). This random rotation “smears” any outliers across the vector, reducing the peak magnitude in any one dimension. In experiments, this trick paid off. The authors tested on synthetic data where 0.1% of Q/K entries were huge outliers (simulating real LLM activations). With incoherent processing, the FP8 quantization error dropped by a factor of 2.6× compared to plain quantization. (In terms of attention, that means FlashAttention-3's FP8 results are far closer to the full-precision result than a baselined FP8 would be.) In practice, they verify this leads to using FP8 while keeping accuracy high. Performance Results. Putting it all together, FlashAttention-3 achieves major speedups. On an NVIDIA H100 GPU:. FP16 mode (using WGMMA and overlap, but doing math in FP16): FlashAttention-3 runs about 1.6–2.0× faster than FlashAttention-2. This matches the theory: if FA2 was hitting only ~35% of peak, and FA3 hits ~75%, that’s roughly a 2× gain. The blog reports around 620–660 TFLOPS end-to-end (H100, head=128, large seq), versus ~350 TFLOPS for FA2, consistent with the pipelining gains above. FP8 mode (using FP8 math with incoherent transform): Performance skyrockets. Because FP8 doubles core throughput, they measure forward speeds near 1.2 PFLOPS (petaFLOPS). In other words, FlashAttention-3 is saturating about 75% of the 1.6 PFLOPS theoretical (1978 TFLOPS) FP8 peak on H100. And thanks to the random rotation trick, this runs with very little loss of accuracy – about 2.6× smaller error than the baseline FP8 attention algorithm. For reference, these gains are also reflected in end-to-end inference: the authors show that a typical attention layer runs roughly 1.6–1.8× faster (measured in milliseconds per layer) compared to FA2 for FP16, and about 2.6× faster in FP8 mode (since we gained both from pipelining and from doubling compute). In practice this means long-context LLM operations (training or autogressive sampling) can run nearly twice as fast just by switching to FlashAttention-3 kernels. Implications and Usage. FlashAttention-3 tackles one of the toughest low-level computational challenges in Transformers by matching the algorithm to the hardware. On Hopper GPUs (and presumably similar architectures like H200), it provides a new “attention trampoline” that the rest of the model can run on. By overlapping tensor-core math with softmax and enabling FP8, it effectively neutralizes the softmax bottleneck and doubles down on the hardware’s matrix power. For roboticists and ML practitioners, this translates to concrete benefits:. Longer context windows at scale: Faster attention means you can afford to train or use longer sequences. For tasks like multi-modal video analysis, embodied planning over long time-horizons, or any use of massive memory stores, this matters directly. An H100 running FlashAttention-3 can handle 2–3× the token length in the same time as before. Lower training/inference cost: Better GPU utilization directly lowers compute cost. If an attention-heavy model runs in half the time, that halves GPU-hours. This frees up resources for more data or bigger models – or simply reduces the bill. FP8 viability: Many new LLMs train and run in FP8 by default. FlashAttention-3 makes one of FP8’s core components (scaled dot-product attention) safely accurate, so practitioners can use FP8 from end to end with confidence. In fact, the authors note FlashAttention-3 has already been integrated into libraries (e.g. a PyTorch FlexAttention backend) to automatically speed up models. In short, FlashAttention-3 is an important “under-the-hood” advance. It is largely transparent to model users (just swap in the optimized attention kernel), but offers a big kick in efficiency. The authors have released both the code and the full paper, and it’s likely to become part of the default attention primitive in frameworks on Hopper hardware. For anyone training or deploying large Transformers on NVIDIA GPUs, it’s a very practical tool: faster runs, longer context, and good accuracy even in 8-bit. Sources: The above is based on the FlashAttention-3 paper and accompanying blog by Shah et al. (2024), which detail the algorithmic changes and report benchmarks. FlashAttention-3 builds on the prior FlashAttention papers and exploits H100-specific features (WGMMA, TMA, FP8) with innovative pipelining. The cited results come from these sources’ experiments and analysis.