FlashAttention-3 on Hopper

Contents 1 Introduction and Motivation 1.1 Why attention is the bottleneck 1.2 The Hopper problem: FA2 leaves most of the H100 idle 1.3 The three ideas of FlashAttention-3 1.4 The code base: from CUTLASS C++ to the CuTe DSL 1.5 Organization of this post 2 The Hopper GPU Architecture: Features Used by FA3 2.1 Chip-level anatomy of the H100 2.2 The SM: warps, warpgroups, and the asynchronous datapaths 2.3 TMA: the Tensor Memory Accelerator 2.3.1 Multicast and clusters 2.4 mbarriers: arrive/wait synchronization in shared memory 2.5 WGMMA: asynchronous warpgroup matrix multiply 2.6 The async proxy and memory fences 2.7 Register reallocation: setmaxnreg 2.8 Threadblock clusters, persistent kernels, and occupancy 2.9 The SFU: why softmax is expensive 2.10 Putting it together: the FA3 hardware checklist 3 Background: Attention, Online Softmax, FA1/FA2 3.1 Attention and the quadratic intermediate 3.2 Safe softmax and its streaming state 3.2.1 The block-combine interpretation 3.2.2 Consequences for implementation 3.3 FlashAttention-1: IO-aware exact attention 3.3.1 Backward recomputation 3.3.2 First-generation work partitioning 3.4 FlashAttention-2: parallelism and work partitioning 3.4.1 Fewer non-matrix operations 3.4.2 A query block becomes an independent thread-block task 3.4.3 Split-Q warp partitioning 3.4.4 Performance envelope and the remaining Hopper gap 3.5 FA1, FA2, and FA3 in one comparison 3.6 Kernel anatomy inherited by FA3 4 The FA3 Forward Pass: Algorithm 4.1 Work decomposition 4.2 Online-softmax state and invariant 4.3 Producer–consumer warp specialization 4.3.1 Pipeline-stage ownership 4.4 Pingpong scheduling: hiding softmax between warpgroups 4.5 Intra-warpgroup overlapping: the 2-stage GEMM–softmax pipeline 4.5.1 Prologue, steady state, and drain 4.6 Softmax details: rescaling, correction, and LSE 4.7 Causal masking and the split loop 4.8 Per-tile computation and data movement 4.9 Tile sizes, head-dimension variants, and register budgets 4.10 Epilogue 5 The FA3 Forward Pass: CuTe DSL Implementation 5.1 A note on the CuTe DSL 5.2 Implementation map 5.3 Kernel configuration 5.3.1 Compile-time variant selection 5.4 Shared memory plan 5.4.1 Tensor views and WGMMA partitions 5.5 Kernel entry: pipelines and role dispatch 5.5.1 Pipeline cursors and phase changes 5.6 The producer: load 5.6.1 From a scheduled tile to copy coordinates 5.7 The consumer: mma 5.7.1 Accumulator ownership and initialization 5.8 Core operation: one \(n\)-block with intra-warpgroup overlap 5.8.1 Why each wait and release is safe 5.9 Pingpong in code: the scheduler barriers 5.10 Softmax module 5.11 Epilogue in code 5.12 Differences from the paper’s CUTLASS C++ kernels 6 FP8 Forward: Layouts, Block Quantization, Incoherent Processing 6.1 FP8 formats and accumulation 6.2 The layout problem: K-major operands and the in-kernel V transpose 6.3 Block quantization 6.4 Incoherent processing: the Hadamard transform 6.5 FP8 kernel schedule and performance shape 7 The FA3 Backward Pass: Algorithm 7.1 Gradient math 7.1.1 The softmax Jacobian, step by step 7.1.2 Why \(D\) can be computed without \(\mathbf{P}\) 7.1.3 Why LSE alone suffices to recompute \(\mathbf{P}\) 7.1.4 Gradients through LSE itself 7.1.5 Where the scale \(\alpha\) is applied 7.2 Tiling and the dQ problem 7.2.1 Why KV-parallel and not Q-parallel 7.2.2 Causal masking: trapezoidal iteration ranges 7.2.3 Three ways to resolve the dQ race 7.3 The preprocessing and postprocessing kernels 7.3.1 Preprocess: D, LSElog2, and zeroing dQaccum 7.3.2 Postprocess: dQaccum \(\to\) dQ 7.4 Warp specialization in the backward 7.4.1 Named barriers of the backward 7.4.2 SMEM budget 7.5 Scheduling the five GEMMs 7.5.1 Operand layouts: who is SS, who is RS 7.5.2 swapAB: transposing the whole GEMM instead of the data 7.6 dQ accumulation across thread blocks 7.7 Backward tile configurations 7.8 Causal masking, varlen, and other mainloop features 7.9 Backward performance 8 The FA3 Backward Pass: CuTe DSL Implementation 8.1 Configuration surface 8.2 SMEM layouts supporting two access orientations 8.3 Host-side: TMA atoms, schedulers, semaphores, PDL 8.4 Role dispatch: two producer warps, two consumer warpgroups 8.5 The producer load loop 8.6 The consumer inner loop: five GEMMs and two pointwise stages 8.7 The dQaccum store agent 8.8 Epilogue: dK/dV stores, and the GQA accumulation path 8.9 The postprocessing kernel 8.10 Preprocess kernel code 9 Tile Scheduling, Variable-Length Sequences, and Masking 9.1 What the tile scheduler controls 9.1.1 Where scheduling appears in the kernel 9.2 The scheduling problem: causal masking skews tile cost 9.2.1 Static, persistent, and dynamic scheduling are not synonyms 9.3 Variable-length sequences (varlen) 9.3.1 How the varlen coordinate map works 9.4 Scheduling and masking: division of responsibility 9.4.1 Correctness invariants for scheduler changes 9.5 GQA and PackGQA 9.6 Masking machinery 9.7 Paged KV and other serving features 10 Flash-Decoding: Split-KV and the Combine Kernel 10.1 Why decoding underutilizes the GPU 10.2 Parallelization over the KV sequence 10.2.1 The combine math 10.3 The CuTe DSL combine kernel 10.4 How the mainloop produces splits 10.5 Choosing num_splits 10.6 Interaction with GQA, paged KV, and varlen 10.7 Performance character and caveats 11 PackGQA: Packing Query Heads into the Query Tile 11.1 Motivation: filling the M tile 11.1.1 What packing changes, and what it preserves 11.2 The packed layout: a zero-copy CuTe transformation 11.2.1 TMA with a packed mode 11.3 The gather/scatter path: PackGQA 11.4 Masking and per-row head indices 11.5 Enablement criteria for PackGQA 11.6 Composition with split-KV, and provenance 12 Performance Analysis and Benchmarks 12.1 Forward pass, BF16 12.2 Backward pass 12.3 FP8 forward 12.4 An issue-time consistency check 12.5 Interpreting throughput curves 12.6 End-to-end implications of attention speedups 13 Caveats and Pitfalls 13.1 Numerical caveats 13.2 Resource and performance caveats 13.3 Minimum validation matrix for kernel changes 14 Sources and Further Reading How the sources were used ...

August 29, 2026 · 162 min · 34331 words · Li Cao