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

1 Introduction and Motivation

FlashAttention-3 (FA3) [1] is the third generation of the exact, IO-aware attention algorithm, specialized for the NVIDIA Hopper (H100) architecture. Its predecessors reorganized attention around the memory hierarchy; FA3 reorganizes it around Hopper’s asynchronous execution facilities.

1.1 Why attention is the bottleneck

The attention mechanism is a core computation in the Transformer. For a single head, given query, key and value matrices \(\mathbf{Q}, \mathbf{K}, \mathbf{V}\in \mathbb{R}^{N \times d}\) (sequence length \(N\), head dimension \(d\)), attention computes

where \(\alpha\) is a scaling factor, typically \(1/\sqrt{d}\), and the softmax is applied row-wise. The arithmetic cost is \(O(N^2 d)\) and — naively — the memory cost is \(O(N^2)\) for materializing \(\mathbf{S}\) and \(\mathbf{P}\) in GPU main memory (HBM). As context lengths have grown from 2K to 128K and toward millions of tokens, this quadratic term dominates both training and inference cost, and the \(O(N^2)\) HBM traffic of the naive implementation dominates its runtime.

FlashAttention [2] (2022) removed the \(O(N^2)\) memory traffic by tiling: it streams blocks of \(\mathbf{K}\) and \(\mathbf{V}\) through on-chip SRAM (shared memory), maintains a running online softmax [4, 5] normalization, and never writes \(\mathbf{S}\) or \(\mathbf{P}\) to HBM. FlashAttention-2 [3] (2023) improved the work partitioning: it parallelized over the sequence-length dimension, minimized non-matmul FLOPs, and reduced shared-memory round trips, reaching ∼70% of peak on the A100 (Ampere) generation.

1.2 The Hopper problem: FA2 leaves most of the H100 idle

When FlashAttention-2 is compiled unmodified for H100, it reaches only about 35% of the H100’s peak throughput (roughly 350 TFLOP/s out of 989 TFLOP/s BF16) [1, 9]. The reason is architectural. Hopper’s headline throughput comes from features FA2 does not exploit:

• WGMMA (warpgroup matrix multiply–accumulate, wgmma.mma_async): asynchronous Tensor Core instructions issued by a warpgroup of 128 threads, which can read operands directly from shared memory and are required to reach the H100’s full Tensor Core throughput. The older mma.sync path used by FA2 achieves only about \(2/3\) of peak [1, 8].

• TMA (Tensor Memory Accelerator): a per-SM copy engine that performs bulk, multidimensional tensor transfers between global and shared memory from a single-thread instruction, freeing all other threads (and their registers and issue slots) for computation.

• Asynchronous execution throughout: both WGMMA and TMA complete asynchronously, but use different completion mechanisms. TMA loads can signal an mbarrier (an arrive/wait barrier in shared memory), while WGMMA completion is managed with commit groups and wgmma.wait_group. Together these mechanisms enable software pipelines in which memory copies, Tensor Core GEMMs, and pointwise work such as softmax overlap.

• FP8 Tensor Cores: doubling throughput to ∼1979 TFLOP/s — but only for kernels able to satisfy FP8 WGMMA’s stringent operand-layout requirements.

Attention is harder to accelerate than plain GEMM because of the softmax sandwiched between the two matrix multiplies. Softmax is performed by the multi-function units (MUFU, on the SFU datapath), whose throughput for MUFU.EX2 (the base-2 exponential) is a small fraction of Tensor Core throughput: for a head dimension of 128, the exponentials alone would take ∼50% as many cycles as the matmuls if executed serially with them [1]. FA3’s central contribution is a set of scheduling techniques that hide softmax (and memory) latency behind the asynchronous GEMMs.

1.3 The three ideas of FlashAttention-3

FA3 [1] introduces three techniques:

1. Producer–consumer asynchrony (warp specialization). The warps of a forward-pass thread block are split into a producer warpgroup, which primarily issues TMA loads, and one or more consumer warpgroups, which perform WGMMA and softmax work. The two sides communicate through a circular shared-memory buffer guarded by mbarriers. Hopper’s dynamic register reallocation (setmaxnreg) shrinks the producer’s register footprint (it barely needs registers, since TMA computes addresses in hardware) and grows the consumers’.

2. Hiding softmax under GEMMs. Two complementary schedules: pingpong scheduling between two consumer warpgroups (while warpgroup A does softmax, warpgroup B’s GEMMs occupy the Tensor Cores, and vice versa, coordinated with named barriers), and intra-warpgroup 2-stage pipelining (within one warpgroup, the softmax of iteration \(j\) overlaps with the WGMMAs of iteration \(j+1\)).

3. Low-precision FP8 with error control. FP8 doubles Tensor Core throughput but raises two problems: layout (FP8 WGMMA requires the second operand to be k-major, forcing an in-kernel transpose of \(\mathbf{V}\)) and accuracy (outlier features inflate quantization error). FA3 answers with block quantization (one scale per tile of \(\mathbf{Q},\mathbf{K},\mathbf{V}\)) and incoherent processing (multiplying \(\mathbf{Q}\) and \(\mathbf{K}\) by a random orthogonal Hadamard-based matrix to spread outliers), achieving \(2.6\times\) lower error than a baseline per-tensor FP8 attention.

The headline results on H100 SXM5: up to 840 TFLOP/s (85% utilization) for BF16 forward, 1.3 PFLOP/s for FP8 forward, \(1.5\)–\(2.0\times\) faster than FlashAttention-2, with the backward pass \(1.5\)–\(1.75\times\) faster than FA2’s [1].

1.4 The code base: from CUTLASS C++ to the CuTe DSL

The kernels evaluated in the FA3 paper were written in CUTLASS/CuTe C++ and live in the hopper/ directory of the official repository [12] (files such as mainloop_fwd_sm90_tma_gmma_ws.hpp, flash_fwd_kernel_sm90.h, mainloop_bwd_sm90_tma_gmma_ws.hpp). Since 2025 the repository also contains a full reimplementation in the CuTe DSL — NVIDIA’s Python-embedded kernel DSL shipped with CUTLASS 4.x [21] — under flash_attn/cute/. This Python implementation (packaged as flash-attn-4; the SM90 kernels are direct ports of the FA3 Hopper kernels, and its file headers credit the FA3 authors) is the reference code base for this post:

All code listings in this post are excerpts of these files (trimmed and annotated where noted; line numbers refer to the repository state of July 2026, commit 2ee80234). Where the CuTe DSL port deviates from the C++ kernels described in the paper, the difference is identified explicitly (see also Section 5.12).

1.5 Organization of this post

Chapter 2 builds up the Hopper architecture knowledge assumed everywhere else. Chapter 3 compresses attention math, online softmax, and FA1/FA2 into the minimum needed to appreciate FA3. Chapters 4 and 5 cover the forward pass — first the algorithm as in the paper, then the CuTe DSL code path by path. Chapter 6 covers the FP8 forward. Chapters 7 and 8 do the same for the backward pass. Chapter 9 covers tile scheduling, variable-length sequences, and masking. Chapters 10 and 11 cover the two decode-oriented techniques — split-KV (Flash-Decoding) with its combine kernel, and PackGQA query-head packing. Chapter 12 collects the measured numbers, and Chapter 13 is a compendium of caveats: numerics, determinism, register pressure, bank conflicts, and supported shapes.

2 The Hopper GPU Architecture: Features Used by FA3

This chapter provides a detailed account of the H100 hardware features used by the FA3 kernels. Readers familiar with TMA, WGMMA, mbarriers, warpgroups, setmaxnreg and the async proxy can skim Table 2.1 and move on.

2.1 Chip-level anatomy of the H100

Table 2.1: NVIDIA H100 SXM5 specifications relevant to FA3 [18, 19].

Three ratios in this table drive kernel design:

• FLOPs : bandwidth. \(989\,\text{TFLOP/s} / 3.35\,\text{TB/s} \approx 295\) FLOPs per byte. An attention kernel at large \(N\) has arithmetic intensity \(O(d)\) per byte of \(\mathbf{K}/\mathbf{V}\) traffic per query block; with tiling it is compute-bound for the head dimensions used in practice. The primary objective is therefore to sustain Tensor Core utilization rather than reduce memory traffic further.

• Tensor core : SFU throughput. Special functions issue at 16 ops/cycle/SM — \(16 \times 132\,\text{SMs} \times 1.83\,\text{GHz} \approx 3.9\) TFLOP/s chip-wide, \(256\times\) less than the Tensor Cores [9] — while BF16 Tensor Cores deliver 4096 FLOP/cycle/SM. For \(\mathbf{Q}\mathbf{K}^\top\) and \(\mathbf{P}\mathbf{V}\) at head dimension \(d = 128\), each attention score requires \(2\times 2\times 128 = 512\) FLOPs of matmul and one exponential: \(512/4096 = 0.125\) cycles of Tensor Core versus \(1/16 = 0.0625\) cycles of MUFU per element — i.e., softmax would consume ∼50% additional cycles if serialized after the GEMMs [1]. This ratio motivates the emphasis on overlapping softmax with GEMMs.

• SMEM : register file. 228 KB SMEM vs. 256 KB registers per SM. FA3 keeps large accumulators (\(\mathbf{O}\) tile, \(\mathbf{dK}/\mathbf{dV}\) tiles) in registers and streams operands through SMEM; both resources are almost fully committed.

2.2 The SM: warps, warpgroups, and the asynchronous datapaths

Each Hopper SM is partitioned into 4 processing blocks (sub-partitions), each with its own warp scheduler, register file bank, Tensor Core, and SFU pipeline. A warp is 32 threads; each sub-partition schedules its resident warps independently. Hopper introduces the warpgroup: 4 contiguous warps (128 threads) that collectively execute warpgroup-wide instructions such as WGMMA. The programming contract is four contiguous warps; code should not infer a fixed one-warp-per-sub-partition placement from that definition. CUDA exposes WGMMA through PTX; in the CuTe DSL, a warpgroup consists of 128 consecutive threads, and helper functions like utils.canonical_warp_group_idx() compute \(\lfloor\texttt{tid}/128\rfloor\).

The execution resources relevant to FA3, and their asynchrony:

The first two rows are Hopper-specific asynchronous facilities. In PTX’s memory model, their shared-memory accesses are associated with an asynchronous proxy, whereas ordinary loads and stores use the generic proxy. A proxy is a memory-consistency classification, not a promise of a physically separate cache or datapath; the distinction matters because cross-proxy accesses require explicit ordering (Section 2.6).

2.3 TMA: the Tensor Memory Accelerator

TMA is a dedicated address-generation and copy facility. A single thread issues cp.async.bulk.tensor.{1..5}d with a tensor map (TMA descriptor): a 128-byte object, created on the host (or device), that encodes base address, tensor shape and strides, the box (tile) shape, the element type, and the SMEM swizzle mode. The TMA unit then:

1. generates all global addresses itself (no per-thread address arithmetic and no registers consumed for addressing, enabling register-starved producer warps);

2. handles out-of-bounds reads by filling zeros (FA3 relies on this: loading a \(\mathbf{K}\) tile that overhangs the end of the sequence is safe — e.g. flash_fwd_sm90.py notes that TMA fills zeros for n_block=-1, whereas the corresponding cp.async access would be invalid);

3. writes into SMEM in a swizzled layout (up to 128-byte swizzle) chosen to make subsequent WGMMA reads bank-conflict-free;

4. on completion, updates the transaction count of an mbarrier, which is how consumers learn the data has arrived.

TMA also performs SMEM\(\to\)GMEM stores (cp.async.bulk.tensor with .global destination, used for the \(\mathbf{O}\)/\(\mathbf{dK}\)/\(\mathbf{dV}\) epilogues), 1-D bulk copies (cp.async.bulk, used for LSE loads in the backward), and reductions: cp.async.bulk.reduce.add atomically adds an SMEM buffer into global memory — FA3’s CuTe DSL backward uses precisely this for accumulating dQ (Chapter 8), instead of per-element atomicAdd.

In the CuTe DSL, descriptor creation is cpasync.make_tiled_tma_atom(op, gmem_tensor, smem_layout, tile_shape) on the host side, and the copy is a cute.copy with a barrier pointer, wrapped in FA3’s code by copy_utils.tma_get_copy_fn. Descriptor prefetch (cpasync.prefetch_descriptor) hides the latency of the first descriptor fetch from global memory.

2.3.1 Multicast and clusters

Thread block clusters (up to 8, typically 2, thread blocks co-scheduled on neighboring SMs) permit distributed shared memory: a block can read the SMEM of its cluster peers, and TMA can multicast one GMEM read into the SMEM of several blocks in a cluster (cp.async.bulk.tensor... .multicast::cluster). The paper’s C++ kernels use cluster size 2 with multicast \(\mathbf{K}/\mathbf{V}\) loads for the forward pass when two blocks process adjacent \(\mathbf{Q}\) tiles of the same head. The CuTe DSL SM90 forward currently runs with cluster_shape_mn = (1, 1) and leaves multicast as a TODO (flash_fwd_sm90.py line 293: # No mcast for now) — one of the relevant C++/DSL differences (Section 5.12).

2.4 mbarriers: arrive/wait synchronization in shared memory

An mbarrier is a 64-bit object in shared memory supporting: init(count), arrive, arrive_and_expect_tx(bytes), wait(phase), plus the TMA-side complete_tx(bytes). It maintains pending and expected arrival counts, a transaction count, and a phase. A phase completes only when its pending-arrival count and transaction count have both reached zero; the object then reinitializes its pending-arrival count for the next phase. CuTe’s SM90 pipeline tracks the alternating phase as a parity bit, and waiters poll it with mbarrier.try_wait.parity. This is how a consumer waits until “buffer stage \(i\) is full.”

The argument to init(count) is the number of arrival units expected per phase—equivalently, the total decrement that the arrive-on operations must supply—not necessarily the number of threads in the CTA. Initialization itself is normally performed once by an elected thread, followed by a CTA rendezvous before any thread uses the object. Those are separate questions: who initializes the object does not determine its expected arrival count. The count must match the convention used by the subsequent arrivals. For example, if one representative lane per warp arrives, the count is the number of participating warps; if every thread arrives, it is the number of participating threads. In the CuTe DSL the expected count travels with a CooperativeGroup: the FA3 kernels bind Agent.Thread and pass the number of arrivals directly, using 1 for a lone TMA-issuing thread, 128 where every producer thread participates in a cp.async copy, and one arrival per MMA warp on the consumer side. cute.arch.elect_one() elects one lane within each executing warp; it does not by itself elect one representative for an entire warpgroup or CTA. A count that is too large prevents phase completion; a count that is too small can complete and reuse the barrier while intended participants still refer to the previous phase.

The canonical TMA pipeline uses two arrays of mbarriers per pipelined buffer: full[s] (producer signals: stage \(s\) contains fresh data) and empty[s] (consumers signal: stage \(s\) has been consumed and may be overwritten). This is exactly what the CuTe DSL’s PipelineTmaAsync encapsulates: producer_acquire waits on empty[s] then performs arrive_and_expect_tx on full[s]; the TMA copy targets full[s]’s transaction count; consumer_wait waits on full[s]; consumer_release arrives on empty[s]. Each side tracks its position with a (index, phase) pair — FA3’s pipeline.py packs both into one Int32 (PipelineStateSimple) so that index/phase extraction is a divmod (implemented with bitwise operations when the stage count is a power of two).

Named barriers are the other synchronization primitive FA3 uses: hardware barriers (barrier.sync/arrive %id, 16 IDs per block) that synchronize the threads that execute the matching barrier operations. The thread-count operand gives the number of participating thread arrivals; it does not encode which threads participate—control flow establishes that set. Both threads executing barrier.arrive and threads executing barrier.sync contribute to the count. Named barriers consume no SMEM and automatically reinitialize after completion, so software need not carry an explicit parity bit as it does for mbarriers. Code must nevertheless keep successive logical generations from overlapping. There are 16 IDs (0–15) per CTA; FA3 conventionally reserves ID 0 for the full-CTA sync_threads() barrier and allocates the rest via an enum (named_barrier.py): epilogue barriers, the pingpong scheduler barriers (WarpSchedulerWG1/2/3), and in the backward, PdS plus the dQFull/dQEmpty pairs that hand the dQ SMEM buffer between compute warpgroups and the store warp.

2.5 WGMMA: asynchronous warpgroup matrix multiply

wgmma.mma_async.sync.aligned.m64nNk16 computes, per instruction, \(C_{64\times N} \mathrel{+}= A_{64 \times 16} B_{16 \times N}\) for BF16/FP16 inputs (\(k{=}32\) for FP8), with \(N \in \{8, 16, \ldots, 256\}\). FA3 relies on the following properties:

1. Warpgroup-wide. All 128 threads of a warpgroup issue the instruction together; the accumulator \(C\) is distributed over the warpgroup’s registers (\(64 \times N\) FP32 elements = \(N/2\) registers per thread).

2. Operand sources: SS or RS. Operand \(B\) must be in shared memory. Operand \(A\) may come from shared memory (“SS” WGMMA) or from the warpgroup’s own registers (“RS”). FA3’s forward uses RS for the second GEMM (\(\mathbf{P}\mathbf{V}\)): the \(\mathbf{P}\) tile just produced by softmax lives in registers, and staging it through SMEM would cost bandwidth and sync (the DSL exposes this as mma_pv_is_rs). The backward uses SS for most GEMMs because the same \(\mathbf{P}/\mathbf{dS}\) tiles must be consumed transposed by other GEMMs, and a register-resident tile cannot be transposed across threads without going through memory.

3. SMEM operands are described by descriptors. A 64-bit matrix descriptor encodes the SMEM address, leading-dimension byte offset, stride-dimension byte offset, and swizzle mode. The operand’s K-major versus MN-major interpretation is selected by the WGMMA instruction’s transpose immediate together with those strides; it is not a separate “orientation bit” in the descriptor. Operands must be arranged in one of the canonical layouts (interleaved, or 32/64/128-byte swizzled “core matrix” layouts). CuTe hides this behind tile_to_shape(make_smem_layout_atom(...)), but the constraint surfaces in FA3 wherever a tensor must be read both normally and transposed (backward: \(\mathbf{Q}\) and \(\mathbf{Q}^\top\), K and \(\mathbf{K}^\top\), dS and \(\mathbf{dS}^\top\) — see Section 8.2).

4. Transposition limits. For FP16/BF16 SMEM operands, WGMMA provides transpose immediates: value 0 selects K-major, while value 1 selects M-major for A or N-major for B. Thus a compatible canonical SMEM layout can often be consumed through either logical orientation without running a separate transpose kernel; this is a descriptor/layout reinterpretation, not a data movement operation. Hopper’s FP8 WGMMA variants do not provide those transpose operands, so their SMEM operands use the K-major form. The FP8 forward therefore rearranges \(\mathbf{V}\) into the layout needed by the second GEMM inside the kernel (Chapter 6).

5. Asynchrony. wgmma.mma_async issues an asynchronous operation and returns control before the matrix multiply has completed. wgmma.commit_group batches preceding uncommitted operations, and wgmma.wait_group N blocks until at most \(N\) of the most recently committed groups remain pending and all older groups have completed. The accumulator and any register-resident A fragments of an operation may not be accessed until a wait covering that operation has completed. The DSL surfaces this as the wg_wait argument on its GEMM helpers: wg_wait=0 drains all prior groups before the helper returns; wg_wait=-1 issues and commits without adding a wait. The latter is used where FA3 overlaps pointwise softmax work with an outstanding GEMM.

6. Register-operand hazards. Before an RS-WGMMA reads registers written by regular instructions, the code must execute wgmma.fence (DSL: warpgroup.fence(), folded into the gemm helpers); before SMEM written by generic stores is read by WGMMA, a fence.proxy.async is needed (next section).

The scheduling constraint is that the accumulator and register-source fragments of an outstanding WGMMA must remain live and unmodified until an appropriate wait_group confirms completion. Deep software pipelining therefore lengthens register live ranges and can increase register pressure. This is one reason the backward pass distributes its five-GEMM chain across consumer warpgroups instead of relying only on a deeper per-warpgroup pipeline.

2.6 The async proxy and memory fences

Hopper distinguishes the generic proxy (ordinary ld/st) from the async proxy used by TMA and WGMMA shared-memory accesses. Conflicting accesses to the same shared-memory location through different proxies are not automatically ordered, even when initiated by the same thread. The rules FA3 must respect are:

• Generic write \(\to\) WGMMA/TMA read: requires fence.proxy.async — in the DSL, cute.arch.fence_view_async_shared(). This fence orders the proxy views; it is not a cross-thread rendezvous, so a warp or named barrier is also needed when different threads produce and consume the data. The forward uses this sequence whenever softmax results are staged into SMEM for an SS-GEMM (tPsP store, then fence, then sync_warp), in the backward after every r2s store of \(\mathbf{P}\)/\(\mathbf{dS}\)/\(\mathbf{dQ}\), and in every epilogue before the TMA store of a tile written to SMEM by register copies.

• TMA load to SMEM \(\to\) generic read: for an mbarrier-tracked TMA load, the corresponding consumer_wait observes completion and provides the required visibility to the generic proxy; no extra proxy fence is needed after that wait.

• TMA load to SMEM \(\to\) WGMMA read: the same consumer_wait prevents the WGMMA from issuing before the TMA load has completed. Both accesses use the async proxy, so this path does not require a generic-to-async proxy fence.

• WGMMA reads of SMEM buffers: a buffer being read by an in-flight WGMMA must not be overwritten; pipelines call consumer_release only after wait_group confirms completion. This is principally an asynchronous-operation lifetime rule, rather than a substitute for a proxy fence. In the DSL kernels, this ordering is explicit: e.g. warpgroup.wait_group(1) … pipeline_k.consumer_release(...) in mma_one_n_block_intrawg_overlap ensures that the older of two pending GEMM groups has completed before releasing K’s stage.

2.7 Register reallocation: setmaxnreg

The kernel’s compile-time and launch constraints establish a maximum per-thread register allocation. On Hopper, setmaxnreg can redistribute part of that allocation through a register pool maintained per CTA: setmaxnreg.dec.sync.aligned.u32 N lowers the executing warps’ absolute per-thread maximum to \(N\) and returns the excess to the pool; setmaxnreg.inc raises the maximum to \(N\) and blocks until the pool can satisfy the request. \(N\) is a multiple of 8 in the inclusive range 24–256. Every warp in the warpgroup must execute the same instruction, and an explicit warpgroup synchronization is required before a subsequent setmaxnreg.

FA3’s warp specialization is only efficient because of this. A TMA-only producer warpgroup needs almost no registers (∼24–56), while consumer warpgroups want the maximum. In the CuTe DSL forward (flash_fwd_sm90.py):

(the dict {1: (256, 56), 2: (240, 24), 3: (160, 32)} at line 215). The backward uses \((240,240,24)\) for its two consumer warpgroups plus producer, or \((256,224,24)\) in the dQ_single_wg configuration where warpgroup 0 does the extra dQ GEMM and needs more registers. The budget constraint is (regs per thread) \(\times\) (128 lanes) \(\times\) (warpgroups) \(\le 64\)K per SM with one resident block, i.e. the per-thread allocations across the 3–4 warpgroups must sum to \(\le 512\) (the code asserts REG_LIMIT = 504 for 2 MMA warpgroups because allocation granularity wastes 8).

2.8 Threadblock clusters, persistent kernels, and occupancy

FA3 launches with min_blocks_per_mp=1 and 384–512 threads per block (3–4 warpgroups), using nearly all of SMEM: the forward hdim-128 pipeline holds 2 stages each of K and V at \(128\times128\times2\) bytes per tile — four 32 KB tiles = 128 KB — plus Q and auxiliary buffers. Occupancy in the classical sense (warps per SM) is deliberately low — 1 block of 384 threads = 12 warps — because latency hiding comes from asynchrony within the block, not from warp oversubscription. The paper’s forward kernel is persistent-capable (grid = number of SMs, tiles pulled from a scheduler); the CuTe DSL SM90 forward currently launches one block per tile (is_persistent=False in flash_fwd_sm90.py), relying on the LPT reordering plus hardware scheduling to get the same effect (Chapter 9).

2.9 The SFU: why softmax is expensive

Each SM sub-partition has an SFU issuing four MUFU operations per cycle (16/SM total). MUFU.EX2 computes \(2^x\) to within a couple of ULP. FA3 exclusively uses base-2 exponentials: instead of \(e^{s - m}\) it computes \(2^{s\cdot\log_2 e - m\cdot\log_2 e}\), folding \(\log_2 e\) into the softmax scale (softmax_scale_log2), saving one multiply per element and using MUFU.EX2 directly. The row-max subtraction is fused as \(2^{s \cdot c - m c} = \texttt{exp2(fma(s, c, -mc))}\). Because FP32 ALU (FFMA) throughput is \(4\times\) MUFU throughput, the FMA component has comparatively low cost; the EX2s themselves are the bottleneck that pingpong and intra-warpgroup overlap hide. On the accuracy side, the fast approximate exp2 introduces a small numerical approximation. Using the same approximated exponentials in the numerator and row sum improves consistency but does not make the error cancel exactly. Its suitability is therefore an empirical accuracy choice for the supported low-precision kernels, not a proof of zero numerical effect.

2.10 Putting it together: the FA3 hardware checklist

Every FA3 design decision in the next chapters maps to one of these hardware facts:

3 Background: Attention, Online Softmax, FA1/FA2

FlashAttention-3 retains the exact tiled-attention algorithm established by FlashAttention-1 (FA1) and the thread-block and warp partitioning introduced by FlashAttention-2 (FA2). Its principal changes concern execution on Hopper: asynchronous data movement, warp specialization, overlap between matrix multiplication and softmax, and an FP8 path. This chapter separates the algorithmic inheritance from those Hopper-specific changes.

3.1 Attention and the quadratic intermediate

For one batch element and one attention head, let \(\mathbf{Q}\in \mathbb{R}^{N_q \times d}\), \(\mathbf{K}\in \mathbb{R}^{N_k \times d}\), and \(\mathbf{V}\in \mathbb{R}^{N_k \times d_v}\). Scaled dot-product attention is

where softmax is applied row-wise and masks are represented by replacing excluded scores with \(-\infty\) before softmax. The computation requires \(O(N_qN_kd + N_qN_kd_v)\) arithmetic. FlashAttention does not reduce this arithmetic complexity and does not approximate the attention matrix.

A conventional implementation treats Equation (3.1) as three separately scheduled operations: a GEMM writes \(\mathbf{S}\) to HBM, a softmax kernel reads \(\mathbf{S}\) and writes \(\mathbf{P}\), and a second GEMM reads \(\mathbf{P}\). The \(N_qN_k\) intermediates dominate memory traffic and saved-activation storage when \(N_q,N_k \gg d,d_v\). For the square case \(N=8192\), an FP32 score matrix occupies \(8192^2 \times 4\) bytes \(=256\) MiB per head. It cannot reside in the 50 MB L2 cache of an H100, and the separate kernels require multiple complete reads or writes of that matrix.

This distinction is central:

• Quadratic arithmetic remains. Dense exact attention still evaluates the relevant score entries.

• Quadratic HBM-resident intermediates are removed. Score and probability tiles exist only in registers or shared memory; global storage contains the linear-size inputs, output, and row statistics.

3.2 Safe softmax and its streaming state

Numerically safe softmax subtracts the maximum of each row before exponentiation. For a score row \(s \in \mathbb{R}^{N_k}\),

The online formulation [4] partitions the row into blocks \(s^{(1)},s^{(2)},\ldots\) and maintains two sufficient statistics: the largest score processed so far, \(m^{(j)}\), and the exponential sum expressed relative to that maximum, \(\ell^{(j)}\). After block \(j\),

The first term changes the reference point of the previous sum whenever a larger maximum appears. The pair \((m,\ell)\) is therefore a lossless summary of all blocks processed so far, apart from ordinary floating-point rounding.

Attention adds an unnormalized output accumulator \(\widetilde{\mathbf{O}}\). For a KV block \(\mathbf{V}^{(j)}\) corresponding to scores \(s^{(j)}\),

After the final block,

Thus every score tile can be discarded after its contribution has entered \((m,\ell,\widetilde{\mathbf{O}})\).

3.2.1 The block-combine interpretation

The same recurrence can combine partial results produced independently. For disjoint key sets \(A\) and \(B\), define their states \((m_A,\ell_A,\widetilde{\mathbf{O}}_A)\) and \((m_B,\ell_B,\widetilde{\mathbf{O}}_B)\). Their union has

This associative state-combination rule explains both streaming within a thread block and split-KV reduction across thread blocks (Chapter 10). Different groupings can produce different last bits because floating-point addition is not associative, but the real-number result is unchanged.

3.2.2 Consequences for implementation

Three consequences recur throughout the Hopper kernels:

• A change in row maximum requires an \(O(d_v)\) rescale of the existing output accumulator. FA3 overlaps this pointwise work with asynchronous WGMMA and, for some configurations, performs the rescale before the \(\widetilde{\mathbf{P}}\mathbf{V}\) GEMM so that the GEMM accumulates into the already rescaled buffer.

• The forward pass writes \(\operatorname{LSE}\) rather than the \(N_q\times N_k\) probability matrix. The backward pass reconstructs each required probability tile from \(\mathbf{Q}\), \(\mathbf{K}\), and \(\operatorname{LSE}\).

• Implementations commonly evaluate the exponential as \(2^{(s-m)c}\), where \(c=\alpha\log_2 e\). The returned \(\operatorname{LSE}\) is converted to natural-log units, so the external mathematical definition is unchanged.

3.3 FlashAttention-1: IO-aware exact attention

FA1 [2] established the principal algorithmic idea: fuse the two GEMMs and softmax into a tiled kernel, retain only a score tile and online-softmax state on chip, and avoid materializing \(\mathbf{S}\) or \(\mathbf{P}\) in HBM. In the original algorithmic presentation, blocks of \(\mathbf{K}\) and \(\mathbf{V}\) form the outer loop. For each KV block, the algorithm visits the query blocks, loads their current output and softmax state, incorporates the KV block, and writes the updated state back. This ordering maximizes reuse of the resident KV block.

Let \(M\) denote usable on-chip SRAM capacity measured in scalar elements. For \(d\le M\le Nd\), the FA1 analysis gives

compared with \(\Theta(N^2+Nd)\) for conventional attention. The reduction is large in the practical regime \(d^2\ll M\). Equation (3.5) is an IO bound, not a statement that each tensor is read only once: the loop order revisits query, output, and row-statistic blocks as successive KV blocks are processed. The FA1 paper also proves an IO lower bound showing that no exact attention algorithm is asymptotically better for every SRAM size in the analyzed range.

3.3.1 Backward recomputation

A conventional training implementation saves \(\mathbf{P}\) for backward. FA1 saves the linear-size output and row normalization statistics instead. During backward, a tile is reconstructed as

This adds recomputation of the score GEMM, but removes the read and write of a quadratic activation. The trade is favorable because Tensor Core arithmetic is much faster than HBM traffic and because the recomputed tile is consumed immediately by the gradient GEMMs.

3.3.2 First-generation work partitioning

The first CUDA implementation assigned one thread block to one (batch, head) pair. That block traversed all query and key blocks belonging to the head. Parallelism was therefore limited to \(B\times H\) thread blocks; long-sequence workloads with small batch or head counts could leave many SMs idle. Within a thread block, FA1 divided K and V across warps while Q was shared. Each warp produced only a partial contribution to the output, so the partials had to pass through shared memory and be reduced across warps. The FA2 paper calls this the split-K warp partition.

3.4 FlashAttention-2: parallelism and work partitioning

FA2 [3] preserves FA1’s exact tiled computation and linear auxiliary memory while addressing three distinct costs.

3.4.1 Fewer non-matrix operations

Tensor Core matrix multiplication has much higher peak throughput than general FP32 arithmetic. A small count of exponentials, reductions, scales, and divisions can therefore occupy a large fraction of kernel time. FA2 uses the unnormalized accumulator in Equation (3.3): only the old accumulator is multiplied by \(e^{m^{(j-1)}-m^{(j)}}\), while the new \(\widetilde{\mathbf{P}}^{(j)}\mathbf{V}^{(j)}\) contribution is added directly. Division by \(\ell\) occurs once, after the final KV block. FA2 also stores the single statistic \(\operatorname{LSE}=m+\log\ell\) for backward instead of storing \(m\) and \(\ell\) separately.

3.4.2 A query block becomes an independent thread-block task

FA2 reverses the forward loop order used in the original FA1 presentation: the outer loop ranges over query blocks, and each query block scans all relevant KV blocks. A thread block owns one \(\mathbf{Q}_i\) and its complete online state until \(\mathbf{O}_i\) is finalized. The tasks

are independent in the forward pass. The grid therefore grows by the number of query blocks and can occupy the GPU even when \(B\times H\) is small. This ordering also writes each output tile only once, although it reloads the KV sequence for each query tile; L2 locality and sufficiently large tiles make that exchange favorable on modern GPUs.

The backward pass uses the complementary ownership: a thread block owns a key/value column block so that \(\mathbf{dK}_j\) and \(\mathbf{dV}_j\) can remain local while it visits query blocks. Contributions to \(\mathbf{dQ}_i\) come from multiple column blocks and require cross-block accumulation. This ownership pattern remains visible in FA3’s FP32 dQaccum and its deterministic semaphore path.

3.4.3 Split-Q warp partitioning

FA2 also changes ownership within a thread block. Instead of partitioning K and V across warps, it partitions rows of Q across warps while each warp can access the required K and V tiles. A warp computes the corresponding rows of both \(\mathbf{Q}\mathbf{K}^\top\) and \(\mathbf{P}\mathbf{V}\) and owns the same rows of \(\mathbf{O}\). No cross-warp reduction of output partials is required. This split-Q partition removes shared-memory writes, synchronization, and reads that were present in FA1’s split-K design.

The term split-Q here describes warp ownership within one attention tile. It is unrelated to split-KV decoding, where independent thread blocks process disjoint ranges of the KV sequence and a separate kernel combines their softmax states (Chapter 10).

3.4.4 Performance envelope and the remaining Hopper gap

On A100, the FA2 paper reports approximately \(2\times\) the performance of FA1 and 50–73% of theoretical peak, depending on direction and shape. The same implementation reaches only about 35% of H100 peak in the FA3 study. The algorithm is not the cause of that regression in efficiency. The FA2 kernel does not exploit Hopper’s TMA producer, warpgroup-level WGMMA, dynamic register reallocation, or asynchronous overlap; it also schedules the GEMMs and softmax largely in sequence. FA3 retains FA2’s query-tile ownership and split-Q structure, then changes how each tile is executed.

3.5 FA1, FA2, and FA3 in one comparison

Table 3.1: The main transition across the three FlashAttention generations.

3.6 Kernel anatomy inherited by FA3

The forward chapters use the following decomposition:

1. Tile mapping: map a logical \((\text{query block},\text{head},\text{batch})\) task to a CUDA block. Chapter 9 adds causal cost ordering, L2 swizzling, and ragged-batch mapping.

2. KV traversal: keep a query tile and its online-softmax state local while visiting the relevant K/V tiles. Causal and local masks restrict this interval.

3. Score and state update: compute \(\mathbf{Q}_i\mathbf{K}_j^\top\), apply masks, update \((m,\ell)\), convert the probability tile to a GEMM operand, and accumulate \(\widetilde{\mathbf{P}}_{ij}\mathbf{V}_j\).

4. Epilogue: divide by \(\ell\), convert to the output dtype, store \(\mathbf{O}_i\), and store \(\operatorname{LSE}_i\) when requested or needed by backward.

Chapter 4 develops the dependencies in this loop; Chapter 5 maps them to the CuTe DSL implementation.

4 The FA3 Forward Pass: Algorithm

This chapter presents the forward-pass design as in the FA3 paper [1], mapped onto the vocabulary of Chapter 2. Chapter 5 then traces the same operations through the CuTe DSL implementation.

4.1 Work decomposition

One CUDA thread block processes one tile of \(B_M\) query rows for one (head, batch): it owns \(\mathbf{Q}_i \in \mathbb{R}^{B_M \times d}\) and produces \(\mathbf{O}_i \in \mathbb{R}^{B_M \times d_v}\) and \(\operatorname{LSE}_i \in \mathbb{R}^{B_M}\). Inside the block, it iterates over \(\mathbf{K}_j \in \mathbb{R}^{B_N \times d}\) and \(\mathbf{V}_j \in \mathbb{R}^{B_N \times d_v}\) tiles, \(j = j_{\max}-1 \ldots j_{\min}\) (descending; for causal attention \(j_{\max}\) is the diagonal). For the standard attention path, one iteration performs the following work:

Here \(\mathcal F\) is the optional score transformation (the identity in standard attention) and \(\mathcal M\) assigns \(-\infty\) to excluded entries. The implementation leaves \(\mathbf{R}_j\) unscaled in the WGMMA accumulator and incorporates \(\alpha\log_2 e\) when evaluating exp2; the equations use \(\mathbf{S}_j\) for the mathematically scaled and masked scores. The block-level tiling follows FA2, while FA3 changes the assignment and temporal ordering of loads, GEMMs, and softmax work.

4.2 Online-softmax state and invariant

The tiled algorithm never materializes the complete score or probability matrix. For each query row, it retains three FP32 quantities across KV tiles: the running maximum \(m_t\), the shifted exponential sum \(\ell_t\), and the unnormalized output vector \(\tilde{\mathbf{O}}_t\). Let \(\mathbf{S}^{(t)}\) denote the scores in the \(t\)th KV tile in processing order. With \(m_{-1}=-\infty\), \(\ell_{-1}=0\), and \(\tilde{\mathbf{O}}_{-1}=0\), the update is

After tile \(t\), these variables satisfy

where \(\mathcal K_t\) is the set of keys processed through tile \(t\). The normalized result is therefore \(\mathbf{O}=\tilde{\mathbf{O}}_T/\ell_T\), independent of the tile order, and \(\operatorname{LSE}=m_T+\log\ell_T\). Masked entries contribute zero because their score is \(-\infty\).

The kernel evaluates the same recurrence in base 2. If the WGMMA accumulator contains the unscaled dot product \(R\), it computes \(2^{(R-m)\alpha\log_2 e}\) and \(2^{(m_{\mathrm{prev}}-m)\alpha\log_2 e}\). This formulation avoids a separate matrix-wide multiplication by \(\alpha\). The factor \(\rho_t\) must rescale both \(\ell_{t-1}\) and \(\tilde{\mathbf{O}}_{t-1}\) before the new tile is accumulated; delaying or omitting either rescale changes the softmax distribution.

4.3 Producer–consumer warp specialization

The thread block launches \(128 \times (1 + W)\) threads: one producer warpgroup and \(W\) consumer (MMA) warpgroups (\(W = 2\) for the standard tile sizes; \(W=1\) or 3 for some head dimensions). Roles:

• Producer warpgroup (after setmaxnreg.dec to 24–56 registers): a single warp of it issues TMA loads. It loads \(\mathbf{Q}_i\) once, then loops over \(j\), loading \(\mathbf{K}_j\) and \(\mathbf{V}_j\) into stage \(j \bmod s\) of a circular SMEM buffer of \(s\) stages (\(s =\) num_stages, typically 2 for \(d = 128\)). It producer_acquires stage \(j \bmod s\) (waits until consumers released it), issues the TMA, and the TMA’s transaction-count completion signals the stage full.

• Consumer warpgroups (after setmaxnreg.inc): execute the loop body (4.1). The two consumer warpgroups split the query rows of the tile: with \(B_M = 128\), warpgroup 1 owns rows 0–63 and warpgroup 2 rows 64–127 (this is the atom_layout_mnk = (tile_m // 64, 1, 1) in the tiled MMA); both consume the same \(\mathbf{K}_j/\mathbf{V}_j\) stages, so the pipeline’s consumer group is “all MMA warps” and a stage is recycled when both warpgroups have released it.

The paper’s Algorithm 1 (CUDA-pseudocode for the two roles) is reproduced here in compressed form:

Algorithm 1 (FA3 forward, warp-specialized).
Producer warpgroup: deallocate registers; load \(\mathbf{Q}_i\) (TMA); for \(j = 0 \ldots\): wait empty[j mod s]; TMA-load \(\mathbf{K}_j, \mathbf{V}_j\) into stage \(j \bmod s\); (completion arrives full[j mod s]).
Consumer warpgroups: allocate registers; wait for \(\mathbf{Q}_i\); for \(j\): wait full[j mod s]; \(\mathbf{S}\leftarrow \mathbf{Q}_i\mathbf{K}_j^\top\) (SS-WGMMA); online softmax; \(\tilde{\mathbf{O}}\leftarrow \text{rescale}(\tilde{\mathbf{O}}) + \tilde{\mathbf{P}}\mathbf{V}_j\) (RS-WGMMA); arrive empty[j mod s]. Finally: normalize by \(\ell\), write \(\operatorname{LSE}\), write \(\mathbf{O}\) via SMEM + TMA store.

K and V receive separate mbarrier pairs (in the DSL: pipeline_k, pipeline_v) even though they are loaded together: GEMM0 only needs \(\mathbf{K}_j\), GEMM1 only \(\mathbf{V}_j\), so distinguishing their readiness lets GEMM0 start before \(\mathbf{V}_j\) arrives and lets \(\mathbf{K}_j\)’s stage be released one GEMM earlier. It also permits the V load of iteration \(j\) to be deferred until after the K load of iteration \(j+1\) in the overlapped schedule below.

4.3.1 Pipeline-stage ownership

Each K or V stage is governed by a full/empty mbarrier pair. The barrier protocol establishes both readiness and exclusive reuse:

1. The producer waits for the stage’s empty barrier before overwriting its SMEM storage.

2. The producer arms the full barrier with the expected byte count and issues the TMA transfer. TMA contributes the transaction bytes on physical completion; instruction issue alone does not make the stage readable.

3. Every consumer warp waits for the full phase before issuing a WGMMA that references the stage.

4. The consumers release the empty phase only after a WGMMA wait proves that the hardware has finished reading that operand. The producer may then reuse the stage in the next pipeline cycle.

The producer and consumers maintain independent stage cursors, but advance them in the same logical tile order. A cursor contains a circular-buffer index and a phase bit; wrapping from stage \(s-1\) to stage 0 toggles the phase, so a wait cannot be satisfied by completion from an earlier reuse of that slot.

The principal lifetimes within one work tile are summarized below.

This separation explains why a single “KV ready” barrier would reduce overlap: K could not be consumed until V completed, and K storage could not be recycled independently of V storage.

4.4 Pingpong scheduling: hiding softmax between warpgroups

With two consumer warpgroups resident on one SM, the hardware scheduler would by default interleave them arbitrarily. FA3 forces a useful pattern with two named barriers (IDs WarpSchedulerWG1, WarpSchedulerWG2): before issuing its GEMMs for iteration \(j\), warpgroup \(w\) waits on its own barrier; right after issuing (before waiting on the WGMMA), it arrives on the other warpgroup’s barrier. Warpgroup 1 arrives once at startup (mma_init) so the chain starts. The effect: the two warpgroups’ Tensor Core phases are serialized and alternate, so whenever warpgroup A is in its softmax phase, the Tensor Cores are busy with warpgroup B’s GEMMs:

        ┌─────────────┬──────────────────────┬──────────────┬─────────────┬──────────────────────┐
WG1     │ GEMM0_j     │ softmax_j            │ GEMM1_j      │ GEMM0_j+1   │ softmax_j+1          │
        └─────────────┴──────────────────────┴──────────────┴─────────────┴──────────────────────┘
        ┌─────────────┬──────────────────────┬───────────────────────┬──────────────────────────┐
WG2     │ softmax_j-1 │ GEMM1_j-1/GEMM0_j    │ softmax_j             │ GEMM1_j/GEMM0_j+1        │
        └─────────────┴──────────────────────┴───────────────────────┴──────────────────────────┘
        ────────────────────────────────────────────────────────────────────────────────────────▶ time

        GEMM* blocks: Tensor Core (WGMMA)        softmax blocks: MUFU/ALU (softmax, exp)

The named barrier only gates the start of GEMM issue; WGMMA asynchrony provides the remaining overlap. Pingpong is enabled in the DSL when use_scheduler_barrier is true — with intra-warpgroup overlap on, that is for \(W \ge 2\) and \(d \le 128\) (for larger head dimension the GEMMs are long enough that softmax hides without help, and the barrier only adds latency).

4.5 Intra-warpgroup overlapping: the 2-stage GEMM–softmax pipeline

Pingpong hides softmax across warpgroups. FA3’s second schedule hides it within one warpgroup by breaking the sequential dependency GEMM0\(_j \to\) softmax\(_j \to\) GEMM1\(_j\): while softmax of iteration \(j\) runs, the Tensor Cores work on GEMM0 of iteration \(j{+}1\) (and GEMM1 of iteration \(j\) is issued only after softmax \(j\) completes, using the \(\tilde{\mathbf{P}}_j\) still in registers). The steady-state loop body (paper Algorithm 2) is:

2-stage pipelined consumer (steady state, iteration \(j\)):

1. wait full_K[j+1]; issue GEMM0\(_{j+1}\) (\(\mathbf{S}_{\text{next}} = \mathbf{Q}\mathbf{K}_{j+1}^\top\)) — async, do not wait;

2. if early O-rescaling is enabled, apply the saved rescale factor; wait full_V[j]; issue GEMM1\(_{j}\) (\(\tilde{\mathbf{O}}\mathrel{+}= \tilde{\mathbf{P}}_j \mathbf{V}_j\)) — async, do not wait;

3. wgmma.wait_group(1) \(\Rightarrow\) GEMM0\(_{j+1}\) done; release \(\mathbf{K}_{j+1}\) stage;

4. softmax on \(\mathbf{S}_{\text{next}}\): row max, \(2^{x}\), and the next rescale factor;

5. wgmma.wait_group(0) \(\Rightarrow\) GEMM1\(_{j}\) done; release \(\mathbf{V}_j\) stage; in the default path rescale \(\tilde{\mathbf{O}}\), otherwise save the factor for the next iteration; convert \(\tilde{\mathbf{P}}_{j+1}\) to FP16/BF16 registers.

The softmax step (MUFU-bound) executes while GEMM1\(_j\) is still in flight. The cost is register pressure: two S-tiles’ worth of state coexist (\(\mathbf{S}_{j+1}\) accumulator + \(\tilde{\mathbf{P}}_j\) operand), plus the O accumulator. The additional live state requires smaller block sizes and can introduce register spilling; in the DSL this constraint appears as the smaller tile_n for large head dimensions (Table 4.1).

Two scheduling constraints are visible in mma_one_n_block_intrawg_overlap, Listing 5.9):

• O-rescale placement. In the default overlapped schedule, GEMM1\(_j\) first completes, softmax\(_{j+1}\) supplies its factor, and the output accumulated through tile \(j\) is then rescaled before GEMM1\(_{j+1}\). For \(d_v > 128\) the DSL instead enables rescale_O_before_gemm: softmax\(_j\) saves its factor in scores_scale, and the next half iteration applies it while GEMM0\(_{j+1}\) is in flight but before GEMM1\(_j\) begins. This moves the long O-vector rescale away from the post-GEMM wait path without permitting simultaneous WGMMA and ALU writes to \(\tilde{\mathbf{O}}\).

• Loop boundaries. Pairing GEMM1\(_j\) with GEMM0\(_{j+1}\) requires a QK-only prologue and a PV-only drain. The DSL implements them in first_half_block_overlap and last_half_block_overlap; the next subsection states the complete sequence.

4.5.1 Prologue, steady state, and drain

Let \(j_0,j_1,\ldots,j_{T}\) denote KV tiles in processing order, regardless of their numerical block indices. The half-iteration offset can then be written without ambiguity:

The first output GEMM uses WGMMA’s zero-initializing form; later output GEMMs accumulate into the persistent FP32 \(\tilde{\mathbf{O}}\) fragment. Empty work tiles and fully masked rows are handled separately so that an uninitialized output fragment cannot reach the epilogue. At no point is a \(B_M\times B_N\) score or probability tile written to global memory: \(\mathbf{R}_j\) is overwritten by \(\tilde{\mathbf{P}}_j\) in registers, and the converted P fragment remains live only until its paired GEMM1 is issued.

Combined with pingpong, Tensor Core issue alternates between warpgroups, while each warpgroup overlaps its softmax work with an independent in-flight GEMM. Measured at \(d{=}128\), \(N \approx 8\)K FP16, pingpong raises the plain warp-specialized kernel from ∼570 to ∼620–640 TFLOP/s, and the intra-warpgroup pipeline to ∼640–661 [9, 1] (Table 12.3). The paper also reports lower performance from a deeper three-stage variant because the added register pressure and compiler scheduling costs exceeded the additional overlap benefit.

4.6 Softmax details: rescaling, correction, and LSE

Per row \(r\) of the S-tile the consumer maintains row_max[r] and row_sum[r] in registers. Per iteration (Listing 5.12 shows the code):

1. Row max. A thread owns a fragment of the row; it reduces its fragment (fmax) and then performs a butterfly reduction over the quad of 4 threads that jointly own the row (shfl.bfly with offsets 1 and 2) — with the SM90 accumulator layout, 4 threads share each row.

2. Guard against \(-\infty\). If an entire row is masked (fully out-of-window rows in local attention or short sequences), the max is \(-\infty\) and \(2^{s - m}\) would be NaN (\(\infty - \infty\)); the code replaces \(m = -\infty\) with 0 for the subtraction (check_inf).

3. Exponentials. \(\tilde p = 2^{(s - m)c}\) via MUFU.EX2, where \(c\) = softmax_scale_log2 \(= \alpha \log_2 e\).

4. Rescale factor. \(\text{rescale} = 2^{(m_{\text{prev}} - m_{\text{new}})c}\). The row sum is folded as \(\ell \leftarrow \ell \cdot \text{rescale} + \operatorname{rowsum}(\tilde p)\). The quad reduction of the row sum is deferred to the finalize step — each of the 4 threads keeps a partial sum; only the final normalization needs the true row sum. This shaves 2 shuffles per iteration off the critical path.

5. Finalize (after the loop): quad-reduce row_sum; compute the reciprocal (rcp.approx.f32); handle empty rows (\(\ell = 0\) or NaN \(\Rightarrow\) output 0, \(\operatorname{LSE}= -\infty\)); produce \(\operatorname{LSE}= (m c + \log_2 \ell)\ln 2\) in natural-log units; multiply \(\tilde{\mathbf{O}}\) rows by \(1/\ell\).

A note on learnable sinks: a per-head “attention sink” value may be folded into \(\ell\) at finalize time (sink_val in the code), implementing learnable-sink attention (as in GPT-OSS-style models) at zero mainloop cost; under PackGQA (Chapter 11) each packed row belongs to a different query head, so the sink is fetched per row.

4.7 Causal masking and the split loop

For causal masks, tile \((i, j)\) falls in one of three regimes: fully visible (\(j B_N + B_N - 1 \le\) diagonal), partially masked (straddles the diagonal), or fully masked (skipped via loop bounds). FA3 structures the \(j\)-loop to keep masking off the steady-state path:

1. Seqlen-masked first iteration: the very last K/V tile (\(j = j_{\max}-1\)) may overhang the end of the sequence; it is processed first, with bounds masking (recall: TMA zero-fills the loads, so masking of scores, setting them to \(-\infty\), is what matters).

2. Causal-masked iterations: for causal/local attention, the next few tiles straddle the diagonal. The count is computed by get_n_block_min_causal_local_mask(): with equal \(B_M = B_N\) one tile straddles; with \(B_M = 192, B_N = 128\) up to \(\lceil 192/128 \rceil + 1\) tiles do.

3. Unmasked steady state: the remaining tiles run the pure pipeline with no mask evaluation at all.

The mask itself is applied on the FP32 accumulator between GEMM0 and softmax. Per element, the causal predicate is \(\text{col} > \text{row} + (\text{seqlen}_k - \text{seqlen}_q)\) (bottom-right alignment). The DSL uses an SM90-specific optimization that converts the column threshold into a per-32-column bitmask and applies it with the R2P (register to predicate) instruction pattern instead of per-element integer comparisons (mask_r2p_lambda in mask.py) — saving ALU work in the two masked regimes.

Any score transformation must precede the final mask. The implementation therefore applies the optional score-modification function to valid accumulator coordinates, then overwrites excluded coordinates with \(-\infty\), and only then computes the row maximum. Reversing the first two operations would allow a general score function to turn a masked value into a finite value. Sequence bounds and causal/local predicates are logically separate: TMA zero filling prevents invalid memory reads, whereas assigning \(-\infty\) prevents padded keys from contributing to softmax.

Local (sliding-window) attention adds a left boundary; the loop then has four segments (seqlen-masked, right-mask, unmasked, left-mask), all precomputed by block_info.py from the window sizes.

4.8 Per-tile computation and data movement

For one dense KV tile, the two GEMMs perform

floating-point operations when a fused multiply–add counts as two. If each input element occupies \(b\) bytes, the streaming K/V traffic for that tile is approximately

excluding descriptor traffic and cache effects. The GEMM arithmetic intensity relative to the streamed K/V operands is therefore \(2B_M/b\) FLOP/byte: \(128\) FLOP/byte for \(B_M=128\) with BF16 or FP16. Q contributes \(bB_Md\) bytes once per query tile, while O and LSE contribute approximately \(bB_Md_v+4B_M\) output bytes once after all KV tiles.

This accounting clarifies the tile-shape tradeoffs. Increasing \(B_M\) reuses each K/V tile across more query rows and raises arithmetic intensity, but it also enlarges the persistent O accumulator and increases the number of consumer warpgroups. Increasing \(B_N\) does not change the idealized K/V arithmetic intensity, but it reduces loop and barrier overhead and supplies a larger WGMMA tile; it simultaneously increases S/P fragment size, softmax work per iteration, and the SMEM required by every K/V pipeline stage. The selected configuration must therefore satisfy SMEM and register limits while leaving enough independent work to cover TMA, WGMMA, and MUFU latency.

4.9 Tile sizes, head-dimension variants, and register budgets

The tile configuration must satisfy the SMEM capacity (\(\mathbf{Q}+ s(\mathbf{K}+ \mathbf{V})\) tiles), register-file requirement (O-accumulator \(B_M \times d_v \times 4\) bytes spread over consumer threads, plus S/P fragments), and load-balance requirement. The DSL uses the following tuned values (_tile_size_fwd_sm90 in interface.py, benchmarked on H100 SXM); most match the paper’s C++ kernels:

Table 4.1: Forward tile configurations on SM90 (CuTe DSL, BF16/FP16). “RS” = \(\mathbf{P}\mathbf{V}\) GEMM reads \(\mathbf{P}\) from registers; “overlap” = 2-stage intra-warpgroup pipeline. All use 1 producer warpgroup plus \(B_M/64\) consumer warpgroups (2 for \(B_M{=}128\), 3 for \(B_M{=}192\)).

(The FA3 paper states no concrete tile sizes — only the \(B_r, B_c\) symbols and SMEM-limit reasoning; the values in Table 4.1 are from the shipping heuristics, C++ hopper/tile_size.h and DSL interface.py, which have been re-tuned over time.)

The table reflects three constraints:

• Small \(d\): the O-accumulator is small, so \(B_M = 192\) (three 64-row MMA atoms, i.e. three consumer warpgroups of 64 rows each) raises arithmetic intensity per \(\mathbf{K}/\mathbf{V}\) byte.

• \(d = 96\), \(B_M = 192\): keeping \(\mathbf{P}\) in registers (mma_pv_is_rs) costs too many registers at this \(B_M\); the DSL comment reports that RS with \(192\)-row tiles reaches approximately \(300\) TFLOP/s instead of \(600\) TFLOP/s, so \(\mathbf{P}\) is staged through SMEM (sP) instead.

• Large \(d\): SMEM for K/V stages scales with \(d \cdot B_N\); \(B_N\) shrinks to keep \(2\) pipeline stages resident within 228 KB, and rescale_O_before_gemm becomes active (\(d_v > 128\)) because the O accumulator (\(128 \times 256\) FP32 = 128 registers/thread just for O at \(d_v{=}256\)) dominates the register file.

Head dimensions supported on SM90: \(8 \le d, d_v \le 256\), any multiple of 8 (non-multiples of 16 are padded to the tile and predicated) — broader than the paper’s C++ kernels, which instantiate \(d \in \{64, 96, 128, 192, 256\}\) (plus hdim-diff combos like \(192/128\) for DeepSeek-style models).

4.10 Epilogue

After the loop: finalize softmax; each consumer thread’s O fragment is converted to the output dtype and staged into SMEM (reusing the \(\mathbf{Q}\) tile’s SMEM — \(\mathbf{Q}\) is dead by then). A named barrier first prevents any writer from reusing the aliased storage until all participating consumers have finished their previous reads. After the register-to-SMEM copies, a fence.proxy.async orders the generic stores before a TMA read, and a second named-barrier generation rendezvouses the writing threads with the store warp. Then one warp (warp 4, the first consumer warp) issues the TMA store of the tile and waits on cp.async.bulk.wait_group. \(\operatorname{LSE}\) is written directly from registers by the threads owning column 0, with bounds predication — unless PackGQA is active, in which case a scatter helper handles the interleaved head layout. The store is ragged-aware for variable-length batches (Chapter 9).

5 The FA3 Forward Pass: CuTe DSL Implementation

This chapter examines the class FlashAttentionForwardSm90 in flash_attn/cute/flash_fwd_sm90.py (∼1550 lines), top to bottom. All listings are excerpts of the real file, trimmed where marked; line numbers refer to the July 2026 tree. The base class FlashAttentionForwardBase (flash_fwd.py) supplies the SMEM layout machinery and the epilogue; softmax.py, pipeline.py, mask.py, block_info.py supply the components.

5.1 A note on the CuTe DSL

The CuTe DSL [21] embeds CUTLASS’s layout algebra in Python: decorators @cute.jit (inlineable device function) and @cute.kernel (kernel entry) trace Python into MLIR and JIT-compile it. cutlass.const_expr(...) marks compile-time branches — everything inside if const_expr(...) is specialized away, serving the same role as if constexpr in the C++ kernels. cute.Tensor carries an engine (pointer) plus a Layout (shape/stride algebra); cute.local_tile, partition_S/D, TiledCopy, TiledMma are direct ports of their CuTe C++ namesakes. Python-level loops with cutlass.range(..., unroll_full=True) become fully unrolled device loops. The resulting source preserves the C++ kernel’s layout and pipeline structure while expressing most specialization through Python and const_expr branches.

5.2 Implementation map

The implementation has two levels. __call__ runs while constructing the specialized kernel: it validates layouts, selects tile and copy atoms, computes launch dimensions, and passes compile-time objects to kernel. The decorated @cute.kernel body defines device control flow. Its producer and consumer branches then call smaller @cute.jit device functions.

The sections below follow this call order. Algorithmic recurrences remain in Chapter 4; this chapter concentrates on tensor layouts, pipeline state, compile-time selection, and synchronization in the source.

5.3 Kernel configuration

Configuration happens in __call__ before launch. The thread/register plan (lines 206–232):

Listing 5.1: flash_fwd_sm90.py lines 206–232 (trimmed): warpgroup and register budgeting.

tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma()
self.num_mma_threads = tiled_mma_qk.size
self.num_threads_per_warp_group = 128
self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group
assert self.num_wg_mma in [1, 2, 3]
self.num_threads = self.num_threads_per_warp_group * (self.num_wg_mma + 1)
self.num_producer_threads = 32
self.num_mma_regs, self.num_producer_regs = {1: (256, 56), 2: (240, 24), 3: (160, 32)}[
    self.num_wg_mma
]
self.use_scheduler_barrier = (
    (self.num_wg_mma >= 2 and self.tile_hdim <= 128)
    if const_expr(self.intra_wg_overlap)
    else (self.num_wg_mma == 2)
)
...
self.rescale_O_before_gemm = self.tile_hdimv > 128 and self.intra_wg_overlap

Everything from Chapter 4 is visible: \(+1\) producer warpgroup; the setmaxnreg budgets per consumer count; pingpong (use_scheduler_barrier) only for \(\ge 2\) consumer warpgroups and \(d \le 128\); early-O-rescale for \(d_v > 128\).

The two tiled MMAs (lines 96–118) encode the GEMM shapes and operand sources:

Listing 5.2: flash_fwd_sm90.py lines 96–118: the two WGMMAs.

def _get_tiled_mma(self):
    tiled_mma_qk = sm90_utils_basic.make_trivial_tiled_mma(
        self.dtype, self.dtype,
        warpgroup.OperandMajorMode.K,   # Q: K-major (contiguous in head dim)
        warpgroup.OperandMajorMode.K,   # K: K-major
        Float32,                        # FP32 accumulate
        atom_layout_mnk=(self.tile_m // 64, 1, 1),
        tiler_mn=(64, self.tile_n),
    )
    tiled_mma_pv = sm90_utils_basic.make_trivial_tiled_mma(
        self.dtype, self.dtype,
        warpgroup.OperandMajorMode.K,   # P: K-major
        warpgroup.OperandMajorMode.MN,  # V: MN-major logical view; no data move
        Float32,
        atom_layout_mnk=(self.tile_m // 64, 1, 1),
        tiler_mn=(64, self.tile_hdimv),
        a_source=warpgroup.OperandSource.RMEM   # <-- RS-WGMMA
        if self.mma_pv_is_rs
        else warpgroup.OperandSource.SMEM,
    )
    return tiled_mma_qk, tiled_mma_pv

Note GEMM1’s B operand: the 16-bit kernel does not move \(\mathbf{V}\) into a second physical tile. It constructs a transposed logical view (layout_utils.transpose_view(sV)) and selects the corresponding MN-major WGMMA interpretation. This is legal for FP16/BF16 WGMMA but the transpose immediates are unavailable for Hopper FP8 WGMMA, which is why the FP8 C++ kernel performs an in-kernel rearrangement (Chapter 6).

5.3.1 Compile-time variant selection

Most high-level conditionals in this class do not produce divergent branches in the device loop. Their values are known during specialization and const_expr removes the unused path. Several flags change storage or instruction selection rather than enabling a small optional operation:

Consequently, performance comparisons between two configurations compare different generated kernels, not two runtime modes of one binary. The flag combination must also be evaluated as a unit: disabling RS adds sP, for example, which changes both SMEM capacity and the synchronization sequence.

5.4 Shared memory plan

Listing 5.3: flash_fwd_sm90.py lines 131–155 (trimmed): shared storage.

# 1 stage * 2 for Q pipeline (full + empty), num_stages*2 for K, num_stages*2 for V
mbar_ptr_Q_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2]
mbar_ptr_K_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2]
mbar_ptr_V_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2]

@cute.struct
class SharedStorageQKV:
    mbar_ptr_Q: mbar_ptr_Q_struct
    mbar_ptr_K: mbar_ptr_K_struct
    mbar_ptr_V: mbar_ptr_V_struct
    sV: sV_struct
    sQ: sQ_struct
    sK: sK_struct
    sP: sP_struct          # only sized > 0 when not mma_pv_is_rs

Sizes for the canonical \(d = 128\), \(B_M = B_N = 128\), 2-stage BF16 config: \(\texttt{sQ} = 128{\times}128{\times}2\,\text{B} = 32\) KB; sK and sV \(= 2 \times 32 = 64\) KB each; sP \(= 0\) (RS mode); mbarriers \(= (2 + 4 + 4) \times 8\) B. Total \(\approx 160\) KB of the 228 KB budget — the remainder allows the DSL runtime’s reserved space and keeps the launch at 1 block/SM. The \(\mathbf{O}\) epilogue tile does not appear: it reuses sQ’s memory (line 535: sO = storage.sQ.get_tensor(...)). Each buffer is 1024-byte aligned (TMA requirement for the largest swizzle).

The K/V smem layouts are built by sm90_utils.make_smem_layout(dtype, ROW_MAJOR, (tile_n, hdim), stage): a 128-byte-swizzled core-matrix atom tiled to the full shape, with the stage index appended as the outermost mode — the canonical WGMMA descriptor layout, matching the swizzle mode the TMA descriptor was created with.

5.4.1 Tensor views and WGMMA partitions

Before TMA atoms are constructed, layout_utils.select reorders the logical modes of fixed-length Q/O from \((B,N,H,D)\) to \((N,D,H,B)\) and K/V in the same manner. This is a view transformation: the engine pointer is unchanged, while the layout determines the address associated with a logical coordinate. cute.local_tile then fixes the batch, head, and block coordinates to expose the two-dimensional tile consumed by a copy atom.

tiled_mma.get_slice(warp_group_idx) selects one 64-row MMA atom, and partition_fragment_ABC derives the per-thread views required by that atom. These operations do not copy tensor elements. They compose layouts that specify which descriptor, register fragment, or accumulator coordinates belong to each thread. For \(B_M=128\), two slices cover disjoint 64-row ranges of Q and O while referencing the same K and V stages.

The stage mode is part of the SMEM layout, so selecting B_idx=state.index changes the descriptor base for a WGMMA without rebuilding the tiled MMA. Similarly, sVt changes only the logical mode order. Its legality depends on the MN-major/K-major capability encoded by the 16-bit WGMMA descriptor; it is not a shared-memory transpose operation.

5.5 Kernel entry: pipelines and role dispatch

Listing 5.4: flash_fwd_sm90.py lines 441–516 (trimmed): descriptor prefetch and pipeline construction.

warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
if warp_idx == 0:                       # prefetch TMA descriptors
    for tma_atom in (tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_O):
        if const_expr(tma_atom is not None):
            cpasync.prefetch_descriptor(tma_atom)

smem = cutlass.utils.SmemAllocator()
storage = smem.allocate(SharedStorage)

tma_warp = ThreadCooperativeGroup(1)                      # producer: 1 thread signals
mma_warps = ThreadCooperativeGroup(self.num_mma_threads // cute.arch.WARP_SIZE)

pipeline_q = pipeline_custom.PipelineTmaAsync.create(
    barrier_storage=mbar_ptr_Q, num_stages=1,
    producer_group=tma_warp, consumer_group=mma_warps,
    tx_count=self.tma_copy_bytes["Q"], defer_sync=True,
)
pipeline_k = pipeline_custom.PipelineTmaAsync.create(
    barrier_storage=storage.mbar_ptr_K.data_ptr(), num_stages=self.num_stages,
    producer_group=tma_warp, consumer_group=mma_warps,
    tx_count=self.tma_copy_bytes["K"], defer_sync=True,
)
pipeline_v = pipeline_custom.PipelineTmaAsync.create(... tx_count=self.tma_copy_bytes["V"] ...)

Three pipelines: \(\mathbf{Q}\) (1 stage — loaded once per work tile), \(\mathbf{K}\) and \(\mathbf{V}\) (num_stages each, with the separate barriers defined in Section 4.3). tx_count is the TMA transaction byte count the mbarrier expects per stage — e.g. for K, \(B_N \cdot d \cdot 2\) bytes. The consumer group is sized in warps: every MMA warp arrives on empty[s] at release.

5.5.1 Pipeline cursors and phase changes

Each PipelineTmaAsync owns two mbarriers per stage: a full barrier for producer-to-consumer readiness and an empty barrier for consumer-to-producer reuse. The corresponding producer and consumer states contain three logical fields: an iteration count, index in \([0,s)\), and the phase expected for that index. Advancing a state performs

and changes the phase when the index wraps. The phase distinguishes two successive uses of the same physical stage.

The main state transitions are:

TMA completion, rather than the producer thread, satisfies the full barrier’s byte expectation. Conversely, the empty barrier counts the consumer group, so the stage is not reusable after only one consumer warpgroup has finished. The one-stage Q pipeline follows the same phase protocol across work tiles; the producer and consumer explicitly toggle their Q phases because the Q cursor does not advance inside the KV loop.

K and V use cloned but intentionally offset cursors in the overlapped path. At a steady-state consumer step, the K cursor identifies the current score tile and the V cursor identifies the preceding probability tile. Treating those cursors as one “KV state” would select the wrong V stage after the prologue.

Role dispatch with register reallocation (lines 580–609):

Listing 5.5: flash_fwd_sm90.py lines 580–609 (trimmed): warp specialization.

if warp_idx < 4:  # Producer warpgroup
    cute.arch.setmaxregister_decrease(self.num_producer_regs)   # e.g. -> 24
    self.load(mQ, mK, mV, sQ, sK, sV, tma_atom_Q, tma_atom_K, tma_atom_V,
              pipeline_k, pipeline_v, pipeline_q, ...)
else:             # Consumer warpgroups
    cute.arch.setmaxregister_increase(self.num_mma_regs)        # e.g. -> 240
    tidx, _, _ = cute.arch.thread_idx()
    tidx = tidx - 128
    self.mma(tiled_mma_qk, tiled_mma_pv, mO, mLSE, sQ, sK, sVt, sP, sO, ...)

5.6 The producer: load

Within the producer warpgroup only warp 0 issues TMA (is_load_warp); the other three warps idle (they exist because register reallocation operates on whole warpgroups). The producer runs its own copy of the tile-scheduler loop, so producer and consumers agree on the tile sequence without communicating. Core of the steady-state loop (lines 833–871), which implements the K-before-V staggering that feeds the 2-stage consumer:

Listing 5.6: flash_fwd_sm90.py lines 784–871 (trimmed): producer loop with intra-warpgroup overlap staggering.

# First iteration: load K on pipeline_k, Q on pipeline_q
pipeline_k.producer_acquire(kv_producer_state)
load_K(block=n_block, producer_state=kv_producer_state, ...)
if warp_idx_in_wg == 0:
    pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase)
    load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0))
    q_producer_phase ^= 1
...
# Steady state (intra_wg_overlap): K of block j goes out together with
# V of block j+1 (the previous block), matching consumer order.
for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1):
    n_block_prev = n_block_max - i - 1
    n_block = n_block_prev - 1
    kv_producer_state_prev = kv_producer_state.clone()
    kv_producer_state.advance()
    pipeline_k.producer_acquire(kv_producer_state)
    load_K(block=n_block, producer_state=kv_producer_state, ...)
    pipeline_v.producer_acquire(kv_producer_state_prev)
    load_V(block=n_block_prev, producer_state=kv_producer_state_prev, ...)
# Tail: V of the last (lowest) block
pipeline_v.producer_acquire(kv_producer_state)
load_V(block=n_block, producer_state=kv_producer_state, ...)

The stagger serves the overlapped consumer schedule (Section 4.5), which consumes K\(_{j+1}\) before V\(_j\). Issuing loads in consumption order maximizes the chance that the wait in step 1 (K) never stalls. Each load_K is a one-line TMA call through the closure built by copy_utils.tma_get_copy_fn; producer_commit is a no-op for TMA pipelines (completion is signaled by the hardware transaction count).

Also visible in the full file: the cp.async fallback path when TMA cannot be used (paged KV cache with page size \(\ne B_N\); PackGQA Q loads with irregular tile shapes), in which case all 128 producer threads participate (PipelineCpAsync), and the producer register budget is raised to 40.

5.6.1 From a scheduled tile to copy coordinates

The producer and consumers instantiate the same tile-scheduler class and independently obtain \((m\_block, head\_idx, batch\_idx)\). No queue entry is passed between roles. Correctness therefore requires both branches to execute the same scheduler progression, including early exits for empty or sparse tiles; otherwise one branch waits on a pipeline generation that the other branch never produces or releases.

For each work tile, SeqlenInfoQK first shifts the tensor engines to the selected batch’s Q and K/V bases. The KV head is \(\lfloor head\_idx/h_q\rfloor\) for ordinary GQA, or is already represented by the packed head coordinate under PackGQA. cute.local_tile then forms Q at coordinate \((m\_block,0)\) and K/V with a free KV-block coordinate. The copy closure receives n_block at issue time, so one preconstructed TMA atom serves every KV iteration of the work tile.

Bounds protection occurs at two different levels:

• A TMA tensor map supplies zero fill for coordinates outside its declared tensor extent. This prevents an invalid load but does not implement attention masking; the corresponding score columns must still become \(-\infty\) before softmax.

• The cp.async and paged-KV paths construct explicit predicates and page-table coordinates. They must avoid forming an invalid page-table address even when a tile contains no valid KV block.

BlockInfo computes n_block_min/max once per work tile, so the steady producer loop issues only the copies that the consumer loop will request.

5.7 The consumer: mma

The consumer partitions SMEM operands into WGMMA fragments once, then loops over work tiles. Setup (lines 966–1007, trimmed):

Listing 5.7: flash_fwd_sm90.py lines 966–1007 (trimmed): consumer setup.

warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group)
wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(warp_group_idx))
wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(warp_group_idx))
_, tSrQ, tSrK = sm90_utils.partition_fragment_ABC(
    wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ, sK)
mma_qk_fn = partial(sm90_utils.gemm_zero_init, tiled_mma_qk,
                    (self.tile_m, self.tile_n), tSrQ, tSrK)
acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC(
    wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt)
mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt)
...
self.mma_init()          # WG1 arrives on the pingpong barrier chain
softmax = Softmax.create(softmax_scale_log2,
    num_rows=acc_O.shape[0][0] * acc_O.shape[1], softmax_scale=softmax_scale)

acc_O is the per-thread slice of the \(B_M \times d_v\) FP32 output accumulator, alive across the whole K/V loop. tOrP is the register fragment through which softmax results feed GEMM1 in RS mode. The B_idx argument that mma_qk_fn/mma_pv_fn take at call time selects the pipeline stage of sK/sVt.

5.7.1 Accumulator ownership and initialization

gemm_zero_init creates a new FP32 score accumulator for every QK GEMM; no score value is carried between KV tiles. After masking and online_softmax, that same distributed register storage contains \(\tilde{\mathbf{P}}\). reshape_acc_to_frgA changes its register layout from the WGMMA accumulator interpretation to the operand-A interpretation, and utils.cvt_f16 writes the converted values to tOrP. The reshape is a register-coordinate transformation; the conversion is the operation that changes element type.

acc_O, by comparison, persists across the KV loop. The peeled-loop helpers carry whether an earlier PV GEMM has initialized it. The first PV uses the zero-initializing WGMMA form, while later PV operations accumulate. The Softmax object also persists: its row_max and row_sum fragments follow the same row ownership as acc_O. Keeping these objects within one consumer warpgroup avoids cross-warpgroup reductions because different consumer warpgroups own disjoint query rows.

The main loop dispatches the three regimes of Section 4.7; here is the skeleton with the overlap path (lines 1110–1189, heavily trimmed):

Listing 5.8: flash_fwd_sm90.py lines 1098–1189 (trimmed): consumer main loop.

n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase)

# First iteration with seqlen masking ("first half block": GEMM0+softmax only)
kv_consumer_state = process_first_half_block(
    n_block=n_block_max - 1, ..., mask_fn=partial(mask_fn, mask_mod=self.mask_mod),
    is_first_block=True)
n_block_max -= 1
# Next couple of iterations with causal masking
if const_expr(self.is_causal or self.is_local):
    n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
        seqlen, m_block, n_block_min)
    for n_tile in cutlass.range(n_block_max - n_block_min_causal_local_mask, unroll=1):
        kv_consumer_state = mma_one_n_block(kv_consumer_state,
            n_block=n_block_max - 1 - n_tile, ...,
            mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False))
    n_block_max = cutlass.min(n_block_max, n_block_min_causal_local_mask)
# The remaining iterations have no masking
for n_tile in cutlass.range(n_block_max - n_block_min_before_local_mask, unroll=1):
    kv_consumer_state = mma_one_n_block(kv_consumer_state,
        n_block=n_block_max - 1 - n_tile, ..., mask_fn=...)
# Release Q so the producer can start the next tile's Q load
pipeline_q.consumer_release_w_index(0)
# Last "half" iteration: the final PV GEMM
kv_consumer_state = process_last_half_block(kv_consumer_state=kv_consumer_state,
                                            zero_init=not O_should_accumulate)
...
row_scale = softmax.finalize(sink_val=sink_val)
softmax.rescale_O(acc_O, row_scale)
self.epilogue(acc_O, softmax.row_sum, mO, mLSE, sO, ...)

The descending value of n_block_max is also the boundary between specialized loop segments. The first call includes sequence-tail masking; the causal/local segment passes a closure with sequence masking disabled; the steady segment omits the causal mask call. These are separate generated loops, so an unmasked iteration does not evaluate a per-iteration mode branch. Local attention adds a final left-boundary segment after the unmasked range.

process_first_half_block establishes the cursor offset and produces the first P fragment without issuing PV. Each middle call consumes the preceding V stage while producing the next P fragment. After all QK operations have completed, pipeline_q.consumer_release_w_index(0) releases Q; the final half-block needs only P and V. This placement allows the producer to begin loading Q for its next scheduled work tile while the current consumers drain the last PV operation and execute the epilogue.

5.8 Core operation: one \(n\)-block with intra-warpgroup overlap

This function implements the FA3 forward inner step corresponding to the five-stage schedule in Section 4.5. It is quoted nearly in full.

Listing 5.9: flash_fwd_sm90.py lines 1411–1477 (lightly trimmed): mma_one_n_block_intrawg_overlap.

@cute.jit
def mma_one_n_block_intrawg_overlap(
    self, smem_pipe_read, n_block, mma_qk_fn, mma_pv_fn,
    pipeline_k, pipeline_v, acc_O, tOrP, smem_copy_params, softmax, seqlen,
    scores_scale=None, score_mod_fn=None, mask_fn=None, check_inf=True,
):
    smem_pipe_read_v = smem_pipe_read.clone()   # V lags one iteration behind K
    smem_pipe_read.advance()
    pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read))
    self.warp_scheduler_barrier_sync()          # pingpong: acquire TC issue turn
    # S = Q @ K.T   (GEMM0 of iteration j+1, async)
    acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1)
    # RescaleOBeforeGemm: rescale O while QK GEMM is in flight, before PV GEMM
    if const_expr(self.rescale_O_before_gemm):
        softmax.rescale_O(acc_O, scores_scale)
    pipeline_v.consumer_wait(smem_pipe_read_v, pipeline_v.consumer_try_wait(smem_pipe_read_v))
    # O += P @ V    (GEMM1 of iteration j, async)
    mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1)
    self.warp_scheduler_barrier_arrive()        # pingpong: hand TCs to the other WG
    warpgroup.wait_group(1)                     # wait GEMM0 (older group) only
    pipeline_k.consumer_release(smem_pipe_read)

    if const_expr(score_mod_fn is not None):
        score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen)
    if const_expr(mask_fn is not None):
        mask_fn(acc_S=acc_S, n_block=n_block)

    row_scale = softmax.online_softmax(acc_S, check_inf=check_inf)  # MUFU work,
    warpgroup.wait_group(0)                     # ... overlapped with GEMM1
    pipeline_v.consumer_release(smem_pipe_read_v)
    tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S)
    tOrP_cur = tOrP if const_expr(self.mma_pv_is_rs) else cute.make_rmem_tensor_like(...)
    # the "to(self.dtype)" conversion fails to vectorize for block sizes other
    # than 128 x 128 ... Call PTX directly instead.
    utils.cvt_f16(tOrP_acc, tOrP_cur)           # paired cvt.f32x2 -> bf16x2
    if const_expr(not self.mma_pv_is_rs):
        tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur)
        cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP)
    if const_expr(not self.rescale_O_before_gemm):
        softmax.rescale_O(acc_O, row_scale)
    else:
        scores_scale.store(row_scale.load())
    if const_expr(not self.mma_pv_is_rs):
        # Fence and barrier to make sure smem store is visible to WGMMA
        cute.arch.fence_view_async_shared()
        cute.arch.sync_warp()
    return smem_pipe_read

Observations, in execution order:

• Two pipeline cursors. smem_pipe_read_v (V, iteration \(j\)) trails smem_pipe_read (K, iteration \(j{+}1\)) — the code-level incarnation of the half-iteration offset.

• wg_wait=-1 twice, then wait_group(1). Both GEMMs are put in flight before any waiting. wait_group(1) returns when at most one of the two most recently committed groups remains pending and every older group has completed. Here that guarantees GEMM0 is complete, so K’s stage can be recycled while GEMM1 may still be pending.

• Pingpong placement. barrier_sync before GEMM0 issue, barrier_arrive right after GEMM1 issue: the warpgroup holds the Tensor Core issue turn during this interval, then yields.

• Softmax sits between wait_group(1) and wait_group(0) — the MUFU/ALU work runs precisely while GEMM1 is outstanding. Placing the softmax work in this interval creates the intra-warpgroup overlap.

• FP32\(\to\)BF16 conversion by inline PTX. utils.cvt_f16 works around a DSL vectorization gap (the comment is from the source): conversions must go 2-at-a-time (cvt.rn.bf16x2.f32) or the kernel wastes ALU issue slots.

• The non-RS path (only for \(d \le 96\)’s \(192\)-row tiles) stages \(\mathbf{P}\) through sP with an async-proxy fence; only sync_warp is sufficient because each warp reads the same SMEM addresses that it wrote; no cross-warp exchange occurs.

5.8.1 Why each wait and release is safe

In this path, each GEMM helper issues and commits a WGMMA group. The wg_wait=-1 argument suppresses an immediate wait; it does not indicate that the instruction has completed. After QK is committed first and PV second, wait_group(1) waits until no more than one committed group remains in flight. QK is therefore complete, while PV may still be executing. This proves three permissions simultaneously: acc_S may be read by ALU instructions, the K descriptor is no longer in use, and the K pipeline stage may be released.

The subsequent softmax accesses acc_S, row statistics, and P-conversion registers, not the PV destination acc_O. It can therefore execute while PV updates acc_O. wait_group(0) establishes completion before V is released or acc_O is read or rescaled. In the early-rescale variant, the saved factor from the preceding softmax step is applied before PV is issued, which also avoids concurrent ALU and WGMMA writes to acc_O.

The scheduler named barrier covers only the WGMMA issue interval. Arrival on the next warpgroup’s barrier occurs after PV issue, so the next warpgroup may issue its WGMMA operations while the current warpgroup performs waits and softmax. Completion ordering remains the responsibility of the WGMMA group waits; a named-barrier arrival is not evidence that either GEMM has finished.

The non-overlapped variant mma_one_n_block (lines 1349–1408) is the same code with one cursor, wait_group(0) after GEMM0, and softmax strictly between GEMM0 and GEMM1 — pingpong then performs all of the latency hiding (this is the path used when intra_wg_overlap=False).

The loop peel halves:

Listing 5.10: flash_fwd_sm90.py lines 1270–1346 (trimmed): the peeled half-iterations.

def first_half_block_overlap(self, n_block, mma_qk_fn, kv_consumer_state,
                             pipeline_k, tOrP, ..., is_first_block=False):
    """Processes the first half block when using intra-warpgroup-overlap"""
    pipeline_k.consumer_wait(kv_consumer_state, ...)
    acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=0)  # sync wait here
    pipeline_k.consumer_release(kv_consumer_state)
    ...
    mask_fn(acc_S, n_block=n_block, mask_seqlen=True)
    row_scale = softmax.online_softmax(acc_S, is_first=is_first_block)
    ...
    tOrP_cur.store(tOrP_acc.load().to(self.dtype))
    ...

def last_half_block_overlap(self, kv_consumer_state, pipeline_v, mma_pv_fn,
                            zero_init, scores_scale=None, softmax=None, acc_O=None):
    """Processes the final PV GEMM when using intra-warpgroup-overlap"""
    if const_expr(self.rescale_O_before_gemm):
        softmax.rescale_O(acc_O, scores_scale)
    pipeline_v.consumer_wait(kv_consumer_state, ...)
    mma_pv_fn(B_idx=kv_consumer_state.index, zero_init=zero_init, wg_wait=0)
    pipeline_v.consumer_release(kv_consumer_state)
    kv_consumer_state.advance()
    return kv_consumer_state

5.9 Pingpong in code: the scheduler barriers

Listing 5.11: flash_fwd_sm90.py lines 1480–1545 (trimmed): pingpong barrier helpers.

@cute.jit
def mma_init(self):
    warp_group_idx = utils.canonical_warp_group_idx(sync=False)
    if const_expr(self.use_scheduler_barrier):
        if warp_group_idx == 1:      # WG1 kick-starts the chain
            cute.arch.barrier_arrive(
                barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1),
                number_of_threads=2 * self.num_threads_per_warp_group)

def warp_scheduler_barrier_sync(self):
    if const_expr(self.use_scheduler_barrier):
        cute.arch.barrier(       # wait for MY barrier
            barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) - 1
                       + utils.canonical_warp_group_idx(sync=False),
            number_of_threads=2 * self.num_threads_per_warp_group)

def warp_scheduler_barrier_arrive(self):
    if const_expr(self.use_scheduler_barrier):
        cur_wg = utils.canonical_warp_group_idx(sync=False) - 1
        next_wg = 1 - cur_wg if const_expr(self.num_wg_mma == 2) \
                  else (cur_wg + 1) % self.num_wg_mma
        cute.arch.barrier_arrive(   # signal the NEXT warpgroup's barrier
            barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg,
            number_of_threads=2 * self.num_threads_per_warp_group)

Each scheduler-barrier generation expects \(2 \times 128\) participating thread arrivals: 128 from the owner warpgroup’s barrier.sync and 128 from the predecessor’s barrier.arrive. The count specifies this total, while the branches above determine which threads participate. With 3 consumer warpgroups, the same protocol forms a round-robin issue order. The hardware automatically reinitializes a named barrier after completion, but the control flow must still prevent one warpgroup from entering the next logical generation prematurely.

5.10 Softmax module

The per-row state lives in two register arrays sized num_rows \(= 2 \times\) (rows per thread of the accumulator layout). Online step (softmax.py lines 126–190, trimmed):

Listing 5.12: softmax.py lines 150–190 (trimmed): online_softmax.

for r in cutlass.range(cute.size(row_max), unroll_full=True):
    acc_S_row = acc_S_mn[r, None].load()                      # (n_block_size)
    row_max_cur = utils.fmax_reduce(acc_S_row,
        init_val=row_max[r] if cutlass.const_expr(not is_first) else None, arch=arch)
    row_max_cur = cute.arch.warp_reduction_max(row_max_cur, threads_in_group=4)
    row_max_prev = row_max[r]
    row_max[r] = row_max_cur
    if cutlass.const_expr(check_inf):
        row_max_cur = 0.0 if row_max_cur == -Float32.inf else row_max_cur
    row_max_cur_scaled = row_max_cur * scale_log2
    acc_S_row_exp = cute.math.exp2(acc_S_row * scale_log2 - row_max_cur_scaled,
                                   fastmath=True)
    if cutlass.const_expr(is_first):
        acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp, init_val=None, arch=arch)
        row_scale[r] = 1.0
    else:
        row_scale[r] = cute.math.exp2((row_max_prev - row_max_cur) * scale_log2,
                                      fastmath=True)
        acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp,
                                          init_val=row_sum[r] * row_scale[r], arch=arch)
    row_sum[r] = acc_S_row_sum
    acc_S_mn[r, None].store(acc_S_row_exp)          # P overwrites S in-place

Note warp_reduction_max(..., threads_in_group=4): the quad reduction for the max, but no quad reduction for the sum (deferred to finalize). The finalize step:

Listing 5.13: softmax.py lines 192–227 (trimmed): finalize — deferred reduction, safe reciprocal, LSE.

# quad reduction for row_sum, deferred from each iteration
row_sum.store(utils.warp_reduce(row_sum.load(), operator.add, width=4))
for r in cutlass.range(cute.size(row_sum), unroll_full=True):
    if cutlass.const_expr(sink_val is not None):
        row_sum[r] += cute.math.exp2(sink_val_cur * LOG2_E - row_max[r] * scale_log2, ...)
    # if row_sum is zero or nan, set acc_O_mn_row to 1.0
    acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r]
    row_scale[r] = (cute.arch.rcp_approx(row_sum[r]
                    if not acc_O_mn_row_is_zero_or_nan else 1.0)) * final_scale
    row_sum[r] = ((row_max[r] * scale_log2 + cute.math.log2(row_sum_cur)) * LN2
                  if not acc_O_mn_row_is_zero_or_nan else -Float32.inf)
    # ^ row_sum now holds the LSE in natural-log units

5.11 Epilogue in code

From the base class (flash_fwd.py lines 330–449, trimmed to the TMA path):

Listing 5.14: flash_fwd.py lines 346–417 (trimmed): epilogue — registers \(\to\) SMEM \(\to\) TMA store.

rO = cute.make_fragment_like(acc_O, self.dtype)
rO.store(acc_O.load().to(self.dtype))                # FP32 -> BF16
cute.arch.barrier(barrier_id=int(NamedBarrierFwd.Epilogue),
                  number_of_threads=self.num_epilogue_threads)  # V reads done
smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx)
cute.copy(smem_copy_atom_O, taccOrO, taccOsO)        # regs -> sO (= sQ memory)
# LSE store: only threads owning column 0 write, predicated on seqlen
if taccOcO[0][1] == 0:
    for m in cutlass.range(cute.size(taccOgLSE.shape[1]), unroll_full=True):
        if t0accOcO[m, 0][0] < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0]:
            taccOgLSE[m, 0] = lse[m]
# O store via TMA, issued by a single warp
cute.arch.fence_view_async_shared()                  # make sO visible to TMA
cute.arch.barrier_arrive(barrier_id=int(NamedBarrierFwd.Epilogue),
    number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE)
if warp_idx == 4:
    cute.arch.barrier(...)                           # all consumers arrived
    store_O()                                        # cp.async.bulk.tensor S2G
    cute.arch.cp_async_bulk_commit_group()
    cute.arch.cp_async_bulk_wait_group(0, read=True)

The first named barrier prevents an aliasing race: sO aliases sQ, and (in the shared-QV configuration Q_in_regs) V — no thread may write the epilogue tile until every consumer warpgroup has finished its last GEMM reading that memory. Its participant count is num_epilogue_threads, the consumer threads that reach this point. The next logical generation uses a larger count, num_epilogue_threads + 32: all consumers arrive after writing their fragments, while the 32-thread store warp waits. The named barrier supplies the cross-thread rendezvous; fence_view_async_shared() supplies the distinct generic-to-async proxy ordering required before the TMA store reads sO.

5.12 Differences from the paper’s CUTLASS C++ kernels

For readers cross-referencing hopper/ (C++) with flash_attn/cute/ (DSL):

• Branding. The DSL package is distributed as FlashAttention-4 (package name flash-attn-4). Its SM90 path is a port of the FA3 C++ kernels — the file headers describe it as a reimplementation of flash_fwd_kernel_sm90.h from CUTLASS C++ to the CuTe DSL — extended with newer features (FlexAttention-style score_mod and mask_mod, block sparsity, learnable sinks, paged KV).

• Clusters/multicast. C++ FA3 uses cluster size 2 + TMA multicast for K/V on the forward; the DSL SM90 forward runs cluster (1,1) (multicast marked TODO).

• Persistence. The C++ forward offers a persistent kernel with a dynamic tile scheduler; the DSL SM90 forward uses one-block-per-tile scheduling with LPT reordering for causal (is_persistent=False); its backward keeps the LPT + L2-swizzle static order.

• FP8. The C++ hopper/ tree implements the paper’s FP8 forward (e4m3 instantiations, in-kernel V transpose). In the DSL, FP8 attention is implemented for SM100 (Blackwell), and interface.py asserts FP8 is only supported on SM100 ... for FA4 CuTe — FP8-on-Hopper remains C++-only (Chapter 6).

• dQ accumulation. C++ FA3’s backward accumulates dQ into global FP32 with red.add through TMA bulk-reduce as well (copy_sm90_bulk_reduce.hpp); the paper text describes atomic accumulation — see Section 8.7 for the exact DSL mechanism.

• Softcap. The C++ kernels support tanh softcapping; in the DSL this is subsumed by the general score_mod mechanism.

• Scheduler-barrier constant. The DSL enables pingpong for \(d \le 128\) with overlap on (line 220); C++ FA3 similarly restricts it per head dim, and both disable it entirely for 1 consumer warpgroup.

• Tile sizes. Re-tuned per implementation: e.g. forward \(d{=}128\) non-causal is \(128\times176\) in C++ tile_size.h but \(128\times128\) in the DSL (_tile_size_fwd_sm90), whose comments document the different Python-side register/SMEM tradeoffs.

6 FP8 Forward: Layouts, Block Quantization, Incoherent Processing

FP8 doubles Hopper’s Tensor Core throughput (989 \(\to\) 1979 TFLOP/s dense), and FA3’s FP8 forward reaches 1.3 PFLOP/s [1, 13]. Reaching that level requires solving two problems that BF16 kernels never see: operand layout (this chapter’s first half) and quantization error (second half).

Implementation status: the FP8 kernels evaluated in the paper live in the CUTLASS C++ hopper/ tree (e4m3 instantiations of mainloop_fwd_sm90_tma_gmma_ws.hpp). The CuTe DSL Python implementation supports FP8 on SM100 (Blackwell) only — interface.py raises "FP8 is only supported on SM100 (compute capability 10.x) for FA4 CuTe" and FP8 backward is unimplemented ("FA4 CuTe FP8 backward is not supported yet (forward-only)"). This chapter therefore describes the Hopper FP8 design per the paper and the C++ code, and notes the DSL’s SM100 equivalents where useful.

6.1 FP8 formats and accumulation

Hopper Tensor Cores support two 8-bit floating formats: e4m3 (4 exponent bits, 3 mantissa bits; range \(\pm 448\)) and e5m2 (\(\pm 57344\), 2 mantissa bits). Attention inputs use e4m3 (precision matters more than range once scaled). FP8 WGMMA accumulates in FP32; however, on Hopper the FP8-path accumulator has been measured to behave as an FP22 format (e8m13) — 13 effective mantissa bits rather than FP32’s 23 [23] — which is one reason long-sum GEMMs (large \(N\)) accumulate slightly more error than an FP16-input WGMMA would; the same reduced precision applies to the summation that combines partial results when the K extent exceeds the instruction’s native 32. This is a caveat inherited by any FP8 attention (Caveat 13.4).

6.2 The layout problem: K-major operands and the in-kernel V transpose

For 16-bit inputs, WGMMA accepts SMEM operands in MN-major or K-major order, so the forward’s \(\mathbf{P}\mathbf{V}\) GEMM can consume a transposed view of V (Section 5.3). FP8 WGMMA, however, requires both operands K-major. For GEMM1, \(\mathbf{O}= \tilde{\mathbf{P}}\mathbf{V}\) with \(\tilde{\mathbf{P}}\in \mathbb{R}^{B_M \times B_N}\): the reduction (“K”) dimension is \(B_N\), i.e. the sequence dimension of V. But \(\mathbf{V}\) arrives from HBM row-major in the head dimension (\(d_v\)-contiguous) — MN-major from the GEMM’s perspective. Options: transpose V in a pre-kernel (extra HBM round trip — rejected), or transpose it inside the kernel between TMA load and GEMM use. FA3 does the latter:

• The producer loads \(\mathbf{V}_j\) (as-is) into SMEM.

• Before first use, the producer warpgroup performs an SMEM\(\to\)RMEM\(\to\)SMEM transpose of the tile into a second buffer \(\mathbf{V}^\top_j\), using ldmatrix (LDSM) with the .trans option and paired stmatrix (STSM), plus byte-level prmt (permute) instructions to shuffle the four 8-bit values inside each 32-bit register (ldmatrix transposes at 16-bit granularity, so FP8 needs the extra byte permutation) [13]. The paper credits the in-kernel transpose idea to the cuDNN team.

• The transpose is folded into the pipeline: in the warp-specialized kernel it is executed under the producer/consumer barrier protocol so it overlaps with GEMMs of other iterations (the paper pipelines it one iteration ahead of use).

A second layout issue is that FP8 WGMMA’s accumulator register layout differs from the FP16 one. In BF16, the FP32 accumulator fragment of GEMM0 can be converted and immediately serve as the RS operand-A fragment of GEMM1 — the layouts line up thread-by-thread. In FP8 they do not: each thread holds the required values at positions that do not match the FP8 operand-A layout. The solution, developed in the Colfax FP8 study [13] and adopted by FA3, is to permute bytes within each thread after the FP32\(\to\)FP8 conversion using __byte_perm/prmt, and to exchange 4-byte words between the 4 threads of a quad (__shfl_sync with per-lane source maps) until the FP8 operand-A register layout is satisfied (the m64nNk32 matrix-A fragment layout in the PTX ISA’s WGMMA section). The net effect is no additional memory traffic; the permutes and shuffles can execute on the ALUs while GEMMs remain in flight. (On SM100/Blackwell this layout restriction is removed: tcgen05.mma takes accumulators in Tensor Memory and both operand layouts natively; the DSL’s FP8 support therefore targets SM100 first.)

6.3 Block quantization

Naive FP8 casts \(\mathbf{Q}, \mathbf{K}, \mathbf{V}\) with one scale per tensor (per head at best). Transformer activations, especially \(\mathbf{K}\), exhibit large per-channel outliers; a single scale can restrict most entries to a small subset of the representable values. FA3 instead quantizes per block: each \(B_M \times d\) tile of \(\mathbf{Q}\) and each \(B_N \times d\) (resp. \(B_N \times d_v\)) tile of \(\mathbf{K}\) (resp. \(\mathbf{V}\)) carries one FP32 scale (q_descale, k_descale, v_descale in the API, per (batch, kv-head) or finer). Because FlashAttention already processes exactly these tiles, the descales fold into existing per-tile scalar work:

• \(\mathbf{S}= \mathbf{Q}_i\mathbf{K}_j^\top\) picks up factor \(\delta_Q \delta_K\), folded into the softmax scale for tile \((i,j)\) by multiplication into softmax_scale_log2.

• \(\tilde{\mathbf{P}}\mathbf{V}_j\) picks up \(\delta_V\), folded into the O rescale already performed per iteration.

The descales therefore require no separate matrix operation in the attention kernel and add only per-tile scalar work. The quantization itself can be fused into the epilogue of the producing operation (e.g. the preceding linear layer or rotary embedding), although that conversion still performs additional arithmetic.

6.4 Incoherent processing: the Hadamard transform

Block scales handle block-level outliers but not a small number of extreme-magnitude entries inside one block. FA3 borrows incoherent processing from the quantization literature (QuIP [6]): pick a random orthogonal matrix \(M \in \mathbb{R}^{d\times d}\) and replace \(\mathbf{Q}\leftarrow \mathbf{Q}M\), \(\mathbf{K}\leftarrow \mathbf{K}M\). Since \((\mathbf{Q}M)(\mathbf{K}M)^\top = \mathbf{Q}M M^\top \mathbf{K}^\top = \mathbf{Q}\mathbf{K}^\top\), attention scores are mathematically unchanged, but each entry of \(\mathbf{Q}M\) is a random \(\pm\)-signed mixture of a full row of \(\mathbf{Q}\), so extreme values are distributed across all \(d\) coordinates before quantization. Concentration reduces the maximum magnitude relative to the RMS, thereby reducing quantization error.

The lower-cost approach — following QuIP [6] and QuIP# [7], which the paper cites explicitly — is \(M = \tfrac{1}{\sqrt d} H D\) with \(H\) a \(d \times d\) Hadamard matrix (\(\pm 1\) entries, orthogonal up to scale) and \(D = \operatorname{diag}(\pm 1)\) random. Then \(xM\) costs \(O(d \log d)\) via the fast Walsh–Hadamard transform rather than \(O(d^2)\), and in practice it is fused into the rotary-embedding kernel of the preceding layer. Because both operations are memory-bound, fusion can add little incremental latency, but the transform still performs additional arithmetic [1]. \(\mathbf{V}\) needs no transform (there is no analogous identity through which a transform of \(\mathbf{V}\) would cancel, and empirical measurements show fewer problematic outliers in V).

Measured effect (paper §4.3: RMSE with \(\mathbf{Q},\mathbf{K},\mathbf{V}\sim \mathcal{N}(0,1)\) plus \(\mathcal{N}(0,100)\) outliers at rate \(10^{-3}\)): FP16 baseline RMSE \(3.2\times 10^{-4}\) (FA2/FA3 FP16: \(1.9\times10^{-4}\), \(1.7\times\) better than the standard implementation because it uses FP32 intermediates); per-tensor-scaled FP8 baseline \(2.4\times 10^{-2}\); FA3 FP8 with block quantization + incoherent processing \(9.1\times 10^{-3}\) — the reported 2.6\(\times\) error reduction. The paper’s ablation shows the two techniques are complementary in different regimes: with these outliers, removing incoherent processing returns the error to \(2.4\times10^{-2}\) (indicating that incoherent processing provides most of the improvement in this regime), while block quantization provides a smaller improvement without requiring the incoherent transform.

6.5 FP8 kernel schedule and performance shape

The FP8 forward retains the FA3 structure (producer/consumer, pingpong, 2-stage overlap) with the following changes: the \(k\)-dimension per WGMMA doubles to 32; tiles grow since FP8 operands are half the bytes (current C++ tile_size.h FP8 configs: \(192\times160\) for \(d\le64\), \(128\times224\) for \(d\le128\), \(128\times128\) for \(d{=}256\)); the V-transpose stage joins the pipeline; and the P-conversion becomes FP32\(\to\)FP8 with the register permutation of Section 6.2. The paper reports up to 1.3 PFLOP/s at \(d{=}256\), \(N{=}16\)k and ∼1.0 PFLOP/s at \(d{=}128\) — i.e. FP8 achieves a lower fraction of its doubled peak (∼67% at best) than BF16 does (∼85%): the exponentials and FP32 softmax pipeline do not accelerate with FP8. Under Amdahl’s law, this unchanged work limits total utilization, and the in-kernel transpose adds shared-memory pressure. Against NVIDIA’s FP8 attention kernels the outcome depends on head dimension: the paper’s measurements put FA3 ahead at \(d{=}64\) and \(d{=}256\), while cuDNN is competitive at \(d{=}128\) [1] (Section 12.3). The scheduling considerations of Chapter 9 apply to both.

One configuration for validation-quality FP8 end-to-end serving uses e4m3 throughout, per-block descales from a calibration pass or dynamic quantization, Hadamard fused with RoPE for Q/K, and BF16 output (out_torch_dtype = bfloat16 — also what the DSL interface does for FP8 inputs on SM100).

7 The FA3 Backward Pass: Algorithm

The backward pass carries the larger share of the training cost: it must recompute the \(N \times N\) probabilities (never stored), run five GEMMs per tile pair instead of two, and resolve a genuine cross-block data dependency (dQ). This chapter derives the math and the FA3-specific execution plan, anchoring each design decision with short excerpts from the CuTe DSL sources (flash_bwd_sm90.py, flash_bwd_preprocess.py, flash_bwd_postprocess.py, and the supporting modules block_info.py, named_barrier.py, copy_utils.py, barrier.py). Chapter 8 then walks the main kernel line by line; the excerpts here are deliberately the pieces that chapter does not quote.

7.1 Gradient math

Given upstream gradient \(\mathbf{dO}\in \mathbb{R}^{N\times d_v}\) and the forward quantities \(\mathbf{S}= \alpha\mathbf{Q}\mathbf{K}^\top\), \(\mathbf{P}= \operatorname{softmax}(\mathbf{S})\), \(\mathbf{O}= \mathbf{P}\mathbf{V}\), the chain rule gives:

Equations (7.1), (7.2), (7.4), (7.5) are the ordinary product-rule gradients of the two GEMMs \(\mathbf{S}= \alpha\mathbf{Q}\mathbf{K}^\top\) and \(\mathbf{O}= \mathbf{P}\mathbf{V}\). The substantive step is (7.3), the softmax Jacobian. The derivation below shows how \(D\) makes the backward tileable at all.

7.1.1 The softmax Jacobian, step by step

Softmax acts row-wise, so fix a row \(i\) and write \(p = \mathbf{P}_{i,:}\), \(s = \mathbf{S}_{i,:}\), \(dp = \mathbf{dP}_{i,:}\) (all length-\(N\) vectors). With \(p_j = e^{s_j} / \sum_k e^{s_k}\), differentiate \(\log p_j = s_j - \log\sum_k e^{s_k}\):

Contract with the upstream gradient \(dp\):

The full Jacobian is a dense \(N \times N\) matrix per row (through the rank-1 term \(p\,p^\top\)), but its contraction collapses to one scalar \(D_i = \langle p, dp \rangle\) per row. That scalar is the entire coupling between different columns \(j\) of row \(i\): once \(D_i\) is known, \(dS_{ij}\) is a pointwise function of \(P_{ij}\) and \(dP_{ij}\). This is exactly what a tiled kernel needs — a KV block can compute its \(B_M \times B_N\) slab of \(\mathbf{dS}\) without seeing any other block’s columns, provided \(D_i\) arrives precomputed.

7.1.2 Why \(D\) can be computed without \(\mathbf{P}\)

Naively \(D_i = \sum_j P_{ij}\, dP_{ij}\) requires the full rows of \(\mathbf{P}\) and \(\mathbf{dP}\) — the matrices the backward does not materialize. But substituting \(\mathbf{dP}= \mathbf{dO}\mathbf{V}^\top\) and \(\mathbf{O}= \mathbf{P}\mathbf{V}\):

\(D\) is computable from two \(N \times d_v\) tensors already sitting in HBM, in a low-cost \(O(N d_v)\) pass — FA3 dedicates the preprocessing kernel (Section 7.3) to it.

7.1.3 Why LSE alone suffices to recompute \(\mathbf{P}\)

The forward saved one float per row: \(\operatorname{LSE}_i = \log \sum_j e^{S_{ij}} = m_i + \log \ell_i\) (Chapter 4). The recomputation is then a single subtraction in the exponent:

Two properties make this numerically safe and computationally inexpensive:

• Numerical safety needs no separate row max: the exponent is \(S_{ij} - \operatorname{LSE}_i = \log P_{ij} \le 0\), so exp2 can never overflow (underflow to \(0\) is the mathematically correct limit). The entire online-softmax machinery — running max, rescaling, two accumulators — disappears from the backward. Keeping \(m_i\) and \(\ell_i\) separately (as FlashAttention-1 did) is strictly redundant.

• The preprocessing kernel stores \(\operatorname{LSE}_i \log_2 e\) (LSElog2), so the mainloop’s per-element work in (7.9) is one FMA plus one EX2.

7.1.4 Gradients through LSE itself

FA3’s DSL backward also supports losses that read \(\operatorname{LSE}\) directly (e.g. distillation objectives). No new kernel is needed, because \(\partial \operatorname{LSE}_i / \partial S_{ij} = e^{S_{ij}}/\sum_k e^{S_{ik}} = P_{ij}\): the extra gradient term folds into \(D\). The derivation ships as a comment at the top of the preprocessing kernel:

Listing 7.1: flash_bwd_preprocess.py lines 5–13: the \(d\operatorname{LSE}\) folding identity.

# Computes D_i = (dO_i * O_i).sum(dim=-1), optionally adjusted for LSE gradient:
#   D'_i = D_i - dLSE_i
# This works because in the backward pass:
#   dS_ij = P_ij * (dP_ij - D_i)                     [standard]
# When LSE is differentiable, d(loss)/d(S_ij) gets an extra term dLSE_i * P_ij
# (since d(LSE_i)/d(S_ij) = P_ij), giving:
#   dS_ij = P_ij * (dP_ij - D_i) + dLSE_i * P_ij
#         = P_ij * (dP_ij - (D_i - dLSE_i))
# The main backward kernel is unchanged; replace D with D' = D - dLSE here.

7.1.5 Where the scale \(\alpha\) is applied

The kernels compute with the unscaled scores where possible and defer \(\alpha = 1/\sqrt{d}\) to the cheapest place: the \(\mathbf{P}\) recomputation folds \(\alpha \log_2 e\) into its FMA constant; \(\mathbf{dV}\) needs no scale at all; \(\mathbf{dK}\) is multiplied by \(\alpha\) in registers just before its epilogue (acc_dK *= softmax_scale); and \(\mathbf{dQ}\)’s scale is deferred all the way to the postprocessing kernel, which multiplies while converting dQaccum from FP32 (Section 7.3).

7.2 Tiling and the dQ problem

FA3 (like FA2 [3]) parallelizes the backward over KV blocks: one thread block owns \(\mathbf{K}_j, \mathbf{V}_j \in \mathbb{R}^{B_N\times d}\) for one (head, batch), loops over all relevant \(\mathbf{Q}_i/\mathbf{dO}_i\) tiles (\(i\) ascending; for causal, starting at the diagonal), and accumulates \(\mathbf{dK}_j, \mathbf{dV}_j\) in registers across the loop — writing them once at the end. Per iteration it computes, in tile form:

7.2.1 Why KV-parallel and not Q-parallel

The forward parallelizes over Q blocks because its only output, \(\mathbf{O}_i\), is private to a Q block. The backward has three outputs with opposite affinities: \(\mathbf{dQ}_i\) sums over \(j\) (KV blocks), while \(\mathbf{dK}_j\) and \(\mathbf{dV}_j\) sum over \(i\) (Q blocks). Whichever dimension the grid parallelizes over, the outputs indexed by the other dimension become cross-block accumulations:

KV-parallel is preferable: two of the three outputs stay in registers and only one needs cross-block treatment. (It also lets G3 and G5 — the two accumulating GEMMs, which under WGMMA can accumulate in place — target the register-resident tensors.)

7.2.2 Causal masking: trapezoidal iteration ranges

Under causal masking, KV block \(j\) only interacts with Q rows \(m \ge j B_N\) (for equal Q/KV lengths), so its inner loop starts at \(m_{\min} = \lfloor j B_N / B_M \rfloor\) and runs to the end of the sequence — each block traverses a trapezoid of the lower-triangular score matrix. The bound computation, shared with local (sliding-window) attention, sits in block_info.py; note the \(\mathrm{seqlen}_q - \mathrm{seqlen}_k\) shift that generalizes to unequal lengths (the causal diagonal is bottom-aligned, Chapter 9):

Listing 7.2: block_info.py lines 57–71 (trimmed): inner-loop bounds for one KV block.

def get_m_block_min_max(self, seqlen_info, n_block):
    m_block_max = cute.ceil_div(seqlen_info.seqlen_q, self.tile_m)
    m_block_min = 0
    if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)):
        n_idx_min = n_block * self.tile_n
        m_idx = n_idx_min + seqlen_info.seqlen_q - seqlen_info.seqlen_k
        m_idx_right = m_idx if const_expr(self.is_causal) else m_idx - self.window_size_right
        m_block_min = max(m_block_min, m_idx_right // self.tile_m)
    if const_expr(self.is_local and self.window_size_left is not None):
        n_idx_max = (n_block + 1) * self.tile_n
        m_idx_left = n_idx_max + seqlen_info.seqlen_q - seqlen_info.seqlen_k + self.window_size_left
        m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m))
    return m_block_min, m_block_max

The trapezoids are unequal: block \(j = 0\) sees (almost) every query, block \(j = n_{\max}\) sees only the last few. This is the mirror image of the forward’s causal imbalance, and the same remedy applies — the LPT (longest-processing-time-first) scheduler launches expensive blocks first (Chapter 9); its fixed traversal order is also what the deterministic mode’s semaphores assume (Section 7.6).

7.2.3 Three ways to resolve the dQ race

\(\mathbf{dK}_j\) and \(\mathbf{dV}_j\) accumulate privately within the block and therefore require no inter-block coordination. In contrast, every KV block \(j\) produces a contribution to every \(\mathbf{dQ}_i\) (within the causal mask), so blocks race on \(\mathbf{dQ}\). The alternatives:

1. Also parallelize a second kernel over Q blocks, recomputing the intermediates again. This doubles the recomputation cost and is not used by the FA2 or FA3 implementations.

2. Have each block atomically add its \(\mathbf{dQ}_i\) contribution to a global buffer. FP16/BF16 atomics would introduce unacceptable precision loss and nondeterministic rounding, so the buffer dQaccum is FP32, sized \(N_{\text{rounded}} \times d\) per head, zero-initialized by the preprocessing kernel. This is FA3’s choice. The “atomic” is coarse-grained: the paper describes atomicAdd; both the FA3 C++ code and the CuTe DSL implement it as a TMA bulk reduce-add (cp.reduce.async.bulk with .add.f32: one instruction adds a whole SMEM staging buffer into dQaccum in global memory, element-wise, with the L2 doing the reduction) — far fewer memory transactions than per-element red.global.add.f32, and no correctness difference. Section 7.6 shows the mechanism.

3. A deterministic variant: same buffer, but block \(j\)’s reduce-add for row-block \(i\) waits on a semaphore until block \(j-1\) (in a fixed order) has completed its add for \(i\) — FP32 additions then happen in a fixed order, at the cost of serialization (Section 7.6; Caveat 13.1).

Because dQaccum is FP32 and the model wants BF16, a small postprocessing kernel reads dQaccum, multiplies by \(\alpha\), and writes dQ in the output dtype (Section 7.3).

GQA: with \(h_q\) query heads per KV head, \(\mathbf{dK}/\mathbf{dV}\) also become cross-block accumulations (the \(h_q\) query heads all contribute). FA3’s DSL backward then switches \(\mathbf{dK}/\mathbf{dV}\) to the same FP32-buffer + bulk-reduce-add treatment (dKaccum/dVaccum), with their own semaphores in deterministic mode.

7.3 The preprocessing and postprocessing kernels

The mainloop is bracketed by two small grid-over-Q-row-blocks kernels that own the FP32 scratch tensors.

7.3.1 Preprocess: D, LSElog2, and zeroing dQaccum

flash_bwd_preprocess.py runs before the mainloop and:

1. loads \(\mathbf{O}_i, \mathbf{dO}_i\) tiles, computes \(D_i = \operatorname{rowsum}(\mathbf{dO}_i \odot \mathbf{O}_i)\) in FP32 (thread-local products, then a width-8 intra-warp shuffle reduction), and stores \(D\) (dPsum) padded to the block size;

2. converts \(\operatorname{LSE}\) to base-2 (LSElog2 \(= \operatorname{LSE}\cdot \log_2 e\)), stored padded — both statistics are laid out per-block-contiguously so the mainloop can fetch them with 1-D TMA (cp.async.bulk);

3. zeroes dQaccum (so the mainloop can reduce-add unconditionally);

4. optionally folds the \(d\operatorname{LSE}\) term of Listing 7.1 into \(D\), i.e. stores \(D_i' = D_i - d\operatorname{LSE}_i\).

The kernel body is quoted in Listing 8.8 (Chapter 8). Each row of \(D\) costs one multiply–add per element of \(\mathbf{dO}\), so the whole kernel is a bandwidth-bound \(O(N d_v)\) pass; it launches with PDL (programmatic dependent launch) so it overlaps the tail of the previous kernel, and the mainloop overlaps it in turn (griddepcontrol_wait before its first LSE load).

This preprocessing requires additional workspace: dQaccum holds \(N_{\text{rounded}} \times d\) FP32 values per (head, batch) — twice the bytes of the BF16 \(\mathbf{dQ}\) it will become — plus \(2 N_{\text{rounded}}\) floats for LSElog2 and dPsum. All are workspace tensors allocated by interface.py, invisible to the caller.

7.3.2 Postprocess: dQaccum \(\to\) dQ

After the mainloop, flash_bwd_postprocess.py converts. One thread block handles one \(B_M \times d\) tile of \(\mathbf{dQ}\): it bulk-loads the flat FP32 tile (contiguous, because the mainloop stored whole staged tiles), scales by \(\alpha\) while converting to the output dtype, and bounces through SMEM so the global write of \(\mathbf{dQ}\) is coalesced in the \((s, d)\) layout:

Listing 7.3: flash_bwd_postprocess.py lines 492–587 (trimmed): the SM90 path of the convert kernel.

# Step 1: load dQaccum gmem -> smem (cp.async, flat 1-D tile)
cute.copy(g2s_tiled_copy_dQaccum, tdQgdQaccum, tdQsdQaccumg2s)
cute.arch.cp_async_commit_group(); cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
# Step 2: smem -> registers, laid out like the mma accumulator
cute.autovec_copy(tdQsdQaccum, tdQrdQaccum)
# Convert from fp32 to fp16/bf16, folding in the softmax scale
rdQ = cute.make_fragment_like(acc, self.dtype)
rdQ.store((acc.load() * scale).to(self.dtype))
# Step 3: registers -> smem (same buffer recast to the narrow dtype)
cute.copy(thr_copy_r2s_dQ, taccdQrdQ, taccdQsdQ)
# Step 4: smem -> registers in gmem-coalesced order
cute.arch.barrier()
cute.autovec_copy(tdQsdQ, tdQrdQ)
# Step 5: registers -> gmem, predicated on seqlen and head_dim
for rest_m in cutlass.range(cute.size(tdQrdQ.shape[1]), unroll_full=True):
    if tdQcdQ[0, rest_m, 0][0] < seqlen_q - m_block * self.tile_m:
        cute.copy(gmem_tiled_copy_dQ, tdQrdQ[None, rest_m, None],
                  tdQgdQ[None, rest_m, None], pred=tdQpdQ[None, rest_m, None])

The register layout in step 2 is shaped by the same tiled-MMA partitioning the mainloop used to produce acc_dQ (the kernel takes AtomLayoutMdQ and dQ_swapAB as parameters for exactly this reason): the flat SMEM tile deinterleaves back into \((s, d)\) coordinates only if read with the layout it was written with. In the GQA path the same kernel family also converts dKaccum/dVaccum to \(\mathbf{dK}/\mathbf{dV}\), including the \(\alpha\) scaling for \(\mathbf{dK}\).

7.4 Warp specialization in the backward

The backward’s resource profile differs from the forward’s: five GEMMs, two pointwise stages, three tensors streamed per iteration (\(\mathbf{Q}_i\), \(\mathbf{dO}_i\), \(\mathrm{LSE}_i\)/\(D_i\)), one large output stream (dQaccum), and much higher SMEM pressure. FA3 organizes a block of \(3\times 128\) threads (configurable to \(4\times 128\)) as:

Producer WG (128 threads, 24 regs). Warp 0: TMA-loads K, V (once per work tile), then \(\mathbf{Q}_i\), \(\mathbf{dO}_i\), \(\mathrm{LSE}_i\), \(D_i\) per iteration. Warp 1: dQaccum store agent — waits on named barriers, issues TMA bulk reduce-adds.

Consumer WG0 (128 threads, 240/256 regs). One half of each \(B_M{\times}B_N\) S/dP tile (rows; columns under SdP_swapAB); GEMMs G1–G5 plus both pointwise stages.

Consumer WG1 (128 threads, 240/224 regs). The other half of each tile; same GEMM sequence, synchronized with WG0 at the PdS barrier.

The diagram shows the common two-consumer configuration. With a 512-thread CTA, a third 128-thread consumer warpgroup participates in the same protocols; the implementation sizes the relevant counts from num_mma_threads rather than assuming exactly 256 consumer threads.

Novelties versus the forward (the paper’s Algorithm 3 explicitly adds a “dQ writer” role beyond producer and consumer, with the accumulation step worded as “using a semaphore, atomically add \(\mathbf{dQ}_i^{(\text{local})}\) to \(\mathbf{dQ}_i\) in global memory” [1]):

• The producer warpgroup hosts two independent agents. Warp 0 produces loads. Warp 1 runs the dQaccum store loop: consumer warpgroups deposit their \(\mathbf{dQ}_i\) tile into an SMEM staging buffer (sdQaccum) and signal dQFull; warp 1 issues the TMA reduce-add to global dQaccum and signals dQEmpty. Stores thus overlap compute, and only one warp ever touches the TMA store path (deterministic issue order — essential for the deterministic mode).

• K and V are loaded once per work tile (they are the resident operands), via TMA with the barrier of the first Q/dO stage carrying their transaction bytes (extra_tx_count), so no separate pipeline is spent on them.

• Q and dO stream through 2-stage pipelines, each stage’s mbarrier also covering the corresponding 1-D TMA loads of \(\mathrm{LSE}_i\) and \(D_i\) (transaction counts add: \(B_M{\cdot}d{\cdot}2 + B_M{\cdot}4\) bytes).

• Statistics live in SMEM, not registers. LSE/D per row are loaded to SMEM, then read into registers with an optional shuffle-distributed scheme (shuffle_LSE): with SdP_swapAB each thread would need \(B_M/4\) statistics; instead 8 threads share them and shfl.sync fetches on demand — a pure register-pressure optimization, enabled for \(d \le 64\).

7.4.1 Named barriers of the backward

The synchronization vocabulary is spelled out in named_barrier.py — each enum value is a hardware named-barrier ID (bar.sync/bar.arrive with an immediate operand, Chapter 2):

Listing 7.4: named_barrier.py lines 28–39: barrier IDs used by the SM90 backward.

class NamedBarrierBwd(enum.IntEnum):
    Epilogue = enum.auto()  # starts from 1: barrier 0 is reserved for sync_threads()
    WarpSchedulerWG1 = enum.auto()
    WarpSchedulerWG2 = enum.auto()
    WarpSchedulerWG3 = enum.auto()
    PdS = enum.auto()
    dQFullWG0 = enum.auto()
    dQFullWG1 = enum.auto()
    dQFullWG2 = enum.auto()
    dQEmptyWG0 = enum.auto()
    dQEmptyWG1 = enum.auto()
    dQEmptyWG2 = enum.auto()

PdS synchronizes all participating consumer/MMA threads (num_mma_threads, commonly 256 and optionally 384) around the shared \(\mathbf{P}/\mathbf{dS}\) SMEM tiles. It serves both a data-ready rendezvous—all contributing register-to-SMEM stores must finish before a transposed WGMMA reads the full tile—and, in a single-stage configuration, a buffer-reuse rendezvous before the next tile overwrites the storage. The preceding fence_view_async_shared() is still required because the named barrier does not itself order generic stores with WGMMA’s async-proxy reads.

Each dQFullWG\(k\)/dQEmptyWG\(k\) pair implements a one-slot mailbox between consumer warpgroup \(k\) (128 threads) and producer warp 1 (32 threads) around one chunk of sdQaccum. Each named-barrier generation therefore uses a count of \(128+32=160\): the count is the total number of thread arrivals from the two branches, not an encoded list of participants. Full means the consumer has filled the chunk; Empty means the TMA store agent has finished reading it and the consumer may overwrite it. This resembles the full/empty ownership protocol of a one-stage mbarrier pipeline, but a named barrier does not track TMA transactions. The store warp must execute cp_async_bulk_wait_group(..., read=True) before arriving at Empty. The WarpSchedulerWG* IDs implement the same pingpong protocol as the forward’s (Section 4.4); the backward keeps them for the S/dP GEMMs but profits less from pingpong than the forward, since the five-GEMM chain already keeps the Tensor Cores busy.

7.4.2 SMEM budget

Why the backward’s tiles are small (Table 7.3) is best seen by adding up shared memory. The backward must simultaneously hold the resident \(\mathbf{K}, \mathbf{V}\), the pipelined \(\mathbf{Q}, \mathbf{dO}\) stages, the staged \(\mathbf{dS}\) (plus \(\mathbf{P}\) in configurations where G3 cannot read it from registers), the FP32 sdQaccum staging buffer, and the LSE/\(D\) stages:

Table 7.1: SMEM accounting (KB) for two BF16 backward configurations, from the layouts in _setup_attributes and _get_shared_storage_cls. Both use 2-stage Q/dO/PdS pipelines and skip sP because G3 consumes \(\mathbf{P}\) from registers (mma_dkv_is_rs).

The \(d{=}128\) non-causal configuration (\(B_M = 80\)) sits at ∼226 KB of the 228 KB limit, within 2 KB of the ceiling (\(B_M = 80\) also forces dQ_swapAB, since 80 is not a multiple of the WGMMA atom’s 64 rows); the causal tuning uses \(B_M = 64\) (∼193 KB). That, not FLOPs, is why \(B_M\) is 64–128 in the backward while the forward enjoys 128–192, and why the backward is the more severely register- and SMEM-limited of the two passes.

7.5 Scheduling the five GEMMs

Within a consumer warpgroup, one iteration issues (Listing 8.6 shows the code; simplified here):

wait Q\(_i\); G1 \(\mathbf{S}= \mathbf{Q}_i\mathbf{K}^\top\) (async) \(\cdot\) load LSE row stats \(\cdot\) wait dO\(_i\); G2 \(\mathbf{dP}= \mathbf{dO}_i\mathbf{V}^\top\) (async, wait_group(1) \(\Rightarrow\) G1 done) \(\cdot\) P1 \(\mathbf{P}= 2^{\ldots}\) (overlaps G2) \(\cdot\) convert \(\mathbf{P}\to\)BF16 \(\cdot\) wait_group(0) (\(\Rightarrow\) G2 done) \(\cdot\) P2 \(\mathbf{dS}= \mathbf{P}\odot(\mathbf{dP}- D)\) \(\cdot\) convert, store to sdS \(\cdot\) fence + named barrier (PdS) \(\cdot\) G3 \(\mathbf{dV}\mathrel{+}= \mathbf{P}^\top \mathbf{dO}_i\) (RS, async) \(\cdot\) G4 \(\mathbf{dQ}= \mathbf{dS}\,\mathbf{K}\) (async; wait_group(1) \(\Rightarrow\) G3 done, release dO) \(\cdot\) G5 \(\mathbf{dK}\mathrel{+}= \mathbf{dS}^\top\mathbf{Q}_i\) (async; wait_group(1) \(\Rightarrow\) G4 done) \(\cdot\) stage \(\mathbf{dQ}\) to SMEM for warp 1 \(\cdot\) wait_group(0) (\(\Rightarrow\) G5 done, release Q).

The exponentials (P1) hide under G2, and the \(\mathbf{dS}\) pointwise (P2) under G3’s issue window; G3/G4/G5 chain with wait_group(1) so two GEMMs can remain pending at parts of the schedule. This is a different overlap strategy from the forward’s: the forward has two GEMMs separated by a softmax and obtains overlap between iterations (intra-warpgroup 2-stage pipelining) and between warpgroups (pingpong). The backward’s five-GEMM chain provides enough intra-iteration instruction-level asynchrony that neither additional scheduling technique is beneficial: the WGMMA queue is kept two-deep within a single iteration, the two pointwise stages slot into the gaps, and the two consumer warpgroups split each tile spatially (half the rows or columns each) rather than alternating iterations. What remains asynchronous across agents is the memory system: TMA loads (warp 0), TMA reduce-add stores (warp 1), and WGMMA all run concurrently with the pointwise math.

7.5.1 Operand layouts: who is SS, who is RS

WGMMA reads operand B from SMEM always, and operand A from SMEM (SS) or registers (RS); 16-bit SMEM operands can be consumed either K-major or MN-major, which is what makes the transposed reads below legal (Chapter 2). For the canonical \(d \le 128\) tunings (SdP_swapAB, no dKV_swapAB):

Table 7.2: The five GEMMs’ operand sourcing in the \(d \le 128\) configurations. “K-major” = contiguous along the contraction axis. Operands are listed in their logical roles; an active swapAB flag (Section 7.5.2) exchanges which one the hardware treats as A and which as B.

Row-wise, the table shows that sQ, sdO, and sK each feed two GEMMs in different orientations (sQ: K-major in G1, MN-major in G5; sdO: K-major in G2, MN-major in G3; sK: K-major in G1, MN-major in G4) — the dual-orientation constraint that forces compromise swizzles (Section 8.2). \(\mathbf{dS}\) is materialized in SMEM regardless, because G4 must read it K-major from a warpgroup-crossing tile; \(\mathbf{P}\) escapes SMEM whenever G3 can take it from registers. The construction of all four tiled MMAs makes these choices explicit:

Listing 7.5: flash_bwd_sm90.py lines 256–304 (trimmed): _get_tiled_mma.

def _get_tiled_mma(self):
    maybe_swap_mn = lambda shape, swap: (shape[1], shape[0], *shape[2:]) if swap else shape
    # S = Q @ K.T, dP = dO @ V.T   -- both operands K-major, from SMEM
    tiled_mma_SdP = sm90_utils_basic.make_trivial_tiled_mma(
        self.dtype, self.dtype,
        warpgroup.OperandMajorMode.K, warpgroup.OperandMajorMode.K, Float32,
        atom_layout_mnk=maybe_swap_mn(atom_layout_SdP, self.SdP_swapAB),
        tiler_mn=(64, ...))
    # dV = P.T @ dO, dK = dS.T @ Q -- A from registers if mma_dkv_is_rs,
    # else the SMEM tile read MN-major (transposed); B always MN-major
    tiled_mma_dK, tiled_mma_dV = [
        sm90_utils_basic.make_trivial_tiled_mma(
            self.dtype, self.dtype,
            warpgroup.OperandMajorMode.MN if not self.mma_dkv_is_rs
                else warpgroup.OperandMajorMode.K,
            warpgroup.OperandMajorMode.MN, Float32,
            atom_layout_mnk=maybe_swap_mn(atom_layout_dKV, self.dKV_swapAB),
            tiler_mn=(64, ...),
            a_source=warpgroup.OperandSource.RMEM if self.mma_dkv_is_rs
                else warpgroup.OperandSource.SMEM)
        for tiler_mn_d in (tiler_mn_dK, tiler_mn_dV)]
    # dQ = dS @ K -- dS K-major, K read transposed (MN-major)
    tiled_mma_dQ = sm90_utils_basic.make_trivial_tiled_mma(
        self.dtype, self.dtype,
        warpgroup.OperandMajorMode.K if not self.dQ_swapAB else warpgroup.OperandMajorMode.MN,
        warpgroup.OperandMajorMode.MN if not self.dQ_swapAB else warpgroup.OperandMajorMode.K,
        Float32,
        atom_layout_mnk=maybe_swap_mn(atom_layout_dQ, self.dQ_swapAB),
        tiler_mn=(64, ...))
    return tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ

The availability of mma_dkv_is_rs is constrained. The constructor (lines 107–112) requires all of: AtomLayoutMSdP == 1, AtomLayoutNdKV == num_wg_mma, SdP_swapAB, and dKV_swapAB off. Under SdP_swapAB the G1/G2 accumulators come out of WGMMA already transposed (\(\mathbf{S}^\top\), \(\mathbf{dP}^\top\) tiles), and if the warpgroup split of the dKV GEMMs matches that of the S/dP GEMMs, each warpgroup already holds exactly the \(\mathbf{P}^\top/\mathbf{dS}^\top\) fragment its G3/G5 need — so A can come straight from registers. All the \(d \le 128\) SM90 tunings satisfy this; the \(d \ge 192\) tunings (SdP_swapAB off, dKV_swapAB on) do not and therefore require an sP staging buffer.

7.5.2 swapAB: transposing the whole GEMM instead of the data

WGMMA fixes M = 64 per atom, and the accumulator is distributed so that each thread owns fragments of rows. When \(B_M < 128\) (e.g. 64 or 80, common in the backward, Table 7.3), assigning the M dimension of \(\mathbf{S}\) to the hardware M would waste atoms or warpgroups. The FA3 kernels instead compute \(\mathbf{S}^\top = \mathbf{K}\mathbf{Q}^\top\) (\(B_N \times B_M\)) by swapping the roles of A and B — SdP_swapAB. Everything downstream transposes accordingly (masks index \((n, m)\), row statistics become column broadcasts, \(\mathbf{dK}/\mathbf{dV}\) vs \(\mathbf{dQ}\) swap which GEMMs need transposed inputs). The DSL carries three independent flags (SdP_swapAB, dKV_swapAB, dQ_swapAB) plus AtomLayout* counts that choose how the 2 warpgroups split each GEMM’s M or N dimension; the per-head-dim tunings are in Table 7.3. In Listing 7.5, the construction is one lambda: maybe_swap_mn exchanges the atom-layout M and N counts, and the operand major modes swap with the operand roles.

7.6 dQ accumulation across thread blocks

The cross-block \(\mathbf{dQ}\) sum of Section 7.2 is executed by a three-step relay, none of whose steps involves a conventional atomic instruction:

1. the consumer warpgroup copies its FP32 acc_dQ fragment into the sdQaccum staging buffer (a plain autovec_copy; the buffer is written and read linearly, so it needs no swizzle) and arrives at dQFullWG\(k\);

2. producer warp 1 waits on dQFullWG\(k\), then a single elected lane issues the bulk reduce-add and a bulk_commit_group;

3. once cp_async_bulk_wait_group(..., read=True) confirms that the TMA operation has finished reading the staging buffer, warp 1 arrives at dQEmptyWG\(k\), releasing the buffer for the next iteration.

These named barriers synchronize only threads in the same CTA. Deterministic ordering among different CTAs that reduce into the same global dQaccum row is a separate job performed by the global semaphore. The instruction at the center is a ten-line PTX wrapper in copy_utils.py:

Listing 7.6: copy_utils.py lines 266–288 (trimmed): the TMA bulk reduce-add.

@dsl_user_op
def cpasync_reduce_bulk_add_f32(smem_ptr, gmem_ptr, store_bytes, *, loc=None, ip=None):
    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
    llvm.inline_asm(
        None,
        [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
        "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
        "l,r,r",
        has_side_effects=True, is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT)

One cp.reduce.async.bulk instruction adds the entire staged chunk (kilobytes) into global dQaccum element-wise; the additions are performed near the L2, and the SM’s involvement ends at issue time. Compared with per-element red.global.add.f32 from registers, this reduces the instruction count by thousands and routes all of a block’s \(\mathbf{dQ}\) traffic through one warp in program order.

That ordering is what the deterministic mode builds on. FP32 addition is not associative, so if two KV blocks’ reduce-adds to the same \(\mathbf{dQ}_i\) rows are applied in different orders on different runs, the results differ in the last bits. FA3’s fix is a per-(row-block, head, batch) semaphore in global memory: block \(j\) may only issue its reduce-add for row-block \(i\) after the semaphore reads \(j\) (blocks increment it when done, so values pass through \(0, 1, 2, \ldots\)). The primitives, from barrier.py, are a spin-wait on an acquire load and a release-ordered increment:

Listing 7.7: barrier.py lines 8–71 (trimmed): global-memory semaphores for deterministic dQ.

@dsl_user_op
def ld_acquire(lock_ptr, *, loc=None, ip=None) -> cutlass.Int32:
    ...  # "ld.global.acquire.gpu.b32 $0, [$1];"

@dsl_user_op
def red_release(lock_ptr, val, *, loc=None, ip=None) -> None:
    ...  # "red.release.gpu.global.add.s32 [$0], $1;"

@cute.jit
def wait_eq(lock_ptr, thread_idx, flag_offset, val):
    flag_ptr = lock_ptr + flag_offset
    if thread_idx == 0:
        read_val = Int32(0)
        while read_val != val:
            read_val = ld_acquire(flag_ptr)

@cute.jit
def arrive_inc(lock_ptr, thread_idx, flag_offset, val):
    flag_ptr = lock_ptr + flag_offset
    if thread_idx == 0:
        red_release(flag_ptr, val)

The acquire/release pair is essential: ld.global.acquire guarantees that once block \(j\) observes the semaphore value, it also observes block \(j{-}1\)’s completed reduce-add; red.release publishes the increment only after the TMA write is visible (cp_async_bulk_wait_group(0) precedes it, Listing 8.7). The expected value must match the traversal order: under the causal LPT schedule blocks visit \(n\) in reverse, so the lock value is computed as \(n_{\max} - 1 - n\) rather than \(n\). Serialization is per row-block only — different \(\mathbf{dQ}_i\) row-blocks still accumulate in parallel — but the dependency chains cost measurable throughput (Caveat 13.1). Section 8.7 quotes the full store-agent loop, including the deadlock risk associated with local attention (blocks that skip an m_block must still bump its semaphore).

7.7 Backward tile configurations

Table 7.3 collects the per-head-dim tunings that the preceding sections referenced.

Table 7.3: Backward configurations on SM90 (CuTe DSL _tile_size_bwd_sm90, matching C++ flash_bwd_launch_template.h tunings). All: 1 producer + 2 consumer warpgroups, \(\mathbf{Q}\)/\(\mathbf{dO}\) 2-stage pipelines unless noted.

“RS dKV” marks the mma_dkv_is_rs configurations of Listing 7.5; dQ_single_wg means consumer WG0 alone computes G4 while WG1 skips it (rebalancing registers to \(256/224\)); “stats shuffled” is the shuffle_LSE/shuffle_dPsum register-pressure optimization (\(d \le 64\) only). The SMEM accounting behind the shrinking tiles is in Table 7.1.

7.8 Causal masking, varlen, and other mainloop features

Causal load imbalance. With KV-block ownership, causal attention makes low-\(j\) blocks expensive (they see almost all queries) and high-\(j\) blocks less expensive — the trapezoids of Listing 7.2. The mask itself is applied only in the iterations near the diagonal (as mask_fn with seqlen + causal components, again with the R2P bitmask optimization of Chapter 4, but indexed \((n,m)\) under SdP_swapAB), and the scheduler counteracts the imbalance with LPT ordering (Chapter 9).

Local (sliding-window) attention. Both loop bounds tighten (Listing 7.2): a right window bounds \(m_{\min}\) like causal does, a left window bounds \(m_{\max}\), so each KV block visits only the band of Q blocks inside the window. The deterministic mode must then account for blocks that never visit an m_block at all (Section 7.6).

Variable-length batches. Varlen selects the SingleTileVarlenScheduler and per-batch SeqlenInfo offsets (Chapter 9); dQaccum, LSElog2, and dPsum are indexed with padded per-batch offsets so each batch’s scratch stays tile-aligned, and with qhead_per_kvhead == 1 the \(\mathbf{dK}/\mathbf{dV}\) epilogue switches to ragged (varlen-aware) TMA stores.

Softcap and score modifiers. The DSL backward has no dedicated softcap code path: interface.py translates softcap into a pair of FlexAttention-style score-modifier callbacks (create_softcap_scoremod / _bwd), and the mainloop invokes score_mod on acc_S right after G1 and score_mod_bwd on acc_dP (given the pre-modification scores) before P2 — the generic hook mechanism of Chapter 5, compiled inline into the kernel.

7.9 Backward performance

On the paper’s H100 measurements the FA3 backward runs 1.45–1.96\(\times\) FlashAttention-2’s backward throughput, around 550–615 TFLOP/s at \(N=16\)K for \(d \in \{64, 128\}\) [1] — roughly 56–62% utilization on the \(2.5\times\)-FLOP accounting (\(10 N^2 d\) per head), versus the forward’s 85%. The structural reasons are the ones this chapter assembled: five dependent GEMMs with two pointwise stages between them, dual-orientation SMEM operands with compromise swizzles, the FP32 dQaccum round trip through L2, and register/SMEM ceilings that cap \(B_M\). Chapter 12 gives the full numbers and comparisons.

8 The FA3 Backward Pass: CuTe DSL Implementation

Reference file: flash_attn/cute/flash_bwd_sm90.py (class FlashAttentionBackwardSm90, ∼1930 lines), plus flash_bwd_preprocess.py and flash_bwd_postprocess.py. As before, listings are trimmed excerpts of the real files.

8.1 Configuration surface

The constructor exposes the full tuning space of Section 7.5:

Listing 8.1: flash_bwd_sm90.py lines 46–143 (trimmed): constructor.

class FlashAttentionBackwardSm90:
    arch = 90
    def __init__(self, dtype, head_dim, head_dim_v=None, qhead_per_kvhead=1,
                 is_causal=False, is_local=False, deterministic=False,
                 tile_m=64, tile_n=128,
                 Q_stage=2, dO_stage=2, PdS_stage=2,
                 SdP_swapAB=False, dKV_swapAB=False, dQ_swapAB=False,
                 AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1,
                 num_threads=384, V_in_regs=False, ..., dQ_single_wg=False):
        ...
        self.num_wg_mma = (self.num_threads // 128) - 1
        self.mma_dkv_is_rs = (AtomLayoutMSdP == 1 and AtomLayoutNdKV == self.num_wg_mma
                              and SdP_swapAB and not dKV_swapAB)
        # Keep LSE and dPsum in each thread, or split them across 8 threads
        # that share them and then shuffle ... reduces register
        # pressure when SdP_swapAB, where each thread needs statistics for
        # (kBlockM / 4) rows.
        self.shuffle_LSE = self.SdP_swapAB and self.tile_hdim <= 64
        self.shuffle_dPsum = self.SdP_swapAB and self.tile_hdim <= 64
        ...
        # dQ_single_wg: WG0 computes the full dQ GEMM, WG1 skips it.
        self.num_wg_dQ = 1 if dQ_single_wg else self.num_wg_mma

The register budgets (lines 428–446): with 2 MMA warpgroups, \((240, 240, 24)\), or \((256, 224, 24)\) under dQ_single_wg — WG0 carries the dQ accumulator so it gets the bigger share; the code asserts the sum \(\le 504\) (REG_LIMIT).

8.2 SMEM layouts supporting two access orientations

The layout construction is the central detail: buffers consumed by WGMMA in both orientations need their swizzle atom restricted to the size that both access patterns can decode:

Listing 8.2: flash_bwd_sm90.py lines 201–252 (trimmed): SMEM layouts for dual-orientation operands.

def _setup_attributes(self):
    # Accommodate both Q and Q^T (and dO and dO^T) in shared memory.
    # Q & dO are used in the SdP Mma and Q^T and dO^T are used in the dKV Mma.
    wg_d_dKV = self.num_wg_mma // self.AtomLayoutNdKV
    self.sQ_layout, self.sdO_layout = [
        sm90_utils.make_smem_layout(self.dtype, LayoutEnum.ROW_MAJOR, shape, stage,
                                    major_mode_size=mms)
        for shape, stage, mms in [
            ((self.tile_m, self.tile_hdim), self.Q_stage, self.tile_hdim // wg_d_dKV),
            ((self.tile_m, self.tile_hdimv), self.dO_stage, self.tile_hdim // wg_d_dKV),
        ]
    ]
    # Accomodate both K and K.T (dQ GEMM reads K^T)
    self.sK_layout = sm90_utils.make_smem_layout(..., (self.tile_n, self.tile_hdim),
        stage=None, major_mode_size=self.tile_hdim // (self.num_wg_dQ // self.AtomLayoutMdQ))
    # There's only V, no V.T, so layout is normal
    self.sV_layout = sm90_utils.make_smem_layout(..., (self.tile_n, self.tile_hdimv), None)
    # Accomodate both S and S.T
    self.sPdS_layout = sm90_utils.make_smem_layout(self.dtype, LayoutEnum.ROW_MAJOR,
        (self.tile_m, self.tile_n), stage=self.PdS_stage,
        major_mode_size=math.gcd(self.tile_n // wg_n_SdP, self.tile_n // wg_n_dKV))
    self.sdQaccum_layout = cute.make_layout(
        (self.tile_m * self.tile_hdim // self.num_wg_dQ, self.num_wg_dQ))

major_mode_size caps the contiguous run length the swizzle atom assumes; a buffer read K-major by one GEMM and MN-major by another must use the gcd of what each would like — a smaller swizzle than optimal, trading a few bank conflicts for correctness in both orientations. This is invisible in the forward (each buffer has one consumer orientation) and a recurring theme in the backward.

The shared-storage struct (lines 306–341) adds sLSE/sdPsum (rounded to 64 rows per stage) and the FP32 sdQaccum staging buffer.

8.3 Host-side: TMA atoms, schedulers, semaphores, PDL

Lines 451–517 create TMA atoms: G2S for \(\mathbf{Q}, \mathbf{K}, \mathbf{V}, \mathbf{dO}\); S2G for \(\mathbf{dK}, \mathbf{dV}\) (only when qhead_per_kvhead == 1 — the GQA path stores via bulk reduce-add instead); LSE and \(D\) are transferred by 1-D cp.async.bulk. The transaction-byte bookkeeping shows how the transfers share a barrier:

Listing 8.3: flash_bwd_sm90.py lines 692–707: Q and dO pipelines carry LSE/dPsum bytes.

pipeline_Q = pipeline.PipelineTmaAsync.create(
    barrier_storage=storage.mbar_ptr_Q.data_ptr(), num_stages=self.Q_stage,
    producer_group=pipeline_producer_group, consumer_group=pipeline_consumer_group,
    tx_count=self.tma_copy_bytes["Q"] + self.tma_copy_bytes["LSE"],
    defer_sync=True)
pipeline_dO = pipeline.PipelineTmaAsync.create(
    barrier_storage=storage.mbar_ptr_dO.data_ptr(), num_stages=self.dO_stage,
    ...,
    tx_count=self.tma_copy_bytes["dO"] + self.tma_copy_bytes["dPsum"],
    defer_sync=False)

Scheduler selection (lines 518–524): varlen \(\Rightarrow\) SingleTileVarlenScheduler; deterministic \(\Rightarrow\) SingleTileLPTBwdScheduler (fixed traversal order is what the dQ semaphores assume); otherwise SingleTileScheduler. The launch (line 621–627) uses use_pdl=True: programmatic dependent launch, paired with griddepcontrol_wait() in the producer before the first LSE load — the mainloop may start while the preprocess kernel is still writing.

8.4 Role dispatch: two producer warps, two consumer warpgroups

Listing 8.4: flash_bwd_sm90.py lines 762–851 (trimmed): four concurrent agents.

if warp_idx < 4:
    cute.arch.setmaxregister_decrease(self.num_producer_regs)   # 24
    if warp_idx == 0:
        self.load(mQ, mK, mV, mdO, mLSE, mdPsum, sQ, sK, sV, sdO, sLSE, sdPsum,
                  tma_atom_Q, ..., pipeline_Q, pipeline_dO, ...)
    if warp_idx == 1:
        self.dQaccum_store(mdQaccum, sdQaccum, block_info, TileSchedulerCls,
                           SeqlenInfoCls, ..., mdQ_semaphore)
else:
    tidx = cute.arch.thread_idx()[0] - 128
    if const_expr(self.num_wg_dQ == self.num_wg_mma):
        cute.arch.setmaxregister_increase(self.num_mma_regs_wg0)   # 240
        self.mma(*mma_args, is_dQ_wg=True)
    else:  # dQ_single_wg: WG0 computes dQ, WG1 skips it
        if warp_idx_in_mma < 4:
            cute.arch.setmaxregister_increase(self.num_mma_regs_wg0)  # 256
            self.mma(*mma_args, is_dQ_wg=True)
        else:
            cute.arch.setmaxregister_increase(self.num_mma_regs_wg1)  # 224
            self.mma(*mma_args, is_dQ_wg=False)

8.5 The producer load loop

For each work tile (= one KV block), K and V are loaded once, sharing the mbarriers of the first Q and dO stages via extra_tx_count; then \(\mathbf{Q}_i\)/LSE\(_i\) and \(\mathbf{dO}_i\)/\(D_i\) stream per iteration:

Listing 8.5: flash_bwd_sm90.py lines 955–993 (trimmed): producer steady state.

first_m_block = m_block_min
pipeline_Q.producer_acquire(producer_state_Q, extra_tx_count=self.tma_copy_bytes["K"])
load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q))
load_Q(first_m_block, producer_state=producer_state_Q)
# Wait for bwd preprocess to finish writing LSE and dPsum
cute.arch.griddepcontrol_wait()
load_LSE(first_m_block, producer_state=producer_state_Q)
pipeline_dO.producer_acquire(producer_state_dO_cur, extra_tx_count=self.tma_copy_bytes["V"])
load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_cur))
load_dO(first_m_block, producer_state=producer_state_dO_cur)
load_dPsum(first_m_block, producer_state=producer_state_dO_cur)
producer_state_Q.advance(); producer_state_dO.advance()

for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1):
    pipeline_Q.producer_acquire(producer_state_Q)
    load_Q(m_block, producer_state=producer_state_Q)
    load_LSE(m_block, producer_state=producer_state_Q)
    pipeline_dO.producer_acquire(producer_state_dO_cur)
    load_dO(m_block, producer_state=producer_state_dO_cur)
    load_dPsum(m_block, producer_state=producer_state_dO_cur)
    producer_state_Q.advance(); producer_state_dO.advance()