FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

Chapter 1: The Blackwell Hardware Model

To understand FA4's design decisions, you need a mental model of the Blackwell GPU's memory hierarchy and execution units. Every optimization in FA4 targets a specific hardware resource. Let's map them out.

Memory hierarchy: four levels

Data on a Blackwell GPU lives in a hierarchy of storage, each level trading capacity for speed:

Tensor cores: 5th generation

Blackwell's tensor cores process 128 × N tiles (N = 128 or 256), compared to Hopper's 64 × N tiles. The throughput is 8192 FLOPs/clock/SM for BF16 — exactly double Hopper's 4096. Each MMA instruction writes its output directly to TMEM asynchronously, meaning computation and other operations can overlap without blocking on register writeback.

The thread hierarchy

From finest to coarsest:

Key functional units and their throughput

2-CTA tensor core mode

Blackwell introduces a new MMA mode where a CTA pair within the same thread block cluster cooperatively executes a single MMA. One CTA initiates the operation; the peer CTA must be launched and active. The paired mode supports M = 128 or 256 by partitioning the A tile in dimension M and the B tile across the two CTAs in dimension N, so each CTA stages only half of operand B in its own shared memory while the hardware consumes the combined B tile. This reduces redundant shared memory capacity and bandwidth.

For FA4's backward pass, this 2-CTA mode is critical: it lets the kernel use M = 256 tiles (vs M = 128 in single-CTA mode), roughly halving the shared memory traffic for operand B and halving the number of global atomic reductions for dQ accumulation.

Chapter 2: Roofline Analysis — Finding the Bottleneck

Before writing a single line of kernel code, FA4's authors do something crucial: a roofline analysis. This means calculating, for each hardware resource, how many clock cycles it needs to process one tile of attention. The resource that takes the most cycles is the bottleneck — and all kernel design effort should go toward either speeding it up or hiding it behind something else.

Forward pass analysis

Let the tile shape along the sequence length of Q and K be M × N, and the head dimension be d. The forward pass performs two MMAs per iteration of the inner loop:

Each MMA requires 2MNd floating-point operations (two matmuls). The total MMA compute time is:

For shared memory traffic, the QK^T MMA reads both operands from SMEM (shared-shared), while PV reads P from tensor memory and V from SMEM (tensor-shared). At 128 bytes/cycle and 2 bytes per BF16 element:

The exponential unit computes exp() for M × N values in the softmax, at 16 ops/cycle:

The numbers for typical tile sizes

Backward pass is worse

The backward pass performs five MMAs per iteration (recomputing S, plus computing dP, dV, dQ, dK). The shared memory traffic becomes the dominant bottleneck:

Shared memory traffic exceeds MMA compute by 30% in 1-CTA mode. Even in 2-CTA mode (which halves some shared memory reads), it still exceeds MMA compute by ~5%. This is why FA4 uses the 2-CTA MMA mode for the backward pass — it is the only way to bring shared memory traffic close to MMA compute time.

Chapter 3: The Pingpong Pipeline

The core scheduling idea in FA4 is the same as FA3: pingpong between two warpgroups. While one warpgroup's tile runs tensor core operations (QK^T or PV), the other warpgroup computes softmax. They alternate, keeping both the tensor cores and the non-matmul units busy simultaneously.

But on Blackwell, this is harder to get right. The accumulators now live in TMEM (not registers), the tile sizes are 128×128 (not 64×128), and the imbalance between MMA speed and everything else is more extreme. FA4 redesigns the pipeline from scratch.

The two warpgroups

Each SM runs one CTA (threadblock) with two warpgroups of 128 threads each. At any given moment:

The superscripts H and L denote the "high" and "low" Q tiles — each warpgroup owns 128 rows of Q. Together they process 256 query rows per SM. The key constraint: the two softmax warpgroups must not overlap in their critical section (the exponential computation), because both use the same MUFU unit on the SM.

TMEM partitioning

This is where the design gets intricate. TMEM must hold the intermediate results that bridge the MMA and softmax phases. FA4 allocates TMEM as follows:

The softmax warpgroup's job

Each softmax warpgroup processes one row of the attention matrix at a time. On Blackwell's larger 128×128 tiles, each thread must hold an entire row of 128 elements in registers. The softmax procedure per warpgroup:

Register pressure

For BF16 inputs, each thread needs: 128 registers for the S row input, 64 registers for output staging, plus miscellaneous temporaries. To reduce register pressure, FA4 stages the output P in quarters: the first three quarters are stored once (triggering the corresponding MMA operations), and the last quarter is stored separately. This careful staging keeps register usage within the 256-register-per-thread limit.

Chapter 4: Software-Emulated Exponential

The MUFU (Multi-Function Unit) computes exp() at 16 operations per clock per SM. For a 128×128 tile, that's 128×128 / 16 = 1024 cycles — matching the MMA compute time exactly. This means even with perfect pingpong overlap, the exponential is right on the edge of being a bottleneck. Any small inefficiency (pipeline bubbles, synchronization) tips it over.

FA4's solution is radical: compute some of the exponentials on the FMA (fused multiply-add) units instead of MUFU. The FMA units are the same ones used for floating-point arithmetic — they can run in parallel with MUFU, effectively doubling the exponential throughput.

The mathematical decomposition

The key identity is:

where ⌊x⌋ is the integer part and x_frac = x - ⌊x⌋ ∈ [0, 1) is the fractional part. The integer part 2^⌊x⌋ can be computed without any floating-point arithmetic — it's just a bit manipulation on the IEEE 754 exponent field. Shift ⌊x⌋ into the exponent bits and you have your answer.

The fractional part 2^x_frac is approximated by a polynomial:

with p₀ = 1.0 and the remaining coefficients chosen to minimize relative approximation error. This polynomial is evaluated using Horner's method on the FMA units — each term is one fused multiply-add.

The algorithm step by step

Accuracy: is it good enough?

The authors benchmark polynomial approximations of different degrees against the hardware MUFU.EX2 instruction:

Partial emulation

Computing ALL exponentials via polynomial emulation would be wasteful — the additional registers for intermediate values and coefficients would increase register pressure and could cause spills. Instead, FA4 uses partial emulation: only 10–25% of the entries in each softmax row are emulated via polynomial, with the remaining entries computed via hardware MUFU.EX2. The exact fraction is tuned empirically based on the ratio of MMA and exponential throughput for a given tile configuration.

This distributes the exponential computation across both MUFU and FMA units, alleviating the bottleneck while keeping register pressure manageable.

python
# Conceptual software-emulated 2^x (simplified)
import struct

def fast_exp2(x):
    # Step 1: Clamp
    x = max(x, -127.0)

    # Step 2: Integer part via bit tricks
    x_floor = int(x)  # In GPU: add magic number + round-down

    # Step 3: Fractional part
    x_frac = x - x_floor   # x_frac in [0, 1)

    # Step 4: Degree-3 polynomial (Horner's method on FMA units)
    # Coefficients from Sollya minimizing relative error
    p3, p2, p1, p0 = 0.07944, 0.2274, 0.6931, 1.0
    frac_part = ((p3 * x_frac + p2) * x_frac + p1) * x_frac + p0

    # Step 5: Combine via IEEE 754 exponent manipulation
    # Shift x_floor into exponent field, multiply by frac_part
    int_part = 2.0 ** x_floor  # In GPU: bit shift on exponent
    return int_part * frac_part

Chapter 5: Conditional Softmax Rescaling

FlashAttention computes attention in blocks. Because we process K/V blocks one at a time, we maintain running statistics for the online softmax: a running maximum m_j and a running sum of exponentials l_j. When a new block has a larger maximum than the running max, we must rescale all previously accumulated output — multiplying O by e^{m_j-1 - m_j}.

The online softmax recap

When processing block j, FA computes:

The intermediate output is updated as:

That rescaling factor e^{m_j-1 - m_j} is a vector multiplication — it must be applied to every element of the accumulated O. On Blackwell, where non-matmul operations are the bottleneck, this rescaling is expensive.

FA4's insight: skip rescaling when it's not needed

FA4 makes two key observations:

Observation 1: Rescaling is only necessary when m_j > m_j-1, i.e., when the current block contains larger attention scores than all previous blocks. In practice, the running maximum stabilizes quickly — after the first few blocks, most iterations don't need rescaling at all.

Observation 2: We can tolerate some "slack" in the rescaling. If m_j - m_j-1 > τ (a threshold), we must rescale. But if the difference is small (m_j - m_j-1 ≤ τ), we can skip rescaling and continue using m_j-1. The statistics will be slightly off, but we fix everything at the end with a final normalization.

The threshold τ is typically set to log₂(256) = 8.0, corresponding to a rescaling factor of 256.0. As long as we keep track of the true statistics, the final output is exact:

At the very end:

Decoupling rescaling from the critical path

Even when rescaling IS needed, FA4 moves it off the critical path. Since accumulators live in TMEM (not registers), rescaling can be performed by a separate "correction" warpgroup that reads the accumulator from TMEM, multiplies by the rescaling factor, and writes it back — all while the main warpgroups continue with the next iteration's MMA and softmax. This decoupling is impossible on Hopper, where accumulators live in registers that are private to each thread.

To avoid warp divergence, FA4 rescales whenever any thread in the warpgroup needs rescaling. This is a conservative choice that trades a small amount of unnecessary rescaling for simpler control flow.

Chapter 6: The Backward Pass

The backward pass is the harder half. It performs five MMAs per iteration (vs two in forward), recomputing S on the fly, and must accumulate gradients dQ, dK, dV. On Blackwell, shared memory traffic is the dominant bottleneck — exceeding MMA compute by 30%.

The five MMAs

Given output gradient dO ∈ R^N×d, the backward pass computes:

where dsoftmax(dP) denotes the row-wise softmax gradient: dS = (diag(p) - pp^T)dp for p = softmax(s). This requires recomputing S = QK^T (since we didn't store it in the forward pass — that's the whole point of FlashAttention). So the five MMAs are:

Three of the five (S^T, dQ, dK) are shared-shared — both operands read from SMEM. This is why shared memory bandwidth dominates the backward pass.

1-CTA pipeline

In 1-CTA mode (M = 128), the pipeline follows the same pingpong principle as the forward pass, but with 5 MMAs + 2 elementwise operations (softmax and its gradient). The pipeline stages across prologue, main loop, and tail:

2-CTA mode: halving shared memory traffic

Even with the 1-CTA pipeline, shared memory is still the bottleneck. FA4 uses Blackwell's 2-CTA MMA mode to address this. With an MMA tile shape of M = 256 and N = K = 128:

• Each CTA loads and stages half of operand B in its own shared memory
• The hardware consumes the combined B tile from both CTAs during the MMA
• This roughly halves the shared memory traffic for operand B
• The dQ reduction dimension is N, which is naturally split across the CTA pair
• Each CTA writes only half the dQ tile, halving the number of global atomic reductions

The two CTAs communicate via Distributed Shared Memory (DSMEM) — a Blackwell feature allowing CTAs in the same cluster to access each other's SMEM. FA4 uses DSMEM to exchange half of the dS tile between the two CTAs, since each CTA needs the full dS for its MMA but computes only half of it.

Deterministic backward pass

The dQ gradient accumulates via atomic additions to global memory, which are nondeterministic (the order of atomic operations varies between runs). For reproducible training in reinforcement learning, FA4 provides a deterministic mode using semaphore locks: each CTA acquires a lock before writing to a common dQ tile, performs its reduction, and releases the lock. Combined with careful CTA swizzling (processing tiles in an order that minimizes stalls), the deterministic backward pass achieves up to 75% of the nondeterministic version's speed.

Chapter 7: CuTe-DSL — A New Way to Write Kernels

Every previous FlashAttention version was written in CUDA C++ using CUTLASS templates — a maze of deeply nested C++ template metaprogramming that took 55 seconds to compile a single kernel. FA4 takes a radically different approach: it is written entirely in CuTe-DSL, a domain-specific language embedded in Python.

What is CuTe-DSL?

CuTe-DSL is the Python frontend for NVIDIA's CuTe (Collective Unrolled Tensor Expression) framework. It provides:

Why Python instead of C++?

The CuTe-DSL programming model is isomorphic to CUTLASS C++ — every C++ CUTLASS construct has a direct CuTe-DSL equivalent. You lose nothing in expressiveness. But you gain enormously in productivity:

Escape hatch to PTX

CuTe-DSL provides direct access to PTX (GPU assembly) as an escape hatch. FA4 uses custom PTX sequences for operations not yet fully exposed in the CuTe-DSL API, demonstrating that the framework doesn't constrain developers to a limited subset of GPU capabilities. This is critical for a performance-critical kernel that needs every last drop of throughput.

Composable primitives

FA4's codebase is designed as a library of composable primitives rather than a monolithic kernel:

• Block-sparse patterns — plug in different masking strategies
• Variable sequence length — varlen support as a composable feature
• Work scheduling — LPT, SPT, and other schedulers as interchangeable modules
• Causal masking — orthogonal to the core attention logic

Developers have already built FlexAttention and block-sparse attention variants on top of FA4 without modifying the core framework. This modularity is a direct benefit of the CuTe-DSL approach — Python's flexibility makes it natural to compose orthogonal features.

python
# Conceptual CuTe-DSL kernel structure (simplified)
from cute_dsl import Tensor, Layout, TiledMMA, TMA

# Define tile shapes
BM, BN, BD = 128, 128, 128

# Define MMA operation (maps to WGMMA instruction)
tiled_mma = TiledMMA(
    instruction="wgmma_bf16_128x128x128",
    warpgroup_size=128,  # 4 warps
)

# Define TMA copy from HBM to SMEM
tma_load_k = TMA(
    src="gmem", dst="smem",
    shape=(BN, BD), dtype="bf16",
    swizzle=128,  # 128-byte aligned for optimal bandwidth
)

# Main loop body (one iteration)
def attention_step(Q_smem, K_smem, V_smem, O_tmem, m, l):
    # 1. QK^T via WGMMA (writes S to TMEM)
    S = tiled_mma(Q_smem, K_smem.T)   # → TMEM

    # 2. Online softmax (reads S from TMEM, writes P to TMEM)
    m_new = max(m, row_max(S))
    if m_new - m > tau:
        O_tmem *= exp(m - m_new)   # conditional rescale
    P = softmax_partial(S, m_new)  # MUFU + FMA emulation

    # 3. PV via WGMMA (accumulates into O in TMEM)
    O_tmem += tiled_mma(P, V_smem)  # → TMEM
    return O_tmem, m_new, l_new

Chapter 8: Benchmarks — Measuring the Payoff

FA4 is benchmarked on a B100 180GB SXM6 (1000W) GPU with CUDA 13.1, PyTorch 2.10.0, and CuTe-DSL 4.4.1. Baselines include cuDNN 9.13.0, cuDNN 9.19.1.2, Triton 3.6, and Gluon (a lower-level GPU language). FA3 does not run on Blackwell due to incompatible Hopper MMA instructions.

Benchmark settings

The setup: BF16 inputs, with/without causal mask, head dimensions 64, 128, and (192, 128) for the DeepSeek V3 configuration. Sequence lengths from 1K to 32K. Total tokens fixed at 32K (batch size adjusts accordingly). Hidden dimension 2048, giving either 32 heads (d=64), 16 heads (d=128), or 16 heads with GQA (d_q=192, d_kv=128).

Forward pass results

Peak throughput: 1613 TFLOPS, approximately 71% of the B200's theoretical 2.25 PFLOPS. The gains are larger for causal attention, which FA4's LPT scheduler handles particularly well.

For DeepSeek V3 configuration (192, 128)

With the GQA head dimensions used by DeepSeek V3 (192 query dims, 128 KV dims, 16 heads), FA4 achieves even larger gains over cuDNN:

Backward pass results

The backward pass shows larger gains over baselines, especially for causal attention:

Scheduling ablations

FA4's LPT (Longest-Processing-Time-first) scheduler provides a measurable 4–8% FLOPS gain for MHA and 7–14% for MQA 8 at head dimension 128, by ensuring that the longest-running worktiles are scheduled first, reducing tail-end imbalance across SMs.

For the deterministic backward pass, the SPT (Shortest-Processing-Time-first) scheduler with CTA swizzling achieves up to 950 TFLOPS for causal attention at 32K sequence length — 75% of the nondeterministic version's speed. This is fast enough for practical use in RL training where reproducibility is essential.

Chapter 9: Connections

FlashAttention-4 sits at the intersection of hardware-aware algorithm design, GPU kernel engineering, and ML systems research. Here are the key threads to follow.

The FlashAttention lineage

Related hardware-aware attention

Key takeaways

Open questions

• How will FA evolve for B300/GB300, which double MUFU throughput to 32 ops/clock? Will the bottleneck shift again?
• Can the polynomial exponential technique be applied to other non-matmul bottlenecks (e.g., layer norm, GELU)?
• Multi-GPU attention (sequence parallelism across NVLink) is not yet addressed — how does FA4's pipeline extend?
• Sparse attention patterns (top-k, block-sparse, sliding window) are supported via composable primitives, but are there scheduling opportunities specific to sparse patterns?