FlashAttention-3 on NVIDIA Hopper: Asynchronous TMA, FP8 Tensor Cores, and Ping-Pong GEMM
The Evolution of Attention: From Memory Wall to Hardware Asynchrony
In standard transformer architectures, the multi-head attention mechanism exhibits quadratic time and memory complexity $mathcal{O}(N^2)$ with respect to sequence length $N$. For years, computing attention required materializing massive intermediate attention matrices ($S = QK^ op$ and $P = ext{softmax}(S)$) into slow GPU High Bandwidth Memory (HBM), creating an severe memory bandwidth bottleneck.
In 2022, Tri Dao and colleagues introduced FlashAttention, introducing online softmax tiling to compute exact attention entirely within fast GPU On-Chip Shared Memory (SRAM), eliminating intermediate HBM reads and writes. FlashAttention-2 (2023) further optimized work partitioning across thread blocks and warps.
However, NVIDIA's Hopper architecture (H100/H800 GPUs) introduced revolutionary hardware features that previous FlashAttention algorithms could not fully utilize. In mid-2024, Tri Dao, Jay Shah, and Together AI researchers published FlashAttention-3 (*FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision*), exploiting Hopper-native hardware primitives to achieve up to 2x faster attention and reaching up to 85% of theoretical GPU tensor core utilization.
Hopper Hardware Innovations Exploited by FlashAttention-3
FlashAttention-3 redesigns the attention kernel around three foundational hardware mechanisms introduced in NVIDIA Hopper:
FlashAttention-3 Pipeline on NVIDIA H100:
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β 1. Asynchronous Data Movement (TMA) β
β Global Memory (HBM3) ββ[Hardware TMA]ββ> Shared Memory (SRAM)β
β (Zero CPU/SM thread overhead) β
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
β 2. Warpgroup GEMM (WGMMA) β
β 128 Threads (4 Warps) execute unified matrix multiplication β
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
β 3. Ping-Pong Kernel Scheduling β
β Warp Group A: Computes GEMM (Q @ K) β
β Warp Group B: Overlaps Softmax & Rescaling β
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
#### 1. Tensor Memory Accelerator (TMA) In prior architectures (Ampere / A100), copying tensor tiles from Global HBM to Shared Memory required compute threads to execute load/store instructions manually, consuming valuable Streaming Multiprocessor (SM) register files and compute cycles.
Hopper introduced the Tensor Memory Accelerator (TMA)βa dedicated hardware copy engine that transfers multi-dimensional tensor tiles asynchronously between HBM and SRAM: - Compute threads issue a single TMA instruction and immediately continue other math operations. - The TMA engine handles memory addressing, boundary padding, and byte layout in hardware. - Memory copies run fully asynchronous, freeing registers for numerical calculations.
#### 2. Warpgroup Matrix Multiply-Accumulate (WGMMA) In Ampere, matrix multiplications were executed at the warp level (32 threads). Hopper introduces Warpgroupsβclusters of 4 warps (128 threads) executing unified WGMMA instructions directly against data stored in Shared Memory without copying tiles into local registers first. This quadruples instruction throughput and slashes register pressure.
#### 3. Interleaved "Ping-Pong" Kernel Scheduling In FlashAttention-2, computing attention required alternating between matrix multiplication (GEMM) and non-GEMM operations (softmax, exponentiation, reduction). Because GEMM instructions use Tensor Cores while softmax instructions use Vector/ALU cores, parts of the GPU were perpetually idle.
FlashAttention-3 introduces Ping-Pong scheduling: - It partitions threads into two specialized warp groups. - While Warp Group 0 executes a matrix multiply ($Q imes K^ op$) on the Tensor Cores, Warp Group 1 concurrently executes softmax reductions and rescaling on the ALU cores. - The groups swap roles in a continuous ping-pong rhythm, achieving near-100% concurrent utilization of both Tensor Cores and Vector ALUs.
Low-Precision FP8 Attention with Block Quantization
To double compute throughput, Hopper introduced hardware support for 8-bit floating-point (FP8) Tensor Cores. However, naive FP8 attention suffers from numerical instability: the dynamic range of attention matrices frequently overflows or underflows 8-bit representations, degrading model perplexity.
FlashAttention-3 introduces Block Quantization with Incoherent Processing: - Attention matrices are dynamically scaled and quantized in compact blocks (e.g., $128 imes 128$) rather than across the entire sequence. - It applies randomized Hadamard transformations to eliminate outlier feature dimensions before FP8 quantization. - This preserves 16-bit numerical accuracy while unlocking the full 2x theoretical throughput of FP8 hardware.
Benchmark Results on NVIDIA H100
| Algorithm | Precision | Speed (TFLOPs/s) | % of H100 Theoretical Peak | Speedup vs FA2 | | :--- | :---: | :---: | :---: | :---: | | PyTorch Standard Attention | FP16 | ~180 TFLOPs | ~18% | 0.25x | | FlashAttention-2 | FP16 | ~350 TFLOPs | ~35% | 1.0x (Baseline) | | FlashAttention-3 | FP16 | ~650 TFLOPs | ~66% | 1.85x | | FlashAttention-3 (FP8) | FP8 | ~1,200 TFLOPs | ~85% | 3.42x |
On FP16 workloads, FlashAttention-3 runs nearly twice as fast as FlashAttention-2. On FP8 workloads, throughput reaches an unprecedented 1.2 PFLOPs/s on a single H100 SXM5 GPU.
Production Impact for LLM Training and Serving
Accelerating attention directly benefits long-context models: - Long-Context Context Windows: Enables efficient training of 100k+ token context windows without spending 80% of compute time in attention kernels. - Serving Throughput: Lowers prefill latency in high-throughput enterprise serving engines like Together AI and cloud infrastructure providers like RunPod. - Synergy with PagedAttention: FlashAttention-3 handles the compute kernel acceleration while vLLM PagedAttention manages KV-cache memory paging.
For implementation details and CUDA source kernels, review the FlashAttention GitHub Repository β.