READ MORE

ArXiv: 2603.05451

🎯 Pitch

Blackwell's tensor cores are 2× faster than Hopper's, but its exponential units and shared memory bandwidth haven't budged—so attention kernels actually slowed down relative to peak compute. FlashAttention-4 redesigns the entire pipeline to bypass these stuck-on-Hopper bottlenecks, hitting 1.3× over cuDNN and 71% of B200's peak throughput by faking exponentials with polynomial math and halving backward-pass atomics.

---

1. Executive Summary

FlashAttention-4 introduces a co-design of attention algorithms and GPU kernel pipelining to address the shifting performance bottlenecks caused by asymmetric hardware scaling on NVIDIA Blackwell GPUs, where tensor core throughput doubles relative to Hopper while shared memory bandwidth and exponential unit throughput remain unchanged. The paper develops three named mechanisms to target these new bottlenecks: a redesigned software pipeline that exploits Blackwell's fully asynchronous MMA operations and larger 128×128 tile sizes to overlap matrix multiplication with softmax computation, software-emulated exponential functions using polynomial approximation on FMA units (increasing exponential throughput beyond the MUFU's fixed 16 ops/clock/SM by offloading 10–25% of exponentials to FMA units), and a 2-CTA MMA mode backward pass that partitions operand B across CTA pairs to reduce shared memory traffic and halves the number of global atomic adds for the dQ gradient. Evaluated on B200 GPUs with BF16 precision, FlashAttention-4 achieves up to 1.3× speedup over cuDNN 9.13 and 2.7× over Triton, reaching 1613 TFLOPs/s — approximately 71% of theoretical peak utilization — while being implemented entirely in CuTe-DSL embedded in Python for 20–30× faster compile times than C++ template-based approaches, establishing that kernel performance on modern accelerators requires explicit algorithmic mitigation of non-matmul bottlenecks only when tensor core throughput grows faster than surrounding functional units.

2. Context and Motivation

The Core Problem: Hardware Scaling Has Broken the Old Assumptions

The fundamental contradiction this paper confronts is that GPUs are getting faster, but attention kernels are not getting faster at the same rate — and the reason is structural, not incidental. Each new generation of NVIDIA datacenter GPUs delivers substantially higher peak tensor core throughput: Hopper H100 provides 1 PFLOPS of FP16/BF16 compute, while Blackwell B200 doubles this to 2.25 PFLOPS. For a naïve analysis, this would suggest that any compute-bound kernel — including attention — should see a proportional speedup. Instead, FlashAttention-3, which was carefully optimized for Hopper's register-file-based tensor cores and 64×128 MMA tile sizes, simply does not run on Blackwell GPUs at all due to lack of forward compatibility for Hopper MMA instructions. Even if it could be ported, a direct translation would leave enormous performance on the table because the hardware balance has fundamentally shifted.

This gap is not a transient engineering issue that will resolve itself with the next GPU generation. The paper identifies a systematic trend in accelerator design:

"tensor core throughput scales faster than other functional units. Blackwell doubles the FP16/BF16 tensor core throughput compared to Hopper... but shared memory bandwidth and exponential unit throughput remain unchanged or scale more slowly. This imbalance shifts the performance bottleneck away from matrix multiplication toward shared memory traffic and non-matmul operations like softmax."

This trend — which the paper terms asymmetric hardware scaling — means that simply increasing peak FLOPS does not translate to proportional application-level speedups. The units responsible for matrix multiplication (tensor cores) are the primary beneficiaries of each new process node and architectural revision because they are the most important for AI workloads. The "supporting" functional units — shared memory, the multifunction unit that computes exponentials, integer ALUs — receive comparatively modest improvements or none at all. The result is an increasingly unbalanced hardware pipeline where the fastest component (tensor cores) must wait on slower components (shared memory reads, exponential evaluations), and the fraction of time spent waiting grows with each generation.

This matters because attention is the computational bottleneck of the Transformer architecture, which powers virtually all modern AI systems. If attention kernels cannot extract close to the theoretical peak throughput of new hardware, then the industry's massive capital investment in each new GPU generation yields diminishing returns. The paper's roofline analysis makes this concrete: for a typical attention forward pass tile configuration of M=N=d=128, the MMA compute requires 1024 cycles, the exponential unit requires 1024 cycles, and shared memory traffic requires 768 cycles — the resources are roughly balanced on Hopper-era hardware assumptions. But for the larger tile sizes that Blackwell's 128×128 MMA tiles enable (M=256, N=d=128), MMA compute and exponential unit both double to 2048 cycles while shared memory traffic increases to 1536 cycles — the bottleneck has shifted, and the exponential unit and shared memory are now the limiting resources, exceeding or matching MMA compute time.

The practical consequence is stark: without explicit algorithmic mitigation of these non-matmul bottlenecks, attention kernels on Blackwell GPUs would leave substantial performance unrealized. The paper's experiments demonstrate that the best existing open-source implementations (Triton with B200-specific instructions) achieve only a fraction of what is possible — FlashAttention-4's 2.7× speedup over Triton is evidence of how much performance is left on the table by kernels that do not account for the shifted bottleneck landscape.

Why This Problem Is Important: The Economics of AI Compute

The importance of this problem extends beyond academic interest in kernel optimization for several concrete reasons.

The industry has already transitioned to Blackwell. Unlike a future architecture that might be years away from deployment, Blackwell-based systems (B200 and GB200) are being deployed at scale across the AI industry at the time of the paper's writing. The authors note that "the AI industry has rapidly transitioned to deploying Blackwell-based systems," citing NVIDIA's architecture technical brief. This means that every organization running large Transformer models on Blackwell hardware — whether for training or inference — is directly affected by the efficiency of attention kernels on this architecture. A 2.7× speedup in the attention layer translates to meaningful improvements in training throughput, inference latency, and total cost of ownership for frontier model development.

Long-context applications amplify the bottleneck. The quadratic scaling of self-attention with sequence length means that as models push toward longer contexts — reasoning over multiple documents, processing entire codebases, or handling high-resolution video — the attention computation dominates runtime to an even greater degree. The paper explicitly connects this to capability: "Scaling attention to longer contexts unlocks new capabilities such as reasoning over multiple documents, modeling entire codebases, and processing high-resolution videos." If attention kernels are inefficient on new hardware, these long-context capabilities become proportionally more expensive, potentially rendering them economically infeasible. Conversely, efficient attention kernels that extract close to peak hardware throughput make long-context applications practical at scale.

Previous approaches do not transfer. The paper makes clear that the prior state-of-the-art — FlashAttention-3 — "primarily targets the NVIDIA Hopper H100 architecture" and that its core techniques depend on Hopper-specific hardware features: register-file-based tensor core accumulators with four threads per row in an interleaved pattern, 64×128 MMA tile sizes, and warp-specialized producer-consumer pipelines designed around Hopper's asynchronous execution model. Blackwell replaces these with fundamentally different primitives: tensor memory (TMEM) for accumulator storage instead of registers, 128×128 MMA tiles (double the area), and fully asynchronous tensor core operations that write directly to TMEM without involving the register file. These are not incremental changes — they constitute a different programming model that requires rethinking the entire kernel design from the ground up.

The paper is explicit about this impossibility of direct porting: "Simply porting existing attention algorithms to this new hardware either leaves significant performance on the table or is impossible due to lack of forward compatibility for Hopper MMA instructions." The "impossible" case refers to the fact that Hopper's MMA instructions (which target registers) cannot execute on Blackwell hardware at all. This means that every technique in FlashAttention-3 that relied on register-level accumulator management, register-pressure-driven scheduling constraints, or Hopper-specific warp-specialization patterns must be redesigned for Blackwell's TMEM-based execution model.

Where Prior Approaches Fall Short

The paper positions itself against a specific landscape of existing attention implementations, each of which falls short on Blackwell hardware in distinct ways.

FlashAttention-3 (Shah et al., 2024). This is the most directly relevant predecessor. FlashAttention-3 achieved state-of-the-art attention performance on Hopper GPUs through three main innovations: warp specialization with asynchronous execution (separating data movement and computation into producer and consumer warps), FP8 low-precision support, and careful scheduling to overlap softmax with matrix multiplication. However, its design assumptions are tightly coupled to Hopper hardware. On Hopper, tensor core accumulators live in registers, which creates extreme register pressure — the kernel must carefully orchestrate which values occupy the limited register file at each point in the pipeline. This register pressure forced FA-3 to effectively serialize much of the backward pass compute graph (compute S, then dP, then dV, then dQ, then dK, with only TMA loads running out of turn). On Blackwell, TMEM eliminates this register pressure by providing a separate 256 KB memory specifically for tensor core accumulators, enabling entirely different scheduling choices that FA-3's architecture cannot exploit.

SageAttention series (Lin et al., 2024a, 2024b, 2025). These approaches achieve speedups through aggressive quantization — INT8 (SageAttention), INT4/FP8 (SageAttention2), and FP4 (SageAttention3). While SageAttention3 does target Blackwell consumer GPUs, the paper notes that "these approaches primarily target consumer GPUs, while most AI compute is deployed on datacenter GPUs." The distinction matters because consumer and datacenter GPUs have different memory hierarchies, different tensor core configurations, and different SM counts — a kernel optimized for a consumer GPU with a small number of SMs and limited memory bandwidth may not scale to the B200's 148 SMs and 192 GB of HBM. Additionally, quantization approaches trade numerical precision for speed; the paper's focus on BF16 precision targets the high-accuracy regime required for training and high-quality inference where quantization error is unacceptable.

Triton (Tillet et al., 2019). Triton provides a higher-level programming model for GPU kernels, and recent versions include B200-specific instructions. However, the paper's benchmarks show that Triton's performance on Blackwell attention workloads is substantially lower than what FlashAttention-4 achieves (2.1–2.7× slower in the forward pass). The gap likely reflects Triton's compiler-based approach, which cannot express the fine-grained warp specialization, explicit TMEM management, and custom pipelining decisions that FlashAttention-4's direct-to-PTX implementation enables. Triton abstracts away these low-level details, which is a productivity win but a performance cost — particularly on hardware with as many novel features as Blackwell.

cuDNN (NVIDIA's vendor library). cuDNN represents the "official" optimized implementation from NVIDIA, with access to internal hardware knowledge and engineering resources. The paper's comparisons to cuDNN 9.13 are the strongest baseline: FlashAttention-4 achieves 1.1–1.3× speedup over this vendor-optimized library. Notably, the paper acknowledges that "since the initial release of our implementation, newer versions of cuDNN have incorporated many of the techniques described in this paper, yielding similar performance to FA4." This suggests both that FlashAttention-4's techniques are the right ones (the vendor independently converged on similar approaches) and that the paper's value is in publicly documenting and analyzing these techniques for the broader community rather than keeping them proprietary.

Standard PyTorch / naive implementations. The paper includes PyTorch as a baseline primarily to show the magnitude of the gap between naïve and optimized implementations, but the comparison is not the focus — a ~10× or greater gap for long sequences is well-established from the original FlashAttention work. The more interesting comparison is against the best available optimized implementations (cuDNN, Triton) where the 1.3–2.7× improvements represent extracting genuinely new performance from hardware that had already received substantial optimization attention.

How This Paper Positions Itself

FlashAttention-4's positioning can be understood along three axes: hardware target, methodological approach, and scope of contribution.

Hardware target: first-to-market Blackwell optimization, not general-purpose. The paper is explicit that its innovations target the B200 and GB200 specifically, not GPU architectures in general. The roofline analysis that motivates the three main techniques (Section 3.1.1 and 3.2.1) is Blackwell-specific: it uses Blackwell's MMA throughput of 8192 ops/clock/SM, Blackwell's MUFU throughput of 16 ops/clock/SM, and Blackwell's shared memory bandwidth of 128 bytes/clock/SM. The specific numbers matter because the bottleneck analysis would produce different results on different hardware — an architecture with faster exponential units or wider shared memory paths would have different resource balance points, and the paper's techniques would need to be re-evaluated. The authors acknowledge this in the conclusion: "some of these algorithms can be extended to other accelerators as compute continues to outpace non-matmul units," but the extension is future work, not a claim of the current paper.

Methodological approach: algorithm-kernel co-design, not pure algorithm or pure engineering. The paper's three main technical contributions — emulated exponentials, conditional softmax rescaling, and 2-CTA MMA mode restructuring — are not purely algorithmic (they would make no sense in a hardware-agnostic analysis) and not purely engineering (they change the mathematical operations performed, not just how operations are scheduled). The exponential emulation changes which functional unit executes the exponential computation, which is an algorithmic choice with numerical accuracy implications. The conditional softmax rescaling changes the online statistics update rule, which is a modification to the attention algorithm's numerical stability mechanism. The 2-CTA backward pass changes how the reduction axis is partitioned, which affects both correctness (atomic add ordering) and performance (shared memory traffic). This co-design philosophy — where hardware characteristics inform algorithmic choices — is the paper's core intellectual contribution, continuing the tradition of the FlashAttention lineage.

Scope of contribution: a framework, not a one-off kernel. While the paper presents a specific B200-optimized attention implementation, its ambition is broader. The implementation in CuTe-DSL (a Python-embedded DSL that compiles to PTX) is presented as a framework for building attention variants, not just a single optimized kernel. The authors describe "independent, composable primitives" for "block-sparse patterns, masking strategies, variable sequence length handling, and work scheduling" that can be "freely combined." This positions FlashAttention-4 as infrastructure — a platform on which future attention variants (FlexAttention, block-sparse, multi-query, grouped-query) can be built with best-in-class performance without requiring each variant to be implemented from scratch in C++ templates.

The paper also explicitly connects to the broader trend in GPU programming toward Python-embedded DSLs that retain low-level control while improving developer productivity. The 20–30× compile time improvement over C++ templates is not just a convenience statistic — it changes the iteration cycle for kernel development from minutes to seconds, enabling experimentation that would be prohibitively slow with traditional approaches. This positions the paper as contributing to both the "what" (how to make attention fast on Blackwell) and the "how" (how to make GPU kernel development more accessible).

Relationship to the FlashAttention lineage. The paper is a direct successor to FlashAttention (Dao et al., 2022), FlashAttention-2 (Dao, 2023), and FlashAttention-3 (Shah et al., 2024), sharing the core IO-aware tiling approach but adapting the pipeline design and algorithmic details to a new hardware generation. Where FA-3's contribution was primarily about exploiting asynchrony and warp specialization on Hopper, FA-4's contribution is about addressing the shifting bottlenecks that emerge when tensor cores outpace their supporting hardware — a problem that FA-3 did not face because Hopper's resources were more balanced for the tile sizes it supported. The paper's abstract frames this progression explicitly: "While FlashAttention-3 optimized attention for Hopper GPUs through asynchronous execution and warp specialization, it primarily targets the H100 architecture." FA-4 extends the lineage to Blackwell, but the extension is not mechanical — it requires fundamentally new algorithmic choices that arise from the asymmetric scaling phenomenon.

3. Technical Approach

3.1 Reader Orientation

FlashAttention-4 is a GPU kernel implementation that computes exact multi-head attention — the core mathematical operation at the heart of every Transformer model — on NVIDIA Blackwell GPUs at speeds approaching the hardware's theoretical limits. The system solves a structural mismatch problem: Blackwell's tensor cores (the matrix multiplication hardware) have doubled in throughput while the supporting hardware units (shared memory for data movement and the exponential unit for softmax computation) have stayed the same speed, meaning that a kernel optimized under old assumptions will spend most of its time waiting on these slower units rather than keeping the tensor cores busy. The solution takes the form of three concrete algorithmic modifications that explicitly reduce demand on the bottlenecked resources — offloading some exponential computations to the underutilized FMA units, skipping unnecessary online softmax rescaling operations, and using a cooperative two-CTA execution mode that nearly halves shared memory traffic and global atomic writes in the backward pass — combined with a redesigned software pipeline that maximizes overlap between computation and data movement under Blackwell's new fully-asynchronous execution model.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components viewed from the highest level:

1. Problem specification — the input tensors (Q, K, V for queries, keys, values), their dimensions (batch size, number of heads, sequence length N, head dimension d), a causal mask flag, and the mathematical operations to be performed (forward pass with two matrix multiplies and a softmax; backward pass with five matrix multiplies, a softmax gradient, and two reductions).
2. Roofline analysis and tile decomposition — before any code runs, the algorithm analytically determines the compute, memory, and exponential throughput requirements for a candidate tile size (M rows of Q, N columns of K, dimension d), comparing each resource's demand against the hardware's supply to identify which unit is the bottleneck and to choose tile dimensions that balance the loads.
3. Forward pass pipeline — a ping-pong schedule where two "warpgroups" (groups of 128 threads) alternate between issuing tensor core MMA instructions for one tile and computing softmax (row-wise max, subtract, exponentiate, rescale, sum) for the other tile, with a separate "correction" warpgroup handling the ongoing output rescaling outside the critical path.
4. Backward pass pipeline — a more complex schedule orchestrating five MMA operations (recomputing the attention score matrix S, computing gradients dP, dV, dS, dQ, dK) across a software pipeline that overlaps the softmax elementwise operations with dQ and dK MMAs from the previous inner-loop iteration, with tensor memory (TMEM) partitioned to share accumulator storage across multiple operations.
5. Scheduling policy — a grid traversal strategy that reorders which sequence-length tiles are processed by which streaming multiprocessors (SMs) to minimize load imbalance, using longest-processing-time-first ordering for causal masking (processing long blocks first so short blocks fill in idle SMs later) and optional pre-sorting of variable-length batches.

Information flows as follows: the host launches a CUDA grid of thread blocks (CTAs) organized over batch, head, and query-tile dimensions. Each CTA iterates over KV tiles in its assigned order, loading the current K and V tiles from global memory into shared memory via TMA (tensor memory accelerator) asynchronous copies. For each inner iteration, the CTA loads a Q tile, executes the forward (or backward) compute graph on that tile using tensor core MMAs for matrix multiplies, the software-emulated or hardware exponential unit for softmax elementwise operations, and shared memory for data exchange, producing output tiles that are written back to global memory. The key innovation is not in what gets computed — exact attention with no approximations — but in which functional unit performs each sub-operation, how those sub-operations are interleaved in time, and how data is routed through the memory hierarchy to minimize time spent on the slowest links.

3.3 Roadmap for the Deep Dive

• First, the forward pass roofline analysis and pipeline design (Section 3.1), because the forward pass is simpler (two MMAs, one softmax) and the roofline directly motivates both the exponential emulation and the conditional rescaling techniques.
• Second, the exponential emulation mechanism (Section 3.1.3), since it is a standalone algorithmic contribution (polynomial approximation of 2^x on FMA units) that interacts with the pipeline design through its effect on register pressure and throughput.
• Third, the conditional softmax rescaling (Section 3.1.4), which modifies the online statistics update rule to skip unnecessary vector rescaling operations — a change to the attention algorithm itself that reduces the number of non-matmul operations.
• Fourth, the backward pass roofline and pipeline (Section 3.2), which is more complex (five MMAs, softmax gradient, two reductions) and introduces the need for the 2-CTA mode.
• Fifth, the 2-CTA backward pass mechanism (Section 3.2.3), which uses Blackwell's cooperative MMA mode and distributed shared memory to reduce shared memory traffic and halve atomic adds, then the deterministic execution variant (Section 3.2.4).
• Sixth, the scheduling policy (Section 3.3), which addresses load imbalance across CTAs and is largely hardware-agnostic, applying to both forward and backward passes.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a hardware-algorithm co-design paper whose core insight is that on modern GPUs where tensor core throughput has out-scaled other functional units, attention kernel performance is bottlenecked by shared memory bandwidth and exponential unit throughput rather than by matrix multiplication speed — and that explicit algorithmic modifications to reduce demand on these bottlenecked resources (while exploiting new hardware features like TMEM and 2-CTA MMAs) can recover performance that a direct port of prior-generation kernels would leave on the table.

---

3.4.1 Forward Pass Roofline Analysis and Tile Configuration

The paper begins with an analytical model of resource requirements for a single forward-pass inner-loop iteration on a tile of dimensions M (query-sequence rows) by N (key-sequence columns) by d (head dimension). The purpose is to identify, before writing any code, which hardware unit will be the bottleneck for a given tile size — this determines what the kernel must optimize for.

MMA compute cycles. The forward pass performs two matrix multiply-accumulate operations per tile: first, computing the attention scores S = QK^T (an M×N output from M×d and d×N inputs), and second, computing the output O = PV (an M×d output from M×N and N×d inputs). Each MMA requires 2MNd floating-point operations. With Blackwell's tensor core throughput of 8192 FLOPs per clock cycle per SM, the total time spent in tensor core computation is:

$T_{\text{MMA}} = \frac{4MNd}{8192} \text{ cycles}$

where the numerator 4MNd accounts for 2MNd operations per MMA times two MMAs.

What it computes: the number of SM clock cycles required to execute both matrix multiplications if tensor core throughput were the only constraint. This is the "ideal" compute time — what the kernel would spend if it had infinite memory bandwidth and instant softmax.

Why this form: the denominator 8192 is the BF16 tensor core throughput per clock per SM on Blackwell (2.25 PFLOPS / 1850 MHz / 148 SMs = 8192), representing the maximum rate at which the hardware can process fused multiply-add operations. The numerator counts the total FMA operations (each multiply-add is two FLOPs, and 2MNd multiply-adds require 4MNd FLOPs). This formulation treats the tensor cores as a fixed-throughput resource and models the time to exhaust a given amount of work — standard roofline methodology adapted to attention's specific operation count.

Shared memory traffic cycles. Not all MMA operations read their operands from the same place in the memory hierarchy. The first MMA (computing S = QK^T) is "shared-shared" (SS): both the Q tile and the K^T tile are read from shared memory. The second MMA (computing O = PV) is "tensor-shared" (TS): the P matrix (attention probabilities) is read from tensor memory (TMEM) where the softmax left it, while the V tile is read from shared memory.

Because Blackwell's MMA instructions operate on tiles of size 128×128 elements, computing an M×N output requires ceil(M/128) × ceil(N/128) individual MMA instructions. Each of these instructions reads its operands from shared memory — and critically, when multiple instructions are required to cover the output, the shared memory operands are read multiple times, amplifying the bandwidth demand.

For the first MMA (QK^T), the total shared memory reads are ceil(M/128) × ceil(N/128) × (128d + 128d) = ceil(M/128) × ceil(N/128) × 256d elements, since each of the MMAs reads a 128×d chunk of Q and a d×128 chunk of K^T. For the second MMA (PV), the total shared memory reads are ceil(M/128) × ceil(d/128) × 128N elements, since each MMA reads an N×128 chunk of V from shared memory (P comes from TMEM, not SMEM).

Assuming M, N, and d are all multiples of 128 (which the paper's tile selection ensures), and with each BF16 element occupying 2 bytes and shared memory bandwidth being 128 bytes per cycle per SM, the total shared memory read cycles are:

$T_{\text{smem}} = \frac{2 \cdot \frac{M}{128} \cdot \frac{N}{128} \cdot 256d + 2 \cdot \frac{M}{128} \cdot \frac{d}{128} \cdot 128N}{128} = \frac{3MNd}{8192} \text{ cycles}$

What it computes: the number of SM clock cycles required to read all shared memory operands for both MMAs. This is the "memory floor" — how long the kernel must spend just moving data from shared memory into the tensor cores, irrespective of how fast the tensor cores can compute.

Why this form: the numerator tracks every 2-byte BF16 element read from shared memory across all MMA instructions, accounting for reuse (or lack thereof). The factor of 3/8192 emerges algebraically from the tile dimensions and SMEM bandwidth — it tells us that the SMEM read time is 3/4 of the MMA compute time (since T_MMA = 4MNd/8192). This fraction is a critical insight: on Hopper, where tile sizes were smaller, the SMEM-to-compute ratio was different, and shared memory was less of a bottleneck relative to compute. The larger tiles that Blackwell enables (because TMEM can hold larger accumulators without register pressure) increase this ratio.

Exponential unit cycles. The softmax operation requires computing e^x for every element of the M×N attention score matrix S. The multifunction unit (MUFU) on Blackwell provides 16 exponential operations per clock cycle per SM. The time required is:

$T_{\text{exp}} = \frac{MN}{16} \text{ cycles}$

What it computes: the number of SM clock cycles required to compute all exponentials for one tile's softmax, assuming all exponentials are executed on the dedicated MUFU hardware unit.

Why this form: the denominator 16 is the MUFU.EX2 throughput — each SM can execute 16 exponential instructions per clock. This is unchanged from Hopper. The numerator MN is the total number of attention scores that need exponentials. Unlike the MMA operations which benefit from Blackwell's doubled throughput (4096→8192), the exponential throughput is stagnant, creating a widening gap between how fast tensor cores can produce the scores and how fast the MUFU can exponentiate them.

Comparative analysis. Table 1 in the paper summarizes these cycle counts for two tile configurations. For M = N = d = 128, the resources are roughly balanced: MMA = 1024 cycles, exponential = 1024 cycles, SMEM = 768 cycles. For the larger tile M = 256, N = d = 128, both MMA and exponential double to 2048 cycles while SMEM increases to 1536 cycles. The key insight is that for the larger tile — which Blackwell's hardware makes feasible because TMEM can hold the 128×128 accumulator tiles without consuming registers — the exponential unit and shared memory are the co-bottlenecks, not the tensor cores. This is a reversal from the Hopper era where tensor core throughput was typically the limiting factor for compute-bound attention kernels.

Design implication. This analysis directly motivates the three forward-pass optimizations. First, because the exponential unit is now a co-bottleneck, the kernel must increase exponential throughput beyond the MUFU's fixed 16 ops/clock — hence the polynomial emulation on FMA units. Second, because shared memory traffic scales with tile size, the pipeline must maximize overlap between SMEM reads and other operations so that SMEM latency is hidden — hence the redesigned ping-pong schedule. Third, because softmax rescaling operations (vector multiplies that renormalize the running output when a new maximum is found) are non-matmul work that competes for shared resources, eliminating unnecessary rescaling steps reduces demand on the bottlenecked units — hence the conditional rescaling.

---

3.4.2 Forward Pass Software Pipeline Design

The forward pass pipeline orchestrates three concurrent activities: tensor core MMA operations (producing attention scores S and output O), softmax computation (maximum-finding, exponential evaluation, row-sum accumulation), and output rescaling (renormalizing previously accumulated output when a new row-maximum is discovered). The design challenge is to interleave these activities so that the tensor cores are never idle waiting for softmax results, and conversely, the softmax computation never starves for data.

Warpgroup assignment. The kernel uses four warpgroups (groups of 128 threads each) within each CTA, assigned to specific roles:

• Two softmax warpgroups — each processes one tile's worth of attention scores in a ping-pong alternation. When warpgroup A is computing softmax on tile i, warpgroup B is computing softmax on tile i+1. Each warpgroup thread owns an entire row of the 128×128 score tile, eliminating the need for inter-warp communication (shuffles) during the row-maximum reduction — a significant simplification over FA-3's design where each row was split across four threads in an interleaved register layout.
• One tensor-core / TMA warpgroup — drives tensor core MMA instructions and asynchronous TMA (Tensor Memory Accelerator) data movement, issuing loads from global memory to shared memory and MMA operations from shared memory to tensor memory.
• One correction warpgroup — performs the output rescaling operation (multiplying the previously accumulated output O by the rescale factor e^(m_old - m_new) when a new maximum is discovered), decoupled from the softmax critical path.

Why four warpgroups: on Hopper/FA-3, the correction step was part of the softmax warpgroup's responsibility, consuming registers and adding latency to the softmax critical path. By moving it to a dedicated warpgroup, FA-4 removes this work from the bottleneck. This is only possible because Blackwell's TMEM serves as a communication channel — the softmax warpgroup writes the rescaling statistics (m, the running maximum, and l, the running normalizer) to TMEM, and the correction warpgroup reads them from TMEM, avoiding the register-pressure constraints that would make this impossible on Hopper.

Ping-pong schedule. The pipeline proceeds in the following order (referenced to Figure 1 in the paper, which shows the schedule for two Q tiles labeled "high" and "low," each corresponding to 128 query tokens):

1. Phase 1: Compute two S tiles. Both Q tiles (S^H = Q^H K^T and S^L = Q^L K^T) are computed via tensor core MMAs as early as possible, writing the results to TMEM. This is possible because the pipeline allocates two S-tile slots in TMEM rather than only one — a deliberate memory allocation choice that enables starting the pipeline with maximal MMA utilization.

2. Phase 2: Softmax on S^H while MMA on S^L. While warpgroup A runs softmax on S^H (finding row-wise maximums, subtracting, exponentiating, computing row sums), the tensor-core warpgroup simultaneously resumes MMA operations on S^L. The softmax writes the resulting attention probabilities P^H to TMEM.

3. Phase 3: MMA on P^H V while softmax on S^L. The tensor-core warpgroup executes the second MMA (O^H = P^H V) while warpgroup B computes softmax on S^L. This overlap is the core of the pipeline: the tensor cores never wait for softmax because there is always another tile's MMA ready to execute.

4. Phase 4: Continue alternating. The pattern continues with subsequent KV tiles, with the softmax and tensor-core activities proceeding on alternating tiles.

TMEM partitioning strategy. The paper describes two viable TMEM partitioning options and explains the choice. With head dimension 128, the kernel needs to allocate accumulator space for two output tiles (each 128×128 BF16 elements = 2 × 128 × 128 × 2 bytes = 64 KB). The 256 KB TMEM thus has ~192 KB remaining for S and P storage. Each S tile (128×128 FP32 elements = 128×128×4 = 64 KB) and each P tile (128×128 BF16 = 32 KB) compete for this space.

The two options are: (Option A) store one S tile and two P tiles (64 + 2×32 = 128 KB, feasible), or (Option B) store two S tiles that overlap in memory with P tiles. The paper chooses Option B because "it allows us to start our software pipeline by immediately computing two S tiles" — the two S tiles are written to TMEM first, then as softmax consumes each S tile to produce a P tile, the P tile overwrites the S tile's TMEM region. This overlap is safe because once softmax has read S, it is no longer needed. This choice also leaves a small TMEM region for communicating rescaling statistics (running max m and normalizer l) from the softmax warpgroups to the correction warpgroup.

Register pressure management. A critical constraint is that each softmax warpgroup thread must hold an entire 128-element row in registers — 128 BF16 values for the input row plus space for the output P row, coefficients, and temporaries. To stay within the 256-register limit, the kernel stages the storage of P in quarters: the first three quarters are written to TMEM (triggering the corresponding PV MMA operations to begin consuming them), and only the last quarter is stored separately. This reduces peak register demand during the softmax computation because registers can be reused as soon as each quarter of P is written.

Synchronization discipline. Just as in FA-3, the two softmax warpgroups must not overlap their critical sections — the part of the computation that evaluates exponentials. The paper states: "we explicitly synchronize the two softmax warpgroups to not overlap in their critical section, which is the part of exponential computation." The reason is that the MUFU exponential unit (and the FMA units that execute the emulated exponentials) are shared resources within an SM — if both warpgroups tried to issue exponential instructions simultaneously, they would contend for the same functional units and serialize anyway, while also consuming more shared resources (register bandwidth, instruction cache). The synchronization ensures that while one warpgroup is in its exponential-critical-section, the other is either in a non-critical phase (loading data, computing max, accumulating sums) or is idle, and the tensor-core warpgroup is free to execute MMAs without contention.

---

3.4.3 Software-Emulated Exponential Functions

The exponential unit (MUFU.EX2) is a fixed-function hardware unit that computes 2^x for BF16 inputs. On Blackwell, its throughput is 16 operations per clock per SM — unchanged from Hopper — while tensor core throughput has doubled. For the forward pass tiles with M=N=128, the exponential unit requires 1024 cycles, which exactly matches the MMA compute time. This means there is no inherent slack: if the exponential computation takes any longer than the MMA computation, it becomes the bottleneck and the tensor cores sit idle. The exponential emulation technique increases effective exponential throughput by offloading a fraction of exponential evaluations to the floating-point FMA (fused multiply-add) units, which have much higher aggregate throughput and are typically underutilized during the softmax phase.

The core decomposition. The technique computes 2^x using the multiplicative property 2^x = 2^(floor(x)) × 2^(x - floor(x)), where the integer part handles the exponent field of the floating-point representation directly through bit manipulation, and the fractional part (in [0,1)) is approximated by a polynomial evaluated on FMA units.

The integer part computation exploits the IEEE 754 binary32 representation: a floating-point number is stored as sign × 2^(exponent - 127) × (1.mantissa). To compute 2^(floor(x)), the algorithm extracts floor(x) as an integer, adds it to the exponent bias (127), and inserts the result into the exponent field of the output — effectively constructing the floating-point representation of 2^(floor(x)) without any arithmetic operations on the mantissa. Specifically:

1. Clamp x to be at least -127 to avoid underflow (inputs smaller than -127 would produce denormalized outputs or zero, which would lose precision unnecessarily).
2. Compute floor(x) using a round-down trick: add the constant 2^23 + 2^22 (which forces the fractional bits into the mantissa bits of the floating-point representation), then subtract it back with round-down rounding mode. The difference gives the integer floor.
3. Compute the fractional part as x_frac = x - floor(x), which lies in [0, 1).
4. Evaluate a polynomial p(x_frac) that approximates 2^(x_frac) on [0,1), using Horner's method with FMA instructions: p(x) = p_0 + x × (p_1 + x × (p_2 + x × p_3)) for a degree-3 polynomial.
5. Combine: shift floor(x) into the exponent field (adding it to the bias) and insert the mantissa bits from the polynomial result. The combination effectively multiplies 2^(floor(x)) × p(x_frac), yielding the final 2^x value.

Polynomial degree and accuracy tradeoff. The paper evaluates polynomial approximations of degrees 3 through 6, comparing against the hardware MUFU.EX2 instruction using 4 million random inputs in [0,1) and measuring error against a double-precision (FP64) reference. Table 2 reports the results:

Key insight from the accuracy table: at BF16 precision, all polynomial degrees from 3 upward produce the same maximum relative error (3.91 × 10^(-3)) as the hardware MUFU.EX2 instruction. This is because the BF16 quantization error (rounding from FP32 to the BF16 8-bit mantissa) dominates the polynomial approximation error. The degree-3 polynomial, which is 600× less accurate than hardware at FP32 precision (8.8 × 10^(-5) vs. 1.52 × 10^(-7)), becomes indistinguishable after rounding to BF16 because the rounding operation truncates mantissa bits beyond the 8th position. The degree-3 polynomial matches the hardware exponential to within 1 ULP (unit in the last place) of BF16 on 99.42% of inputs — sufficient for attention computation where the softmax output is consumed at BF16 precision in the subsequent matrix multiplication.

The paper selects the degree-3 polynomial as the default. Higher degrees (4-6) close the FP32 accuracy gap to within 2× of hardware but require additional FMA instructions (each degree adds one multiply-add), increasing register pressure and latency without improving the BF16-level output. The degree-3 polynomial strikes the optimal balance between accuracy and instruction count for this use case.

Partial emulation strategy. The polynomial approximation is not applied to all exponentials in a softmax row. The paper applies emulation to only 10-25% of the entries in each row, with the remaining entries computed via the hardware MUFU.EX2 instruction. There are two reasons for this:

First, the emulated exponential consumes more registers than the hardware instruction (to hold the polynomial coefficients and intermediate values) and has higher instruction latency (several FMA instructions vs. one MUFU instruction). Using emulation for all exponentials would increase register pressure to the point of causing register spills to stack memory, which would negate the throughput benefit.

Second, the MUFU and FMA are separate functional units that can execute in parallel. By sending ~10-25% of exponentials to the FMA units while the MUFU handles the remaining ~75-90%, the effective exponential throughput increases because both units operate simultaneously. The exact fraction is "tuned empirically based on the ratio of MMA and exponential throughput for a given tile configuration" — when the MMA throughput is much higher than the exponential throughput (which is the case for larger tiles), a higher emulation fraction is warranted.

Why this is an algorithmic contribution, not just engineering. The exponential emulation changes which functional unit performs a given mathematical operation — it exploits the fact that the hardware has multiple types of computation units with different throughput characteristics, and the standard softmax implementation routes all exponentials through the slowest unit (MUFU) while leaving the faster FMA units idle. By re-routing some exponentials through polynomial evaluation on FMA, the effective exponential throughput increases beyond the MUFU's fixed 16 ops/clock limit. This is possible only because (a) the BF16 precision requirement is low enough that a cheap polynomial approximation matches the hardware's accuracy, and (b) the FMA units have sufficient spare capacity during the softmax phase. On hardware where either of these conditions is false (higher precision requirements or fully saturated FMA units), the technique would not be beneficial.

---

3.4.4 Conditional Online Softmax Rescaling

FlashAttention uses the "online softmax" algorithm to compute attention in tiles without accessing the entire sequence at once. This algorithm maintains running statistics as it processes successive blocks of keys and values, renormalizing the accumulated output whenever a new row-maximum value is discovered. The conditional rescaling technique reduces the frequency of these renormalization operations by tolerating small discrepancies in the running maximum.

The standard online softmax update rule. Let m_{j-1} be the running row-wise maximum after processing blocks 1 through j-1, l_{j-1} be the running row-wise sum of exponentials (the normalizer), and O_{j-1} be the running output (the weighted sum of value vectors). When processing block j, the attention scores S_j = Q K_j^T are computed, and the new running statistics become:

$m_j = \max(m_{j-1}, \text{rowmax}(S_j))$

$\ell_j = e^{m_{j-1} - m_j} \ell_{j-1} + \text{rowsum}(e^{S_j - m_j})$

$O_j = e^{m_{j-1} - m_j} O_{j-1} + e^{S_j - m_j} V_j$

What it computes: the updated running maximum m_j (taking the elementwise maximum of the old running max and the new block's row-max), the updated normalizer l_j (rescaling the old sum by the ratio of old-to-new exponential maxima and adding the new block's exponentiated scores), and the updated output O_j (rescaling the accumulated output and adding the new block's contribution). The rescaling factor e^(m_{j-1} - m_j) is less than or equal to 1 — it "downweights" previous contributions when a new, larger maximum is found, maintaining numerical stability by ensuring that the largest exponentiated value is always 1 (or close to it).

Why this form is necessary. Without rescaling, the exponentiated scores e^(S_j) could overflow or lose precision because the raw attention scores can be large (proportional to sqrt(d)). The running-max subtraction normalizes the scores so that the maximum exponentiated value in each row is 1, keeping all values in [0, 1] and preventing overflow. The rescaling factor e^(m_{j-1} - m_j) is the mechanism that reconciles the old (possibly too small) normalization with the new, corrected one.

The conditional modification. The paper makes two observations:

First, rescaling is only necessary when m_j > m_{j-1} — that is, when the new block contains a value larger than any seen so far. If the new maximum equals the old maximum, the rescaling factor is e^0 = 1 (no change), and the rescaling multiplication is a no-op. Many blocks do not increase the running maximum, especially for later blocks in the sequence where the maximum has already been found.

Second, even when m_j > m_{j-1}, the kernel can tolerate some "slack" — it can delay rescaling until the discrepancy exceeds a threshold τ. The paper introduces:

$O_j = \begin{cases} e^{m_{j-1} - m_j} O_{j-1} + e^{S_j - m_j} V_j & \text{if } m_j - m_{j-1} > \tau \\ O_{j-1} + e^{S_j - m_{j-1}} V_j & \text{otherwise} \end{cases}$

where τ is typically set to log_2(256) = 8.0, corresponding to a rescaling factor threshold of 256.0.

What it computes: the same output O_j, but with two execution paths. When the new maximum exceeds the old maximum by more than τ (the difference in natural-log space corresponds to a factor of more than 256× in linear space), the standard rescaling path executes — the old output is down-weighted and the new contributions are normalized by the new maximum. When the new maximum is within τ of the old maximum, the rescaling is skipped — the old output is kept at its existing normalization (using m_{j-1}) and the new contributions are normalized by m_{j-1} instead of m_j. The running maximum m_j is not updated in the "skip" case; m_{j-1} continues to be used.

Why this form is correct. At the end of all blocks, the final output is normalized by the true final maximum m_final and the true final normalizer l_final:

$\text{Output} = \frac{1}{\ell_{\text{final}}} O_{\text{final}}$

Any intermediate rescaling that was skipped is corrected in this final step, because l_final incorporates the accumulated scaling factors and the final maximum. The conditional path accumulates some values at a slightly-offset normalization (using m_{j-1} when the true current maximum is m_j), but the final normalization step compensates. The threshold τ = 8.0 (factor 256) is chosen to ensure that the accumulated values do not underflow below BF16's representable range before the final correction — values scaled by up to 256× are still well within BF16's dynamic range.

Avoiding warp divergence. The conditional rescaling is applied per row of the attention matrix, but GPU warps execute in lockstep — all 32 threads in a warp execute the same instruction. If some threads in a warp need rescaling while others don't, the warp would diverge (take the rescaling branch on some threads, the skip branch on others), serializing the execution and losing the performance benefit. The paper addresses this: "to avoid warp divergence, we rescale when any of the threads in the warp needs rescaling." This conservative policy means that if at least one row in the warp's responsibility has a new maximum exceeding the threshold, the entire warp executes the rescaling path. The overhead of unnecessary rescaling on a few rows is smaller than the cost of warp divergence.

Practical impact. The conditional rescaling reduces the number of vector-multiply operations (each rescaling requires multiplying an entire M×d output tile by a scalar factor e^(m_{j-1} - m_j), which is M×d multiply operations) during the online softmax computation. In the common case where attention scores are dominated by a few large values (which is typical for trained Transformers, where attention is sparse), the maximum is found early in the sequence and subsequent blocks do not trigger rescaling. The paper does not report the exact reduction in rescaling operations, but the roofline analysis in Section 3.1.1 identifies these non-matmul operations as contributors to the softmax bottleneck, and reducing them directly reduces demand on the shared functional units (FMA units, register bandwidth) that are the limiting resources.

---

3.4.5 Backward Pass Roofline and Software Pipeline

The backward pass is substantially more complex than the forward pass: it must compute five matrix multiplications (recomputed S, then gradients dV, dP, dS, dQ, dK), a softmax gradient (involving elementwise multiplication, summation, and the identity dS = (diag(p) - pp^T) dP), and two global memory reductions (for dQ accumulation across sequence blocks). The roofline analysis in Section 3.2.1 models these requirements for the tile configuration M = N = d = 128, comparing against Blackwell's MMA (8192 FLOPs/clock), shared memory (128 bytes/clock), and exponential (16 ops/clock) throughputs.

MMA compute cycles. With five MMAs each requiring 2MNd operations:

$T_{\text{MMA}} = \frac{10MNd}{8192} \text{ cycles}$

For M=N=d=128, this equals 2560 cycles.

What it computes: the tensor core time for all five MMAs. The factor of 10 (five operations × 2 FLOPs per multiply-add) is 2.5× the forward pass's factor of 4 (two operations × 2 FLOPs per multiply-add), reflecting the backward pass's higher arithmetic intensity.

Shared memory traffic cycles. Among the five MMAs, three are shared-shared (SS) — both operands read from shared memory — and two are tensor-shared (TS) — operand A from TMEM, operand B from shared memory. The SS MMAs contribute reads from shared memory: the S recomputation (2Md + 2Nd for reading Q and K^T), dP (2Md + 2Nd for reading V and dO^T), and dQ (2MN + 2Nd for reading dS and K), totaling 2Md + 3Nd + MN elements. The TS MMAs, dV and dK, contribute 2Md total reads. Additionally, the kernel writes dS (M×N BF16 = MN/64 cycles) and dQ (M×d FP32 = 8Md bytes or Md/16 cycles) to shared memory.

$T_{\text{smem}} = \frac{4Md + 3Nd + MN}{64} + \frac{MN}{64} + \frac{Md}{16} \text{ cycles}$

For M=N=d=128, this equals 3328 cycles — approximately 30% more than the MMA compute time (2560 cycles). This makes shared memory bandwidth the primary bottleneck for the backward pass.

What it computes: the total shared memory access time for all MMA operands, intermediate gradient writes, and reduction traffic. The first fraction covers MMA operand reads (each BF16 element is 2 bytes, and SMEM bandwidth is 128 bytes/clock, so element count / 64 gives cycles), the second fraction covers the dS write (MN BF16 elements = MN/64 cycles), and the third covers the dQ write and TMA read (4 bytes per FP32 element, written and read = 8Md bytes = Md/16 cycles).

Why this form matters. The SMEM time (3328 cycles) exceeds the MMA compute time (2560 cycles) by 768 cycles — almost exactly the exponential unit time (1024 cycles, from T_exp = MN/16). This means that even if the kernel could perfectly overlap MMA and exponential computation, the shared memory traffic would still be the long pole, and any reduction in SMEM traffic directly improves end-to-end performance. This motivates the 2-CTA MMA mode, which reduces shared memory operand reads by nearly half.

Exponential unit cycles. The backward pass requires exponential operations for the softmax recomputation and its gradient:

$T_{\text{exp}} = \frac{MN}{16} \text{ cycles}$

For M=N=128, this equals 1024 cycles — significant but not the primary bottleneck compared to SMEM.

Pipeline design (1-CTA mode). The backward pass must execute five MMA operations and two elementwise operations (softmax, dsoftmax) in an order that respects data dependencies while maximizing overlap. The paper's key insight is that on Blackwell, TMEM enables much more flexible scheduling than FA-3's register-file constraints allowed. In FA-3, accumulator tiles were held in registers, and the severe register pressure forced an effectively serial schedule: compute S, then dP, then dV, then dQ, then dK, with only TMA loads running significantly out of turn.

In FA-4, TMEM provides a separate accumulator storage that doesn't consume registers. The pipeline is structured as an iteration over KV tiles in the outer loop: for each KV tile, the kernel loads K and V, then streams over Q tiles in the inner loop. The scheduling key is to overlap the softmax/dsoftmax elementwise computation (which uses the MUFU and FMA units) with MMA operations that don't depend on its output.

Specifically, for iteration i of the inner loop (processing Q tile i against the current KV tile), the kernel:

1. Launches the MMA to recompute S_i = Q_i K^T (producing the attention score tile).
2. Launches the MMA for dP_i = dO_i V^T (computing the gradient of the attention probabilities), which depends on the recomputed S_i only through the softmax (which hasn't run yet — this is an out-of-order execution enabled by TMEM: the MMA writes to TMEM, and the dP computation will read from TMEM once softmax has written P_i there).
3. Runs the softmax-dsoftmax elementwise operations on S_i (from step 1) to produce dS_i.
4. Launches the MMA for dQ_i = dS_i K and dK_i = dS_i^T Q_i, using dS_i produced in step 3.

The critical overlap pattern is: while the MMA for dQ_i-1 and dK_i-1 from the previous iteration are still executing (using the previous iteration's dS), the kernel starts the dS elementwise computation for the current iteration, using separate TMEM regions so there is no conflict. This exploits TMEM's larger capacity compared to the register file — TMEM can hold accumulator tiles for multiple iterations simultaneously, whereas registers in FA-3 could only hold the current iteration's accumulators.

TMEM partitioning for backward pass. The backward pass has more accumulator tiles than TMEM can hold simultaneously. With head dimension 128, each accumulator tile (128×128 FP32) occupies 64 KB. Five such tiles would require 320 KB, exceeding the 256 KB TMEM capacity. The paper's partitioning:

• TMEM block at offset 0: shared between S (computed first) and P (computed by softmax from S). S is written by MMA, consumed by softmax, then P overwrites S. P is consumed by subsequent MMAs.
• TMEM block at another offset: shared between dP, dS, and dQ. dP is written by MMA, consumed by dsoftmax to produce dS. dQ is written by MMA for the current iteration's dQ tile (which is small enough to fit alongside these).

This "overlay" strategy — reusing TMEM regions for different purposes at different times — is only possible because of the careful pipeline ordering that ensures each tile is consumed before its region is reused. The paper's Figure 2 illustrates the computational graph and the software pipeline order across the prologue (warmup), main loop (steady state), and tail (draining) phases.

---

3.4.6 2-CTA Backward Pass: Reducing Shared Memory Traffic and Atomic Adds

The single-CTA backward pass leaves shared memory bandwidth as the bottleneck (3328 cycles vs. 2560 cycles of MMA compute). The 2-CTA MMA mode, a new Blackwell feature, allows two CTAs (thread blocks) within the same thread block cluster to cooperatively execute a single MMA instruction. This enables two performance improvements: (1) each CTA loads and stages only half of operand B in its shared memory, reducing total shared memory traffic, and (2) the dQ gradient accumulation, which requires atomic additions to global memory, can be restructured to halve the number of atomic operations.

How 2-CTA MMA mode works. In standard (1-CTA) MMA mode, a single CTA executes an MMA instruction that reads operand A (M×K tile) and operand B (K×N tile) from the CTA's own shared memory, producing an M×N accumulator. In 2-CTA mode, the M dimension of the output tile is partitioned across the two CTAs: each CTA owns M/2 rows of the output and has its own accumulator in its own tensor memory. The N dimension of operand B is also partitioned: each CTA loads and stages only N/2 columns of B in its own shared memory, and the hardware logically concatenates the two halves to form the full B operand for the multiply. This means that each CTA's shared memory holds only half of B, reducing the per-CTA shared memory traffic for that operand.

Shared memory traffic reduction. In the backward pass, the MMA for recomputing S, computing dP, computing dV, and computing dK all use an MMA tile shape of M=256 and N=K=128 (with the M=256 dimension split across the two CTAs, so each CTA handles M=128). For these MMAs, each CTA stages only half of operand B (N=64 per CTA instead of N=128), approximately halving the shared memory reads for operand B. Across the five MMAs, this reduces total SMEM traffic, bringing it closer to (though still slightly above) the MMA compute time. Table 3 notes that in the 2-CTA setting with M=256, N=d=128, "shared memory traffic exceeds MMA compute time by approximately 5%" — a substantial reduction from the 30% excess in the 1-CTA case with M=128.

Restructuring the dQ computation. The dQ gradient presents a special challenge. In the FlashAttention backward pass, each CTA processes a fixed KV tile (parallelized over N CTAs, each responsible for a different portion of the sequence length) and streams over Q tiles (M) in the inner loop. The dQ gradient for a given Q tile accumulates contributions from all KV tiles, requiring a global memory reduction — each inner-loop iteration atomically adds its contribution to a dQ buffer in global memory.

In the 1-CTA mode, each CTA computes a full dQ tile of M×d and performs M×d atomic adds to global memory. In the 2-CTA mode, the natural partitioning would give each CTA an M/2 × d portion of dQ, but the atomic reduction would still need to reduce over the N dimension (the KV sequence length). The paper identifies a conflict: "the reduction dimension of the dQ MMA is N, which is naturally split across the CTA pair." This means the dQ MMA tile shape in the 2-CTA mode would be (M/2, N) × (N, d), using N as the reduction dimension — but N is the dimension being reduced over (accumulated across outer-loop iterations), and splitting N across CTAs would mean each CTA only sees half the reduction, producing an incorrect partial result.

Using distributed shared memory (DSMEM) to repack dS. The solution uses DSMEM — shared memory that is accessible across CTA pairs within the same thread block cluster — to exchange data between the two CTAs. The procedure, illustrated in Figure 3:

1. Each CTA computes dS for its assigned rows (M/2 rows per CTA) and writes it to its own shared memory.
2. The CTAs exchange half of their dS data via DSMEM, so that each CTA ends up with a dS tile of shape (M/2, 2N) — its own M/2 rows but all 2N columns of dS (double the original N, because it now has both CTAs' dS data concatenated along the column dimension).
3. Each CTA executes the dQ MMA with tile shape (M/2, 2N) × (2N, d), accumulating an (M/2, d) tile in tensor memory. The reduction dimension is now 2N (spanning the full KV sequence block), so the reduction is complete and correct.
4. Each CTA writes its (M/2, d) tile to global memory via atomics, performing M/2 × d atomic adds — half the number of the 1-CTA case.

Why this works. The key is that dS has been repacked from being partitioned along the reduction axis (N, which would be split across CTAs) to being partitioned along the non-reduction axis (M, which is already split). Each CTA now holds a full copy of dS along the reduction dimension (2N columns), so it can perform the complete reduction independently within its own tensor core, producing a valid (M/2, d) partial result. The communication cost is the DSMEM transfer of the dS halves — but this is a single exchange per inner-loop iteration, and the latency is hidden by the software pipeline.

Pipeline reordering for DSMEM latency hiding. The 2-CTA backward pass reorders its pipeline relative to the 1-CTA version to hide the DSMEM transfer latency. Specifically, the dP MMA for the current tile (iteration i) is computed before the dQ MMA for the previous tile (iteration i-1). This means that while the dQ MMA from the previous iteration is executing (including the DSMEM transfer to repack dS from iteration i-1), the current iteration's dP MMA can execute concurrently because it uses different TMEM regions and different shared memory operands. The elementwise dS computation for the current tile runs in parallel with the dQ MMA from the previous tile, achieving the overlap that the pipeline is designed for.

TMEM reuse in 2-CTA mode. The dQ tile in the 2-CTA mode is only (M/2, d) = (64, 128) = 16 KB in FP32, which is small enough to fit in TMEM alongside other tiles. The paper reuses the same TMEM region that held S (and later P) for the dQ accumulator, since S and P are no longer needed by the time dQ is computed.

Halving global atomic adds. Each CTA now accumulates dQ for only M/2 rows instead of M rows, so the number of global atomic add operations in each inner-loop iteration is halved. Global atomics are expensive for two reasons: they are high-latency operations that go through the L2 cache (not just shared memory), and they introduce nondeterminism because the order in which CTAs complete their atomics affects the floating-point accumulation result (floating-point addition is not associative). Reducing the atomics count by 2× directly improves both performance and the degree of nondeterminism.

---

3.4.7 Deterministic Backward Pass

The standard backward pass introduces nondeterminism in the gradient computation because multiple CTAs writing to the same dQ (and in grouped-query attention, dK and dV) tile via global atomic adds can complete in any order. Floating-point addition is not associative, so different orderings produce slightly different results. For applications requiring bit-identical reproducibility — such as debugging, reinforcement learning where policy gradient estimates must be consistent, or scientific computing — this nondeterminism is unacceptable.

Lock-based serialization. The deterministic mode uses a semaphore lock to serialize the global reductions. For each CTA that needs to write to a common tile:

1. The CTA must acquire a lock (a semaphore counter in global memory) for that tile.
2. CTAs acquire the lock in a predefined order (not the order they happen to finish).
3. The acquiring CTA reads the current value of the tile from global memory, performs its reduction (adding its contribution), writes the result back, and releases the lock by incrementing the semaphore counter.
4. A memory fence instruction ensures device-wide visibility of the semaphore write, so the next CTA sees the updated lock value.

Performance cost and mitigation. The lock-based approach has two sources of overhead: the memory fence instruction itself (which serializes memory operations across the device), and stalls when a CTA waits for a previous CTA to release a lock it needs. The paper addresses stalls through "shortest-processing-time-first" (SPT) scheduling. For causal masking, the KV blocks are launched in descending order (longest KV blocks first), query blocks are traversed in ascending order starting from the causal diagonal, and dQ reductions are ordered by descending query block index. This ensures that "no CTA is stalled on its first dQ write" because the CTA that produces results earliest (shortest computation) finishes its atomic section before CTAs with longer computations need the same lock.

For non-causal attention, the paper applies swizzling over the head and batch dimensions to reduce stalls — CTAs operating on different heads or batches do not contend for the same dQ tiles, so interleaving their execution reduces the probability that two CTAs targeting the same tile are co-scheduled and cause one to stall.

Performance impact. The deterministic backward pass achieves "up to 75% the speed of the nondeterministic backward pass of the 1-CTA backward pass" (Figure 7), and Figure 8 shows the non-causal deterministic variant with similar behavior. The 25% overhead for determinism is low enough to be practical for training, and the careful scheduling (LPT + reverse mblock order vs. naive) makes a substantial difference — the paper's Figure 7 shows that the SPT/LPT scheduling variants significantly outperform the naive approach with no batch/head swizzle.

---

3.4.8 Scheduling Policy for Load Imbalance

Attention kernels are naturally load-imbalanced in two situations: causal masking (where the upper triangle of the attention matrix is masked out, so blocks near the beginning of the sequence have fewer valid KV positions to attend to than blocks near the end), and variable sequence lengths (varlen) in batched inference (where different sequences in a batch have different lengths, so some CTAs process much more work than others). The scheduling policy uses longest-processing-time-first (LPT) ordering to minimize makespan — the total time for the slowest SM to finish.

Causal masking LPT. The standard attention grid is organized as (mblocks, heads, batches), processed left-to-right (increasing KV position). For causal attention, this means SMs process worktiles from shortest to longest — the early blocks attend to few positions, the later blocks attend to many — which is optimal for cache locality but terrible for load balancing: SMs that finish their short-work tiles early sit idle while SMs processing long-work tiles continue running.

The paper's LPT scheduler instead processes mblocks (KV position tiles) in reverse order — longest first, shortest last. The outer loop is over batches, and within each batch, the scheduler swizzles over heads in sections that fit within the L2 cache. This means: divide heads into groups whose KV data fits in L2, process all query heads within each KV head group (for MQA/GQA), and within each group traverse mblocks in descending order (longest-processing-time-first). Empirically, this yields "4-8% FLOPS gain for MHA and 7-14% for MQA 8" on H200 GPUs — the LPT ordering is not Blackwell-specific and was validated as an improvement on Hopper hardware as well.

Variable sequence length LPT. For varlen (variable-length sequences), the batch dimension introduces load imbalance because different sequences have different lengths. The paper sorts batches in descending order of their maximum per-worktile execution time using a preprocessing kernel that writes a virtual-to-actual batch index mapping. The attention kernel then traverses batches in this sorted order. Because the metadata (sort mapping) is small and can be cached, the sorting itself introduces no performance loss — it is a one-time preprocessing cost amortized over the attention computation.

Grid organization and L2 cache management. The paper emphasizes that batches must be the outermost dimension of the grid traversal, and head swizzling must respect L2 cache capacity. If all KV heads were loaded first before varying over batches, the total KV head data might exceed L2 cache size, causing thrashing — KV tiles evicted from L2 would need to be reloaded from HBM for each batch. By swizzling over heads in sections that fit within L2, the scheduler ensures that KV head data loaded for one section benefits all batches before being evicted. This is a tradeoff between load balancing (which favors LPT ordering) and cache locality (which favors processing all batches for the same KV data before moving on), and the paper's section-based swizzling balances these concerns.

4. Key Insights and Innovations

Innovation 1: Asymmetric Hardware Scaling as a First-Class Design Concept

The paper's most distinctive intellectual contribution is not any single kernel optimization, but the explicit naming and systematic treatment of asymmetric hardware scaling as the central design constraint for modern GPU kernels. Prior to this work, the dominant mental model for GPU kernel optimization was that each new hardware generation simply provides more of everything — more FLOPS, more bandwidth, more of each functional unit — and a well-optimized kernel from the previous generation would see proportional speedup when ported. The FlashAttention lineage through version 3 largely operated under this assumption: the innovations were about exploiting new modes of execution (asynchrony in FA-3) rather than compensating for imbalances in how different units scaled.

This paper makes a diagnostic claim that is both simple and profound: the ratio of tensor core throughput to non-matmul throughput is increasing, and this trend is systematic, not accidental. The evidence is concrete — Blackwell's tensor core throughput doubled from Hopper (4096→8192 ops/clock/SM) while the MUFU exponential unit (16 ops/clock/SM), shared memory bandwidth (128 bytes/clock/SM), and other functional units remained unchanged. The authors characterize this as a structural trend in accelerator design: "increasing the throughput of the most important components (typically matrix multiply units) to get higher performance under similar power / silicon area constraint." This is not a bug or an oversight — it is a deliberate architectural choice that reflects the economics of chip design, and it means that the imbalance will persist or worsen in future generations, not self-correct.

What makes this a conceptual contribution rather than mere observation is the roofline methodology that operationalizes it. The paper's roofline tables (Tables 1 and 3) do not just note that certain units are slower — they compute the cycle counts for each resource under specific tile configurations and show where the crossover points are. For the forward pass with M=N=d=128, MMA compute and exponential unit both require 1024 cycles — balanced. For the larger tile M=256, N=d=128, both double to 2048 cycles — still balanced, but now shared memory traffic has grown to 1536 cycles, making it a co-bottleneck. For the backward pass with M=N=d=128, shared memory traffic (3328 cycles) exceeds MMA compute (2560 cycles) by 30%. These numbers are not hand-wavy approximations — they are derived from the hardware's published specifications and microbenchmarking measurements, and they directly predict which optimizations will matter and by how much.

This contrasts sharply with prior work in the FlashAttention lineage, where roofline analysis was used primarily to motivate the IO-awareness principle (avoiding HBM reads/writes) rather than to identify which on-chip resources were the bottleneck. FA-1 and FA-2 were about reducing global memory traffic; FA-3 added asynchrony and low-precision; FA-4 is the first to systematically identify that even within the SM, the balance between functional units has shifted and must be explicitly managed. The paper's contribution is not that asymmetric scaling exists — hardware architects have known this — but that it has become the dominant factor determining attention kernel performance, and that kernel designers must treat it as a first-class design constraint rather than an afterthought.

The significance of this framing extends beyond the specific B200 optimizations. It provides a conceptual vocabulary — "asymmetric hardware scaling," "bottleneck shift," "non-matmul resource" — that can structure the analysis of any future hardware generation. If the next architecture doubles tensor core throughput again while leaving shared memory unchanged, the analysis framework (compute cycle counts for each resource under candidate tile sizes, identify the bottleneck, design algorithmic mitigations) applies directly. The paper's title itself — "Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling" — signals that asymmetric scaling is the problem statement, not a parenthetical observation.

Innovation 2: Algorithmic Mitigation of Non-Matmul Bottlenecks

The conventional approach to a hardware bottleneck is to optimize the kernel's use of the bottlenecked resource — better scheduling, more overlap, reduced contention. FlashAttention-4 takes a more radical approach: it changes which functional unit executes a given mathematical operation, effectively increasing the throughput of the bottlenecked resource by routing some of its work to underutilized alternative units. This is algorithm-hardware co-design in a stronger sense than the prior literature — it modifies the algorithm (what mathematical operations are performed and with what precision) to match the hardware's actual resource balance, rather than only modifying the schedule (when operations execute).

Exponential emulation is the clearest example. The conventional approach to the exponential bottleneck would be to hide its latency through better overlap — compute exponentials in parallel with other work, use larger tiles to amortize the cost, or accept that it limits throughput. The paper instead asks: is there another functional unit on the chip that can compute exponentials, even approximately? The FMA units — which execute fused multiply-add instructions — are typically underutilized during the softmax phase (they are used for the rescaling operations, but those are a small fraction of total FMA capacity). By implementing a polynomial approximation of 2^x using FMAs, the kernel can effectively increase the SM's exponential throughput beyond the MUFU's fixed 16 ops/clock.

What distinguishes this from a mere engineering trick is the precision analysis that justifies it. The paper explicitly measures the accuracy of polynomial approximations of varying degrees against the hardware MUFU.EX2 instruction at both FP32 and BF16 precision (Table 2). The finding that a cheap degree-3 polynomial (8.8 × 10^(-5) max FP32 relative error) becomes indistinguishable from hardware after BF16 rounding — because the BF16 quantization error of ~3.9 × 10^(-3) dominates — is the intellectual move that makes the optimization viable. Without this analysis, one might reasonably object that replacing hardware exponentials with polynomial approximations introduces unacceptable numerical error. The paper demonstrates that the error is quantized away by the output precision, making the approximation provably lossless at the precision the downstream computation uses. This transforms the technique from a heuristic speed-accuracy tradeoff into a principled optimization with formal justification.

The conditional softmax rescaling (Section 3.1.4) is a different kind of algorithmic mitigation: it reduces the number of non-matmul operations rather than increasing their throughput. The standard online softmax algorithm unconditionally rescales the running output every time a new maximum is found — a vector multiply that consumes FMA bandwidth and contributes to the non-matmul bottleneck. The paper observes that rescaling can be deferred without affecting final correctness, because the final normalization step will correct any intermediate discrepancies. This is a modification to the attention algorithm itself — the update rule for the running statistics — not just to its implementation. The threshold τ = log_2(256) = 8.0 is chosen to keep intermediate values within BF16's dynamic range, so the correctness guarantee is not asymptotic but quantitative.

Together, these two techniques represent a philosophical shift in how kernel optimizers approach hardware constraints. Rather than treating the hardware's functional unit allocation as fixed and optimizing within it, the paper treats it as malleable — the kernel can decide, per-operation, which functional unit to use based on the current resource balance and precision requirements. This is a form of functional unit load balancing that operates at the algorithm level, not the instruction scheduler level. The partial emulation strategy (10–25% of exponentials on FMA, the rest on MUFU) is a concrete instance: the kernel dynamically allocates exponential computations across two functional units to maximize aggregate throughput, analogous to how a compute-optimal workload balancer would allocate work across heterogeneous processors.

Innovation 3: The 2-CTA Cooperative MMA Mode as a Reduction Restructuring Primitive

The 2-CTA backward pass reconfiguration (Section 3.2.3) is the paper's most architecturally novel contribution — it takes a new hardware feature (Blackwell's 2-CTA cooperative MMA mode) and uses it to solve a problem (shared memory bandwidth and atomic adds) for which the feature was not originally designed. Blackwell's 2-CTA mode was introduced to support larger MMA tile sizes (M=256) by allowing two CTAs to cooperatively feed a single tensor core operation, partitioning operand B across their shared memories. The paper's insight is that this same mechanism, combined with distributed shared memory (DSMEM) for data exchange between CTAs, can restructure the reduction axis of the dQ gradient computation to eliminate a fundamental conflict between the MMA tile partitioning and the global reduction pattern.

The conflict arises because the dQ gradient accumulates over the KV sequence length dimension (N), but the natural 2-CTA partitioning splits N across the CTA pair — meaning each CTA would only see half the reduction, producing an incorrect partial result. The conventional solution would be to avoid using 2-CTA mode for the dQ MMA entirely, reverting to 1-CTA mode and accepting the associated shared memory and atomics overhead. The paper instead recognizes that the conflict can be resolved by repacking the dS tensor via DSMEM before the dQ MMA executes, so that the reduction dimension is not split across CTAs.

This is a conceptual contribution because it reframes the 2-CTA feature from a simple throughput multiplier (bigger tile = more FLOPS per instruction) into a reduction restructuring primitive that can change the communication pattern of distributed reductions. The paper demonstrates that, with appropriate data repacking and pipeline reordering, the 2-CTA mode can simultaneously reduce shared memory traffic (each CTA stages half of operand B), halve the number of global atomic adds (each CTA owns half the M-dimension of the output), and maintain correctness of the reduction (the full N-dimension is available to each CTA via DSMEM exchange). These three benefits — reduced SMEM bandwidth, reduced atomics, and correct reduction — are achieved simultaneously through a single architectural mechanism, which is an elegant example of hardware-algorithm co-design.

The significance is amplified by the fact that global atomic adds are a pervasive source of nondeterminism and performance overhead in attention kernels. The dQ gradient accumulation requires atomic updates to global memory because multiple CTAs (processing different KV tiles) contribute to the same dQ output. These atomics are expensive (L2 cache latency, not just SMEM), nondeterministic (floating-point addition is non-associative, so the order of atomics affects the result), and scale with sequence length. Halving the atomics count is a direct improvement on all three fronts — performance, determinism, and scalability. The deterministic backward pass (Section 3.2.4) builds on this reduction, using a semaphore lock to serialize the remaining atomics with only ~25% overhead, enabled in part by the atomics count already being halved.

This technique is not incremental — it represents a material advance over FlashAttention-3's backward pass, which was effectively serialized due to register pressure and had no mechanism for reducing atomics. It is also not purely engineering because it requires reasoning about the mathematical structure of the reduction (which dimension is being reduced, how the partitioning interacts with it) rather than just scheduling existing operations more efficiently.

Innovation 4: Hardware-Accelerator DSL Implementation as a Strategic Enabler

While the CuTe-DSL implementation (Section 4) could be dismissed as "just" an engineering choice, the paper makes a substantive claim that deserves recognition as an innovation: embedding GPU kernel development in a Python-based DSL with JIT compilation changes the economics of kernel optimization in ways that enable the other innovations in the paper. The 20–30× compile time reduction over C++ template-based approaches (Table 4) is not merely a developer convenience — it changes what kinds of design exploration are feasible.

The paper's techniques — particularly the exponential emulation fraction and the pipeline ordering for the backward pass — require extensive empirical tuning. The fraction of exponentials to emulate (10–25%) is "tuned empirically based on the ratio of MMA and exponential throughput for a given tile configuration." The pipeline reordering for DSMEM latency hiding requires iterating on different schedules to find one that achieves maximal overlap. Each tuning iteration requires recompiling the kernel. In a C++ template-based workflow (like FlashAttention-3's), each compilation might take minutes — turning a tuning session of 50 iterations into hours. In the CuTe-DSL workflow, the same iterations take seconds, making tight feedback loops practical.

This connects to a broader claim the paper makes about accessibility and ecosystem. The traditional approach to GPU kernel optimization — writing performance-critical code in C++ templates with inline PTX assembly — has an exceptionally high barrier to entry. The paper explicitly states that CuTe-DSL "enables researchers and engineers with just a few months of GPU programming experience to contribute meaningful extensions without requiring deep expertise in C++ template metaprogramming." If this claim holds, it means the FlashAttention-4 framework functions not just as a high-performance kernel but as a platform for attention variant research — a composable set of primitives (block-sparse patterns, masking strategies, variable-length handling, work scheduling) that can be combined to produce new attention variants with near-peak performance without requiring each variant to be implemented from scratch.

This is a distinctive positioning relative to prior work. FlashAttention-3 was a C++ template library that was difficult to extend; the paper's evidence is that developers successfully built FlexAttention and block-sparse variants on top of FA-4 "without modifying the core framework." The claim that the framework enables composition across orthogonal primitives — rather than requiring each new attention variant to be a separate kernel — is a software architecture contribution that has more in common with compiler design than with traditional kernel optimization. The vision of "independent, composable primitives" that "compile down to efficient GPU kernels" represents a shift from optimizing individual attention implementations toward building an optimizing compiler for attention variants, where the compiler (CuTe-DSL + JIT) can apply the paper's hardware-specific optimizations (exponential emulation, 2-CTA mode, pipeline scheduling) to any attention variant built on the framework.

The paper's statement that cuDNN has since incorporated many of the described techniques is indirect validation of this point: the techniques are now available to cuDNN users, but only for the attention variants that cuDNN explicitly supports. The CuTe-DSL framework offers the same techniques for arbitrary attention variants composed from the primitive library, which is a different value proposition — breadth of applicability rather than peak performance on a fixed set of variants.

5. Experimental Analysis

Evaluation Methodology

• Hardware and precision. All benchmarks are run on a B200 GPU (B100 180GB SXM6, 1000W) with BF16/FP16 inputs. The paper specifies CUDA 13.1, FlashAttention 2.8.3, Triton 3.6, PyTorch 2.10.0, CuTe-DSL 4.4.1, and cuDNN 9.13 (with additional comparisons to cuDNN 9.19.1.2 which incorporates some FA4 techniques) — establishing that the comparisons are against contemporary, Blackwell-aware baselines rather than stale implementations.

• Dataset. The paper evaluates on synthetic attention workloads parameterized by sequence length and head dimension, not on a fixed benchmark dataset like MATH. The "test set" is the set of (seqlen, head_dim, causal) configurations swept. This is standard for kernel benchmarking — the goal is to measure computational throughput on representative tensor shapes, not end-to-end model accuracy. The key parameters are: sequence length from 1K to 32K, total tokens fixed at 32K (by varying batch size inversely with sequence length), hidden dimension 2048, and head dimensions of 64, 128, and (192, 128) for the DeepSeek V3-inspired configuration.

• Metrics. The primary metric is TFLOPs/s (tera-floating-point-operations per second), computed from the mathematical operation count of exact attention divided by measured wall-clock runtime. For the forward pass, FLOPs = 4 × seqlen² × head_dim × num_heads; for causal masking, this is divided by 2 to account for the approximately half of entries that are masked. For the backward pass, FLOPs = forward_FLOPs × 2.5 (since there are 2 matmuls in forward and 5 in backward, due to recomputation). The paper also reports TFLOPs/s as a percentage of theoretical peak (2.25 PFLOPS = 2250 TFLOPs/s for BF16 on B200), yielding a utilization metric. Runtime is measured as wall-clock time averaged over 10 runs after 5 warmup iterations.

• Baselines. The paper compares against five baselines, spanning from vendor-optimized to high-level frameworks:

  • cuDNN 9.13 (and 9.19.1.2 in Appendix A): NVIDIA's vendor-optimized deep learning library, representing the best closed-source implementation available. The paper notes that cuDNN 9.19.1.2 incorporates some FA4 techniques after collaboration between the teams.
  • Triton 3.6 (Tillet et al., 2019): A high-level GPU programming language and compiler that includes B200-specific instructions. This represents the best open-source, non-CuTe-based implementation.
  • Gluon (Triton Team, 2024): A lower-level GPU programming language with finer control than Triton, included as an intermediate baseline between Triton's high-level abstractions and CuTe's full low-level control.
  • FlashAttention-2.8.3 (Dao, 2023): The Hopper-optimized predecessor. FlashAttention-3 is noted not to run on B200 ("simply does not run... due to lack of forward compatibility for Hopper MMA instructions"), so FA2 serves as the FlashAttention lineage baseline.
  • PyTorch 2.10.0: The standard PyTorch attention implementation, included primarily as a reference point for the magnitude of the optimization gap (which is well-established from prior work).

• Generation budget / compute accounting. The paper measures raw kernel execution time, not "generations" — there is no sampling or search budget. The comparison metric is throughput (TFLOPs/s) on identically-shaped tensor operations. Each baseline executes the same mathematical operation (exact attention forward or backward) on the same tensor shapes. The paper does not need to account for "difficulty estimation cost" or "generation budget" — the comparison is purely about how efficiently each implementation maps the same mathematical operation onto the B200 hardware.

• Cross-validation / statistical protocol. For runtime measurements, the paper uses 5 warmup runs followed by 10 timed runs, reporting the average. For the deterministic backward pass scheduling ablation, multiple scheduling strategies (naive, LPT, LPT with reverse mblock order, SPT) are compared directly on identical workloads. For the polynomial exponential accuracy evaluation (Table 2), 4 million random inputs in [0,1) are tested against a double-precision (FP64) reference. For the roofline analysis, the hardware throughput numbers (8192 ops/clock/SM for MMA, 128 bytes/clock/SM for SMEM, 16 ops/clock/SM for MUFU) come from published NVIDIA specifications and independent microbenchmarking results (Luo et al., 2025), not from the paper's own measurements.

Main Quantitative Results

Forward Pass Throughput

The headline result for the forward pass is that FlashAttention-4 achieves 1.1–1.3× speedup over cuDNN 9.13 and 2.1–2.7× over Triton across sequence lengths from 1K to 32K with head dimension 128 (Figure 4, both non-causal and causal panels). The paper also reports an absolute throughput of up to 1613 TFLOPs/s, corresponding to approximately 71% of the B200's theoretical peak of 2.25 PFLOPS for BF16.

Breaking this down by configuration:

• Non-causal, head dimension 128 (Figure 4, left): FlashAttention-4 reaches approximately 1600 TFLOPs/s at sequence length 8K and above, while cuDNN 9.13 achieves roughly 1200–1300 TFLOPs/s in the same range. The gap is relatively stable across medium and long sequences (4K–32K), suggesting that FA4's advantage comes from architectural choices (pipeline design, TMEM usage) rather than sequence-length-dependent effects. Triton achieves approximately 600–750 TFLOPs/s, or roughly 2.1–2.7× slower than FA4. At shorter sequences (1K–2K), all implementations show lower absolute throughput due to kernel launch overhead dominating the shorter compute time; the relative gaps are somewhat compressed.

• Causal, head dimension 128 (Figure 4, right): The relative gains are larger in the causal case, which the paper attributes to the LPT scheduler. FlashAttention-4 reaches roughly 1400–1500 TFLOPs/s at 8K and above, while cuDNN 9.13 achieves approximately 1000–1100 TFLOPs/s — a larger relative gap than in the non-causal case. Triton similarly shows a larger gap (roughly 500–600 TFLOPs/s, 2.3–2.7× slower). The LPT scheduler's effect is most pronounced for causal attention because the natural grid ordering (left-to-right over KV positions) creates severe load imbalance that LPT directly addresses.

• Causal, head dimension (192, 128) (Figure 5): This DeepSeek V3-inspired configuration (16 heads, 192 query dimensions, 128 key/value dimensions) is evaluated only against cuDNN, showing that FA4 achieves marginally higher throughput across the sequence length sweep. The paper reports this in a separate figure (Figure 5) with the y-axis unlabeled in the provided image, making precise TFLOPs numbers difficult to extract, but the qualitative pattern matches the head-dimension-128 results.

• Forward pass with head dimension 64: The paper states it evaluates head dimension 64 but the corresponding figures are not included in the main text. The appendix (not shown in the provided paper content) likely contains these results.

A critical contextual note appears in the Figure 4 caption: "Since the initial release of our implementation, newer versions of cuDNN have incorporated many of the techniques described in this paper, yielding similar performance to FA4." The paper includes cuDNN 9.19.1.2 results (Appendix A) showing that the gap has largely closed in the latest vendor library, which the authors present as validation that their techniques are correct — the vendor independently converged on similar approaches. This means the reported 1.3× speedup is a snapshot in time against cuDNN 9.13 specifically, not a permanent advantage over all future cuDNN versions.

Backward Pass Throughput

The backward pass results (Figure 6, both non-causal and causal with head dimension 128) show FA4 achieving consistent throughput advantages across long sequence lengths. The paper highlights that these results "demonstrate the effectiveness of our 2-CTA backward pass."

• Non-causal, head dimension 128 (Figure 6, left): FlashAttention-4 achieves approximately 1400–1550 TFLOPs/s at sequence lengths 4K–32K, consistently outperforming all baselines. cuDNN 9.13 reaches roughly 1050–1150 TFLOPs/s, Triton reaches roughly 550–600 TFLOPs/s, and FlashAttention-2 achieves roughly 650–750 TFLOPs/s. The gap between FA4 and the baselines is slightly larger than in the forward pass — consistent with the backward pass being more complex (five MMAs vs. two) and thus benefiting more from the 2-CTA mode's shared memory traffic reduction.

• Causal, head dimension 128 (Figure 6, right): Similar pattern: FA4 reaches roughly 1300–1400 TFLOPs/s at 4K–32K, with cuDNN at roughly 1000–1100 TFLOPs/s and Triton at roughly 500–600 TFLOPs/s. FA2's performance drops noticeably at longer sequence lengths in the causal case, likely due to load imbalance that FA4's LPT scheduler addresses.

The paper does not provide backward pass results for head dimension 64 or for the (192, 128) configuration in the main text.

Deterministic Backward Pass

The deterministic backward pass (Figure 7, causal attention, head dimension 128; Figure 8 in Appendix A includes non-causal) achieves up to 75% the speed of the nondeterministic backward pass of the 1-CTA backward pass. This is presented as an ablation of scheduling strategies within the deterministic kernel:

• SPT (shortest-processing-time-first): Achieves the highest throughput, approximately 1050 TFLOPs/s at 16K–32K sequence lengths.
• LPT with reverse mblock order: Slightly lower, roughly 1000 TFLOPs/s in the same range.
• LPT (standard, without reverse mblock ordering): Approximately 950 TFLOPs/s.
• Naive (no batch/head swizzle): The lowest, around 800–850 TFLOPs/s.

The gap between the best (SPT) and worst (naive) scheduling strategies is roughly 20–25%, demonstrating that careful scheduling of the lock-based serialization is essential to achieving practical deterministic performance. The paper's SPT strategy — launching KV blocks in descending order, traversing query blocks in ascending order from the diagonal, and ordering dQ reductions by descending query block index — is specifically designed so that "no CTA is stalled on its first dQ write," which minimizes the time any CTA spends waiting for a semaphore lock held by another CTA.

For the non-causal deterministic backward pass (Appendix A, Figure 8), the paper shows a comparison between "batch/head swizzle" and "naive" (no swizzle), with swizzling providing a clear throughput advantage (roughly 100–150 TFLOPs/s higher across sequence lengths) by reducing contention — CTAs operating on different heads or batches don't contend for the same dQ tiles, so interleaving them reduces stall probability.

LPT Scheduling Gains on Hopper (Not Blackwell)

The paper reports that the LPT scheduling strategy (Section 3.3) was validated on Hopper GPUs as well: "for BF16 and head dimension 128 we obtain 4-8% FLOPS gain for MHA and 7-14% for MQA 8 as measured on an H200 GPU." This is significant because it demonstrates that the scheduling improvement is not Blackwell-specific — it benefits prior-generation hardware — and that the gain is larger for multi-query attention (MQA) where the head-level load imbalance is more pronounced. However, these numbers are reported inline in the text (Section 3.3) without a corresponding figure, so the experimental protocol (sequence lengths, batch sizes, comparison baseline) is not specified.

Ablation Studies and Robustness Checks

• Polynomial degree vs. accuracy for exponential emulation (Table 2): The paper evaluates polynomial approximations of degrees 3–6 for the 2^x emulation on 4M random inputs, measuring maximum relative error at FP32 precision and after BF16 rounding. The degree-3 polynomial achieves 8.8×10^(-5) max FP32 relative error (600× worse than hardware's 1.52×10^(-7)) but becomes indistinguishable after BF16 rounding (3.91×10^(-3) for both, since BF16 quantization error dominates). Higher degrees close the FP32 gap (degree 5: 2.9×10^(-7), within 2× of hardware) but provide no BF16-level benefit. The paper selects degree 3 as optimal — sufficient accuracy at minimal instruction cost.

• Partial emulation fraction (10–25%): The exact fraction of exponentials emulated is tuned empirically per tile configuration. The paper does not report ablation curves showing sensitivity to this fraction — the reader cannot assess whether 10% vs. 25% makes a 1% or 10% difference in overall throughput. The stated range is qualitative guidance, not a rigorous sensitivity analysis.

• Conditional rescaling threshold τ = 8.0 (log_2(256)): The threshold for skipping softmax rescaling is set to 8.0, corresponding to a rescaling factor of 256. The paper does not report ablation results for different thresholds, sensitivity to this choice, or evidence that the specific value of 8.0 is optimal rather than simply conservative. The justification is that values scaled by up to 256× remain within BF16's dynamic range — a correctness argument rather than a performance-optimality argument.

• Deterministic backward pass scheduling strategies (Figures 7, 8): The paper compares SPT, LPT with reverse mblock order, LPT, and naive scheduling under causal attention, and batch/head swizzle vs. naive under non-causal attention. This is the most thorough ablation in the paper, showing monotonic improvement as scheduling sophistication increases. The SPT strategy (best) achieves roughly 25% higher throughput than naive (worst) at long sequence lengths.

• 1-CTA vs. 2-CTA backward pass: The paper's roofline analysis (Table 3) shows that 2-CTA mode reduces shared memory traffic overhead from exceed-MMA-by-30% to exceed-MMA-by-5%, but the paper does not present direct runtime comparisons of 1-CTA vs. 2-CTA backward pass kernels in isolation. The overall backward pass results (Figure 6) demonstrate the 2-CTA kernel's performance, but the marginal contribution of the 2-CTA mode specifically (vs. other backward-pass optimizations like the pipeline redesign) cannot be isolated from the presented data.

• LPT scheduler on Hopper (H200): The 4-8% (MHA) and 7-14% (MQA 8) FLOPS gains are reported without a figure, without specification of sequence lengths or batch sizes, and without indicating whether the baseline is FA3 with standard scheduling or some other kernel. This makes the result suggestive but not independently evaluable.

• Compile time comparison (Table 4): FlashAttention-3 (C++ templates) requires 60–120 minutes to precompile the full set of kernels for different attention variants and configurations, while FlashAttention-4 (CuTe-DSL) requires 3–4 minutes — a 20–30× reduction. This is not a runtime performance ablation but a developer productivity metric. The paper notes that FA2 and FA3 "require precompiling hundreds of kernels for different attention variants," while FA4's JIT compilation compiles only the kernel variants actually used.

• Numerical accuracy of the full attention computation: The paper does not report end-to-end numerical accuracy comparisons between FA4 and a reference implementation (e.g., PyTorch FP32 attention) for the full forward and backward passes. The exponential emulation accuracy is evaluated in isolation (Table 2), and the conditional rescaling is argued correct via the final normalization step. However, the accumulated effect of all algorithmic modifications (emulated exponentials, conditional rescaling, BF16 precision throughout, the 2-CTA reduction restructuring) on end-to-end attention output error or gradient error is not quantified. This is a significant omission for a paper that modifies the attention algorithm's numerical path — the claim that the emulation is "provably lossless at the precision the downstream computation uses" applies to the exponential operation in isolation but does not account for interaction effects across the full attention computation.

Critical Assessment

Does the Paper Demonstrate That Asymmetric Hardware Scaling Creates Shifting Bottlenecks?

The roofline analysis (Tables 1 and 3) makes a compelling analytical case: for specific tile configurations, the cycle counts for MMA, shared memory, and exponential operations show that shared memory and exponentials are co-bottlenecks with (or exceed) tensor core compute on Blackwell. The analysis is explicit about its simplifications — it "does not consider all resources in the GPU (e.g., floating point math, register bandwidth, L2 bandwidth)" — which is honest but means the analysis might miss secondary bottlenecks that emerge once the primary ones are addressed.

The experimental results support the inference that addressing these bottlenecks improves performance (FA4 outperforms baselines that don't target them), but they do not isolate the contribution of each bottleneck mitigation. The paper does not present an ablation that, for example, disables exponential emulation while keeping the new pipeline to measure how much throughput is lost specifically to the exponential bottleneck. Without such ablations, the causal chain from "identified bottleneck" to "technique addresses bottleneck" to "measured improvement" relies on the roofline analysis rather than direct experimental decomposition.

The claim that shared memory traffic "exceeds MMA compute time by approximately 30%" in the 1-CTA backward pass and by "approximately 5%" in the 2-CTA case (Table 3) is analytically derived, but the paper does not present measured SMEM bandwidth utilization to confirm that the kernel is actually saturating the 128 bytes/clock/SM limit. If the kernel is not saturating SMEM bandwidth — for example, due to bank conflicts or addressing patterns — then the analytical bottleneck calculation overstates the problem and the 2-CTA mode's benefit would be correspondingly overestimated.

Does FlashAttention-4 Actually Achieve 1.3× Over cuDNN and 2.7× Over Triton?

Yes, with important temporal and configuration qualifications. The forward pass results (Figure 4) clearly show FA4 above cuDNN 9.13 across all sequence lengths ≥ 4K for both causal and non-causal attention with head dimension 128. The specific multiplier varies by configuration — the 1.3× figure appears to be the upper end, with 1.1× being more typical — but the advantage is consistent.

The paper's own Figure 4 caption undercuts the permanence of this advantage: cuDNN 9.19.1.2 "yields similar performance to FA4" after incorporating FA4's techniques. This means the 1.3× number is a snapshot comparing FA4 to a previous cuDNN version, not a structural advantage over the best available implementation at publication time. This is not a weakness of the paper — the authors are transparent about it and frame the convergence as validation — but readers should not cite "1.3× faster than cuDNN" without the version qualifier.

The 2.7× over Triton is a more durable claim because Triton's higher-level abstractions inherently limit the degree of low-level control available. However, Triton's throughput varies substantially with sequence length (Figure 4 shows Triton at roughly 600 TFLOPs/s at 16K vs. FA4 at 1550 TFLOPs/s — closer to 2.6× — but the gap narrows at shorter sequences). The 2.7× figure appears to be the maximum across the sweep, not the typical value.

Missing: head dimension 64 results. The paper mentions evaluating head dimension 64 but does not include those figures in the main text. Head dimension 64 is extremely common in practice (many models use 64 or 128 heads with dimension 64), and the roofline analysis would predict different bottleneck characteristics at d=64 (SMEM traffic scales with d, so smaller d reduces SMEM pressure relative to MMA). Without these results, the paper's claims about bottleneck mitigation cannot be assessed for head dimension 64 — the exponential emulation and 2-CTA mode might be less impactful (or more impactful) at smaller head dimensions, and the reader has no data to evaluate this.

Is the 1613 TFLOPs/s (71% Utilization) Claim Meaningful?

Yes, as an efficiency metric for the hardware, but with caveats about what it measures. The 71% utilization number is computed as 1613 / 2250, where 2250 TFLOPs/s is the theoretical peak BF16 tensor core throughput of the B200. This is a standard way to report kernel efficiency — what fraction of the hardware's theoretical maximum compute throughput is the kernel actually achieving?

The caveat is that peak utilization is calculated against tensor core throughput specifically, but the paper's central claim is that tensor cores are not the bottleneck — shared memory and exponential units are. If the kernel is bottlenecked by SMEM bandwidth or exponential throughput, then 100% tensor core utilization is physically impossible regardless of kernel quality. The 71% number should therefore be interpreted as "the kernel achieves 71% of what would be possible if tensor cores were the only constraint," which is exactly what the paper's roofline analysis predicts — the other 29% is time spent waiting on SMEM and exponential units that cannot be overlapped, not inefficiency in the kernel's scheduling.

A more informative metric would be utilization against the actual bottleneck resource — what fraction of peak SMEM bandwidth is the kernel achieving? What fraction of MUFU capacity? The paper's roofline analysis provides the analytical framework for this but the experimental results do not include these hardware counter measurements. NVIDIA's profiling tools (Nsight Compute) can report SMEM bandwidth utilization and functional unit utilization, which would directly validate the bottleneck analysis. The paper does not report such measurements.

Do the Experiments Support the Claim That Exponential Emulation Increases Throughput?

Indirectly. The forward pass results (Figure 4) show that FA4 as a whole outperforms baselines, and the roofline analysis (Table 1) identifies the exponential unit as a co-bottleneck. The exponential emulation technique is analytically well-motivated and the accuracy analysis (Table 2) shows it is lossless at BF16 precision. However, the paper does not present an ablation comparing FA4 forward pass with and without exponential emulation. Without this, the reader cannot determine how much of the throughput gain comes from exponential emulation specifically versus the pipeline redesign, conditional rescaling, or other implementation details.

This is a significant methodological gap because the exponential emulation is presented as one of the paper's three named innovations. The technique is novel and analytically sound — replacing a subset of hardware exponentials with polynomial approximations on FMA units — but its quantitative contribution to the measured speedup is unverified.

Does the 2-CTA Backward Pass Deliver Its Claimed Benefits?

The backward pass results (Figure 6) show FA4 outperforming baselines, but the marginal benefit of the 2-CTA mode specifically is not experimentally isolated. The paper's analytical model (Table 3) predicts that the 2-CTA mode reduces SMEM traffic overhead from 30% to 5% of MMA time. This is a prediction, not a measurement. The actual backward pass kernel includes multiple simultaneous changes from FA3/FA2 — the redesigned pipeline, TMEM-based accumulator management, 2-CTA mode, and the LPT scheduler — and Figure 6 measures the combined effect.

The claim that 2-CTA mode "halves the number of global atomic adds" is a structural property of the algorithm — each CTA owns M/2 rows instead of M rows, so atomics per CTA are halved — but the performance impact of atomics reduction depends on how much the atomics were bottlenecking the kernel in the first place. If the kernel was already limited by SMEM bandwidth, reducing atomics might have minimal effect on throughput even though it's a genuine improvement in the algorithm.

The deterministic backward pass (Figure 7) is more informative because it isolates the scheduling strategy within a single kernel design. The 20–25% gap between naive and SPT scheduling demonstrates that the lock-based serialization overhead is real and that the SPT scheduling strategy effectively mitigates it. However, this evaluates the deterministic kernel's scheduling, not the 2-CTA mode's SMEM or atomics benefits.

Missing Experiments and Analyses

Several experiments would strengthen the paper's claims:

1. Hardware counter measurements (SMEM bandwidth utilization, MUFU utilization, tensor core utilization) for FA4 kernels. These would directly validate the roofline analysis by showing which resource is actually saturated at peak throughput. If the kernel achieves 71% tensor core utilization but only 60% SMEM bandwidth utilization, the bottleneck is something other than SMEM (perhaps instruction issue rate or register bandwidth), and the analytical model would need revision.

2. Ablation of exponential emulation: forward pass throughput with emulation enabled vs. disabled, at multiple sequence lengths. This would quantify the technique's contribution and reveal whether it matters more at some sequence lengths than others (the exponential bottleneck should be more severe at longer sequences where the M×N tile is larger).

3. Ablation of conditional softmax rescaling: forward pass throughput with rescaling always-on vs. conditional. This would quantify how many rescaling operations are actually skipped in practice and what throughput benefit results.

4. Head dimension 64 results in the main text. This is a gap in coverage rather than an ablation, but it matters for practical assessment since many production models use d=64.

5. End-to-end numerical error measurement: Compare the output of FA4's forward and backward passes against a high-precision reference (e.g., FP32 or FP64 PyTorch attention) to quantify the accumulated numerical error from all algorithmic modifications combined (emulated exponentials, conditional rescaling, BF16 throughout, 2-CTA reduction restructuring). The paper is careful about the exponential emulation's per-operation accuracy, but attention involves a chain of operations where errors can compound — softmax exponentiation, normalization, matrix multiplication with V, and the backward pass's gradient through the softmax. Without an end-to-end error measurement, the claim that the emulation is "sufficient for attention computation" is analytically supported but not experimentally verified.

6. Comparison to FlashAttention-3 on Hopper to isolate hardware-vs-algorithm contribution. Since FA3 does not run on Blackwell, this would require running both FA3 and FA4 on Hopper (where both can execute) to determine how much of FA4's improvement comes from better algorithms (exponential emulation, conditional rescaling) vs. Blackwell-specific features (TMEM, 2-CTA mode, larger tiles). This is partially addressed by the LPT scheduler results on H200 (4–8% gain), but the other techniques are not evaluated on Hopper.

7. Sensitivity analysis for the emulation fraction. The paper states that 10–25% of exponentials are emulated, tuned empirically. Showing throughput as a function of emulation fraction (e.g., 0%, 10%, 20%, 50%, 100%) would reveal how sensitive performance is to this parameter and whether the 10–25% range is a broad optimum or a narrow peak.

Conditional Nature of the Claims

The paper's claims are appropriately scoped to the hardware and configurations tested — B200 GPU, BF16 precision, head dimensions 128 and (192,128), sequence lengths 1K–32K. The paper does not claim that these techniques work on other hardware (consumer GPUs, prior-generation datacenter GPUs, non-NVIDIA accelerators), though it speculates that "some of these algorithms can be extended to other accelerators as compute continues to outpace non-matmul units."

The strongest claims (1.3× over cuDNN, 2.7× over Triton) are empirically supported for the tested configurations but are explicitly time-bound — cuDNN has since incorporated the techniques and achieved similar performance. The 71% utilization claim is analytically constrained by the roofline: it represents what is achievable given the non-matmul bottlenecks, not what would be achievable if those bottlenecks were eliminated.

The most undersupported claim is that the exponential emulation and conditional rescaling techniques are the drivers of the observed speedup rather than being incidental optimizations in a kernel whose performance is primarily determined by the pipeline redesign (which exploits Blackwell's asynchronous MMA and TMEM). The paper's narrative attributes the gains to the bottleneck-mitigation techniques, but the experimental design does not decompose the contribution of each technique. This is a common limitation in systems papers where the artifact is a complete kernel, not a modular collection of independently evaluable components, but it means the paper's architectural prescriptions ("identify non-matmul bottlenecks and mitigate them algorithmically") rest more on the analytical roofline than on experimental validation that the mitigations are causally responsible for the speedup.

6. Limitations and Trade-offs

Black-Box Throughput: No End-to-End Numerical Error Quantification

The paper introduces multiple techniques that modify the attention algorithm's numerical path — software-emulated exponentials using polynomial approximation (Section 3.1.3), conditional online softmax rescaling that defers normalization operations (Section 3.1.4), and a backward pass that restructures the gradient reduction using 2-CTA cooperative MMAs and DSMEM data exchange (Section 3.2.3). Each technique is analyzed in isolation for accuracy: the exponential emulation measures per-operation error against FP64 reference (Table 2), the conditional rescaling is justified by a correctness argument about final-step normalization, and the 2-CTA restructuring is structurally equivalent to the original gradient computation. However, no experiment measures the accumulated end-to-end numerical error of the complete forward and backward passes against a high-precision reference implementation.

This is consequential because attention computation involves a chain of operations where errors can compound. The softmax normalization involves exponentiation, subtraction by the row maximum, summation, and division — small per-element errors in the exponentials could interact with the row-max subtraction to produce larger relative errors in the normalized attention probabilities. The backward pass then computes gradients through this softmax, and non-associativity in the floating-point reductions (particularly the global atomic adds for dQ, even when halved by 2-CTA mode) means different reduction orderings produce different gradient values. The paper's claim that the degree-3 polynomial emulation is "sufficient for attention computation where the softmax output is consumed with BF16 precision" (Section 3.1.3) is analytically supported for the exponential operation in isolation, but the claim does not account for how that error propagates through the subsequent matrix multiplication with V and the backward pass gradients.

What evidence exists. Table 2 measures the exponential emulation's per-operation accuracy on 4M random inputs, showing the degree-3 polynomial matches hardware to within 1 BF16 ULP on 99.42% of inputs. Nothing measures the output of the full softmax(QK^T/sqrt(d))V against a reference, nor the gradient outputs dQ, dK, dV against reference gradients. The conditional rescaling's correctness is argued analytically ("at the end of the computation, all accumulated values are renormalized by the true maximum... This maintains the correctness"), not experimentally verified. The paper does not report whether models trained with FA4's backward pass converge identically to models trained with a reference implementation.

Mitigation status. The paper does not acknowledge this gap. The accuracy analysis is scoped to the exponential function alone, and the end-to-end numerical behavior is left as an open question. For a paper whose techniques modify the attention algorithm's numerical path — not just its schedule — this is a notable omission. A researcher deciding whether to adopt FA4 for training would need to independently verify that the numerical differences do not affect training dynamics or final model quality.

---

The B200-Specific Roofline: No Evidence of Generality to Other Hardware

The paper's entire analytical framework — the roofline analysis in Section 3.1.1 and 3.2.1 that identifies shared memory and exponential units as the shifted bottlenecks — uses Blackwell-specific hardware constants: tensor core throughput of 8192 ops/clock/SM, MUFU throughput of 16 ops/clock/SM, shared memory bandwidth of 128 bytes/clock/SM. These numbers are derived from the B200 and GB200 specifications and microbenchmarking. The three main techniques (exponential emulation, conditional rescaling, 2-CTA backward pass) are motivated by the specific bottleneck ratios these numbers produce for the tile configurations studied (M=N=d=128 and M=256, N=d=128).

The paper acknowledges in its conclusion that "some of these algorithms can be extended to other accelerators as compute continues to outpace non-matmul units," but this is speculative — no experiments are conducted on any hardware other than the B200, not even on other Blackwell-family GPUs with different SM counts or memory configurations, and not on Hopper GPUs except for the LPT scheduler (which is a scheduling technique, not a bottleneck-mitigation technique). The LPT scheduler results on H200 (4–8% for MHA, 7–14% for MQA 8, reported in Section 3.3 without a figure or experimental protocol) do not validate the exponential emulation, conditional rescaling, or 2-CTA mode on prior-generation hardware.

The consequence is that the paper does not establish whether asymmetric hardware scaling is the actual driver of FA4's improvements, or whether those improvements would appear on any modern GPU regardless of bottleneck ratios. Consider two alternative hypotheses: (1) The gains come primarily from the pipeline redesign that exploits Blackwell's fully-asynchronous MMA and TMEM — features that would improve throughput even if the exponential unit and SMEM bandwidth had scaled proportionally with tensor cores. (2) The exponential emulation and conditional rescaling provide throughput benefits that are independent of the specific bottleneck ratio — polynomial exponentials on FMA units might outperform MUFU even on hardware where the MUFU is proportionally faster, because FMA units have higher aggregate throughput in any architecture where tensor cores don't saturate them.

Without experiments varying the bottleneck ratios (e.g., by running on a GPU with different MUFU throughput or different SMEM bandwidth, or by artificially throttling resources in simulation), the paper cannot distinguish between these hypotheses. The roofline analysis identifies what should be the bottleneck — but the experimental results show only that FA4 as a whole is faster than baselines, not that addressing the identified bottlenecks is the reason it is faster. The cuDNN team's independent incorporation of similar techniques (noted in the Figure 4 caption) is weak evidence of generality — cuDNN targets the same B200 hardware, so convergence could reflect shared hardware constraints rather than universal principles.

What evidence exists. All benchmarks in Section 5 are on a B200 GPU (specified as B100 180GB SXM6 in Appendix A.1). The roofline numbers are Blackwell-specific. The LPT scheduler results on H200 are reported inline without a figure, making the protocol and baseline unclear. The paper does not run FA4 on an H100 or H200 to show which techniques transfer and which do not.

Mitigation status. The paper acknowledges the hardware specificity in the conclusion but treats it as future work rather than a limitation to be addressed. The statement that "some of these algorithms can be extended" is a forward-looking claim that the paper does not attempt to validate. For a practitioner deciding whether to invest in adopting FA4's techniques for non-Blackwell hardware, the paper provides no empirical guidance.

---

Decomposability Gap: Individual Technique Contributions Are Not Experimentally Isolated

The paper presents FlashAttention-4 as a complete kernel that incorporates multiple simultaneous innovations: a redesigned software pipeline exploiting asynchronous MMA and TMEM (Sections 3.1.2 and 3.2.2), software-emulated exponential functions (Section 3.1.3), conditional softmax rescaling (Section 3.1.4), the 2-CTA backward pass with DSMEM data exchange and halved atomics (Section 3.2.3), deterministic execution with lock-based serialization (Section 3.2.4), and LPT scheduling for load imbalance (Section 3.3). The paper's narrative attributes the performance gains to the bottleneck-mitigation techniques — the exponential emulation and conditional rescaling address the exponential bottleneck, the 2-CTA mode addresses the SMEM bottleneck — but no ablation experiments isolate the contribution of any single technique.

The experiments in Section 5 compare FA4 against external baselines (cuDNN, Triton, FA2, PyTorch), showing that the complete FA4 kernel is faster. They do not compare FA4 with and without exponential emulation, with and without conditional rescaling, with 1-CTA vs. 2-CTA backward pass, or with standard vs. LPT scheduling on B200. The only technique-level ablation is the deterministic backward pass scheduling comparison (Figure 7), which varies the scheduling strategy within the deterministic kernel — but this ablates a scheduling choice, not a bottleneck-mitigation technique.

The consequence is that the paper's central architectural prescription — "identify and mitigate non-matmul bottlenecks algorithmically" — rests on the analytical roofline rather than on experimental evidence that mitigating those bottlenecks is causally responsible for the measured speedup. A reader cannot determine whether the exponential emulation contributes 2% or 20% of the total gain, whether the conditional rescaling matters at all in practice (how many rescaling operations are actually skipped?), or whether the 2-CTA mode's SMEM traffic reduction translates to measurable throughput improvement.

This gap matters for several reasons. First, it affects reproducibility and adoption: a practitioner porting FA4's ideas to a different hardware target or a different attention variant needs to know which techniques are essential and which are incidental. Second, it affects scientific understanding: the paper's contribution is framed as demonstrating that asymmetric hardware scaling creates shifting bottlenecks that require algorithmic mitigation, but without isolating the mitigations, the paper demonstrates only that a well-engineered Blackwell kernel is faster than prior kernels — which could be true for many reasons unrelated to asymmetric scaling. Third, the cuDNN team has since incorporated "many of the techniques" (per the Figure 4 caption), but it is unclear which subset of techniques was necessary to achieve similar performance — the techniques could be partially redundant, with cuDNN achieving the same throughput using only a subset of FA4's innovations.

What evidence exists. No FA4-vs-FA4 ablation experiments are presented. The forward pass results (Figures 4, 5) and backward pass results (Figure 6) compare FA4 against external baselines. The deterministic backward pass (Figures 7, 8) ablates scheduling strategies within the deterministic kernel. The LPT scheduler results on H200 (inline in Section 3.3) lack a figure and protocol specification.

Mitigation status. The paper does not acknowledge this as a limitation. The techniques are described in detail and the roofline analysis motivates them analytically, but the experimental design does not validate the causal link between specific bottleneck mitigations and measured throughput gains.

---

Head Dimension 64 Results Are Described but Not Shown

The paper states in Section 5 that the evaluation includes "head dimensions 64, 128, and (192, 128)," yet the main text figures show only head dimension 128 (Figures 4, 6, 7) and the DeepSeek V3 configuration (192, 128) compared against cuDNN only (Figure 5). No head dimension 64 results appear in the main paper, and the appendix (which is truncated in the provided content) may or may not contain them.

This is consequential because head dimension 64 is extremely common in practice — many production Transformer models use 64 or 128 attention heads with dimension 64 (e.g., the original Transformer configuration, many LLaMA variants). The roofline analysis in Section 3.1.1 expresses shared memory traffic as a function of d, with T_smem proportional to MNd. At d=64, SMEM traffic is halved relative to d=128 for the same M and N, which would change the bottleneck balance — the exponential unit and tensor cores might become relatively more important, and the 2-CTA mode's SMEM traffic reduction might be less impactful. The exponential emulation's benefit might change because the ratio of exponential operations (proportional to MN, independent of d) to SMEM traffic (proportional to MNd) is different at d=64.

Without head dimension 64 results, the paper's claims about bottleneck mitigation cannot be assessed for a configuration that represents a substantial fraction of real-world attention workloads. A practitioner using head dimension 64 cannot determine from the paper whether FA4's techniques provide similar speedups, or whether the bottleneck analysis would identify different limiting resources at that head dimension.

What evidence exists. The paper mentions head dimension 64 as part of the evaluation scope (Section 5, "Benchmark settings") but the corresponding figures are absent from the main text. The appendix reference in Appendix A mentions "Additional Details on Experiments and Benchmarking" but the provided paper content ends before any d=64 figures appear. It is possible these results exist in the full appendix; within the provided content, they are missing.

Mitigation status. The paper does not flag this as a gap. The roofline analysis (Tables 1 and 3) assumes M=N=d=128 or M=256, N=d=128 — the d=64 case is not analytically modeled. If the d=64 results exist but are omitted from the main text due to space, this is a presentation choice rather than a technical limitation, but it still leaves the reader unable to evaluate a practically important configuration.

---

FLOPs-Based Throughput May Overstate Gains When Tensor Cores Are Not the Bottleneck

The paper's primary metric is TFLOPs/s, computed from the mathematical FLOP count of exact attention divided by measured wall-clock time. This is standard for kernel benchmarking and directly measures computational throughput. However, the paper's central claim is that tensor cores are not the bottleneck — shared memory bandwidth and exponential unit throughput are. Measuring performance in tensor-core-centric units (TFLOPs/s) when the limiting resource is something other than tensor cores creates a subtle reporting distortion.

Specifically, the paper reports that FA4 achieves "up to 1613 TFLOPs/s (71% utilization)" where utilization is 1613/2250 of the B200's theoretical peak BF16 tensor core throughput. But if the kernel is bottlenecked by shared memory bandwidth, then 100% tensor core utilization is physically impossible — the tensor cores must idle some fraction of the time waiting for operands to arrive from shared memory. A kernel achieving 71% tensor core utilization might simultaneously be achieving, say, 95% of peak shared memory bandwidth utilization, which would represent a higher efficiency against the actual bottleneck resource.

The 1.3× speedup over cuDNN, expressed as a ratio of TFLOPs/s, could partially reflect differences in how effectively each kernel saturates shared memory bandwidth rather than differences in tensor core utilization. If cuDNN achieves lower TFLOPs/s because it is more severely bottlenecked by shared memory (e.g., due to less aggressive pipelining of SMEM reads), then the speedup is real but its interpretation changes — FA4 is faster because it mitigates the SMEM bottleneck, not because it uses tensor cores more efficiently. The TFLOPs/s metric conflates these mechanisms.

What evidence exists. The paper's roofline analysis (Tables 1 and 3) provides the analytical framework to express performance in terms of the bottleneck resource, but the experimental results do not include hardware counter measurements (SMEM bandwidth utilization, MUFU utilization, tensor core utilization) that would validate which resource is actually saturated. The 71% utilization number is reported but not decomposed by resource.

Mitigation status. This is partially a presentational issue — TFLOPs/s is the standard metric for kernel benchmarking, and the paper's roofline analysis provides the conceptual framework for understanding that not all of the "missing" 29% is kernel inefficiency. However, the paper does not report alternative metrics (e.g., effective SMEM bandwidth achieved, fraction of peak MUFU throughput, or time-in-MMA vs. time-in-softmax breakdowns from a profiler) that would give a more complete picture of where the kernel spends its time. A reader relying solely on the TFLOPs/s numbers might misinterpret the 71% figure as indicating room for further optimization of tensor core usage, when the roofline analysis suggests the remaining headroom is in shared memory or exponential throughput.

---

No Training Convergence or Model Quality Validation

The paper evaluates FlashAttention-4 as a standalone kernel, measuring throughput on synthetic attention workloads. No experiment trains a real Transformer model using FA4 and compares convergence behavior, final model quality, or training stability against a reference attention implementation. This is standard for kernel benchmarking papers — the original FlashAttention paper also reported primarily kernel throughput, with training experiments appearing in later work — but it leaves open the question of whether the algorithmic modifications in FA4 affect the optimization trajectory.

The backward pass introduces specific numerical differences from a reference implementation: (1) the degree-3 polynomial exponential approximation in the softmax produces slightly different attention probabilities (within 1 BF16 ULP on 99.42% of elements, but the remaining 0.58% could differ by more); (2) the conditional rescaling skips intermediate normalization steps, accumulating values at slightly different scales until the final correction; (3) the 2-CTA mode restructures the dQ gradient reduction, changing the order of floating-point additions, which are non-associative; (4) the nondeterministic backward pass (the default) produces different gradient values on each run due to atomic add ordering. Even the deterministic backward pass, which serializes atomics via semaphore locks, may produce different numerical values than a reference implementation due to the 2-CTA reduction restructuring.

These differences are individually small, but their interaction with the stochastic gradient descent optimization process is unknown. Training is a chaotic dynamical system where small per-step gradient perturbations can either cancel out (no effect on final quality) or amplify (different local minima, different generalization). The paper's accuracy analysis (Table 2) shows BF16-level equivalence for the exponential emulation in isolation, but this does not guarantee training-equivalence of the full system.

What evidence exists. The paper provides no training experiments. Section 5 evaluates kernel throughput on synthetic tensors. The exponential accuracy table (Table 2) measures per-operation error, not training-relevant metrics.

Mitigation status. The paper does not acknowledge this as a limitation or suggest it as future work. This is understandable for a systems paper focused on kernel performance, but it is a gap for practitioners who need to decide whether to adopt FA4 for training. The cuDNN team's incorporation of similar techniques (per the Figure 4 caption) provides weak indirect evidence — if NVIDIA's vendor library includes these optimizations by default, they have presumably been validated on training workloads — but the paper itself does not provide this validation.

7. Implications and Future Directions

How This Work Changes the Landscape

FlashAttention-4 introduces an explicit and operationalized conceptual framework for asymmetric hardware scaling as a first-class design constraint in GPU kernel development. Before this paper, the dominant approach to optimizing attention for a new GPU generation was to exploit new modes of execution—asynchrony in FA-3, low-precision in the SageAttention series—while implicitly assuming that if you could keep the tensor cores fed, performance would follow. FlashAttention-4 demonstrates that this assumption has broken down: when tensor core throughput doubles while shared memory bandwidth and exponential unit throughput remain unchanged, the bottleneck shifts away from matrix multiplication toward the supporting functional units, and no amount of scheduling optimization within the existing algorithm can fully compensate.

The methodological shift this introduces is that kernel optimization must now include algorithmic modification of which functional unit executes which operation, not just how operations are scheduled. The exponential emulation technique (Section 3.1.3) is the signature example: rather than accepting the MUFU's fixed 16 ops/clock/SM as a hard ceiling on exponential throughput, the paper re-routes 10–25% of exponential evaluations to the FMA units—a different functional unit with different throughput characteristics—using a polynomial approximation that is provably lossless at BF16 precision. This is not scheduling (it does not change when exponentials execute), and it is not pure engineering (it does not simply optimize the existing MUFU code path). It is an algorithmic decision—what mathematical operation to compute, on which hardware unit, with what precision—made in response to a hardware resource imbalance.

The significance of this shift extends beyond attention kernels. Any computation that involves both matrix multiplications (which benefit from tensor core scaling) and elementwise operations (which do not) faces the same structural challenge on architectures where tensor cores out-scale other units. The paper provides a diagnostic methodology—the three-resource roofline analysis comparing MMA compute, shared memory traffic, and exponential unit throughput against hardware constants—that can be applied to any candidate algorithm on any hardware target. This is analogous to how the Chinchilla scaling laws (Hoffmann et al., 2022) provided a framework for reasoning about the pretraining compute tradeoff, but applied at the individual-kernel level and to the balance of on-chip functional units rather than the model-size-vs-data tradeoff.

The practical consequence is a reframing of what "optimized for hardware X" means. Before this paper, an attention kernel "optimized for Blackwell" would have meant primarily exploiting Blackwell's new features—TMEM, 2-CTA MMA mode, larger tiles. FlashAttention-4 argues that optimization for a new hardware generation must also include compensating for what did not scale—identifying which units are now the bottleneck because they failed to keep pace with tensor cores, and modifying the algorithm to reduce demand on those units. This is a more demanding standard, because it requires analyzing the full resource balance rather than only the headline features, but it is also what distinguishes a kernel that achieves 71% utilization from one that achieves 50%.

The paper also resolves a latent tension in the kernel optimization literature between hardware exploitation and hardware portability. The FlashAttention lineage has been criticized for being increasingly tied to specific GPU generations (FA-2 for Ampere, FA-3 for Hopper, FA-4 for Blackwell), with each new version requiring a near-rewrite. The paper's CuTe-DSL implementation strategy (Section 4) partially addresses this by embedding the kernel in a Python-based DSL that supports composable primitives for different attention variants—but the more fundamental contribution is the analytical framework that makes the rewrite systematic rather than ad hoc. The roofline analysis with hardware-specific constants tells you why the old kernel is suboptimal on the new hardware, and the bottleneck-mitigation strategies (offload to underutilized units, reduce operations on bottlenecked units, restructure reductions) provide a playbook for what to do about it. This makes the per-generation rewrite less of a black art and more of an engineering process—still hardware-specific, but now principled.

The paper also subtly shifts the burden of proof for kernel optimization claims. Prior work could claim a kernel was optimized by showing it achieved high tensor core utilization. FlashAttention-4 demonstrates that on modern hardware, high tensor core utilization is neither sufficient (the kernel might be bottlenecked elsewhere) nor necessary (the kernel might correctly trade off tensor core utilization for higher throughput on the actual bottleneck resource). The appropriate metric is end-to-end throughput relative to the bottleneck resource's peak, not relative to tensor core peak. The paper's roofline analysis provides the framework for this, but the experimental section does not fully operationalize it—a gap that future work should address by reporting hardware counter measurements alongside TFLOPs/s.

Which research directions become more attractive? Hardware-algorithm co-design for specific bottleneck ratios is now a clearly defined research area with a demonstrated methodology. Work that assumes uniform hardware scaling—that each generation provides proportionally more of everything—is implicitly challenged by this paper's evidence that scaling is structurally asymmetric. The cuDNN team's independent incorporation of FA4's techniques (noted in the Figure 4 caption) validates that the approach is correct, but it also means the low-hanging fruit for B200 attention optimization has been picked—future work should target either the next hardware generation (where the bottleneck ratios will be different), other architectures (AMD, Intel, custom accelerators) where the ratios differ, or other Transformer operations (FFN layers, embedding lookups) that may face similar asymmetric scaling challenges.

Follow-Up Research This Work Enables

Hardware counter validation of the roofline bottleneck predictions. The paper's roofline analysis (Tables 1 and 3) predicts that shared memory bandwidth and exponential unit throughput are the co-bottlenecks on Blackwell for the tile configurations studied, but the experimental results report only TFLOPs/s. A direct follow-up would use NVIDIA Nsight Compute or similar profiling tools to measure actual SMEM bandwidth utilization (bytes/clock vs. the 128 bytes/clock/SM peak), MUFU utilization (ops/clock vs. the 16 ops/clock/SM peak), and tensor core utilization (ops/clock vs. the 8192 ops/clock/SM peak) during FA4 kernel execution across the full sequence length sweep. This would either validate the analytical model (if the measured bottleneck resource is near saturation while others have headroom) or reveal that the actual bottleneck is something the model omitted—register file bandwidth, instruction issue rate, L2 cache latency—which would refine the methodology for future hardware generations.

End-to-end training convergence study with FA4 backward pass. The paper's accuracy analysis is scoped to the exponential emulation in isolation (Table 2), but the backward pass introduces multiple sources of numerical difference from a reference implementation: polynomial exponential approximation in the softmax gradient path, conditional rescaling that changes intermediate normalization, 2-CTA reduction restructuring that changes the order of floating-point additions, and nondeterministic atomic add ordering. A strong follow-up would train a language model (e.g., a 1B–7B parameter Transformer) from scratch on a standard corpus (C4, the Pile) using FA4's forward and backward passes, and compare training loss curves, downstream benchmark performance, and gradient statistics against an FP32 reference implementation. The key question is whether the per-step numerical differences (individually within 1 BF16 ULP for 99.42% of exponentials) accumulate or cancel over training. A negative result—meaningful divergence in training dynamics or final model quality—would not invalidate FA4 but would establish that the accuracy claims need to be evaluated at the training-trajectory level, not just the per-operation level.

Ablation study decomposing FA4's forward pass speedup into technique-level contributions. The paper attributes forward pass gains to the pipeline redesign, exponential emulation, and conditional softmax rescaling collectively, but no experiment isolates their individual contributions. A follow-up would implement configurable ablation flags in the FA4 forward kernel and measure throughput at representative sequence lengths (4K, 16K, 32K) with: (a) hardware exponentials only (emulation disabled), (b) unconditional rescaling (always rescale when max changes), (c) both disabled (pipeline improvements only), and (d) all enabled (current FA4). This would quantify whether the exponential emulation contributes, say, 5% or 20% of the total gain, and whether the conditional rescaling matters at all on realistic attention score distributions (where the maximum is typically found early). The results would guide practitioners on which techniques are essential to adopt versus incidental optimizations, and would validate or challenge the paper's central claim that non-matmul bottleneck mitigation is the primary driver of the speedup.

Head dimension 64 characterization and analytical model extension. The paper's roofline analysis and main results focus on head dimension 128, but d=64 is extremely common in production models. Shared memory traffic scales linearly with d (T_smem ∝ MNd), so at d=64 the bottleneck balance shifts—SMEM traffic is halved, potentially making the exponential unit or tensor cores proportionally more limiting. A follow-up would extend the analytical model to d=64, run the full FA4 benchmark suite at this head dimension, and compare whether the relative gains over cuDNN and Triton are larger, smaller, or similar to the d=128 case. If the gains are substantially different, it would demonstrate that the optimal set of techniques is head-dimension-dependent, which would motivate adaptive technique selection based on problem shape rather than a one-size-fits-all kernel configuration.

Application of the asymmetric scaling framework to other Transformer operations. Attention is not the only Transformer component where matrix multiplications dominate alongside elementwise operations. The feed-forward network (FFN) layers involve large matrix multiplications interleaved with activation functions (GELU, SiLU), which—like softmax—use specialized functional units that may not scale with tensor cores. Layer normalization and RMS normalization involve reductions and elementwise operations that compete for shared memory bandwidth. A follow-up study would apply the paper's three-resource roofline methodology to these operations on Blackwell hardware, identifying whether they face similar asymmetric scaling bottlenecks and whether analogous algorithmic mitigations (e.g., polynomial approximation of activation functions on FMA units, conditional rescaling for normalization statistics, tensor memory staging for FFN intermediates) yield throughput improvements. This would establish whether asymmetric hardware scaling is an attention-specific problem or a general challenge for all Transformer kernels on modern GPUs.

Porting and evaluation on non-NVIDIA accelerators with different resource balance points. The paper's techniques are motivated by Blackwell-specific bottleneck ratios, but the general principle—identify which unit failed to scale and offload work to underutilized alternatives—should apply to any architecture with heterogeneous functional units. A follow-up would implement FA4-style exponential emulation, conditional rescaling, and pipeline redesign on an AMD MI300X (which has different tensor core throughput, different shared memory bandwidth, and potentially different exponential unit characteristics) or an Intel Gaudi accelerator. The key measurement is whether the same techniques that improve throughput on B200 also improve throughput on hardware with different bottleneck ratios—if yes, the techniques may be universally beneficial regardless of the specific ratio (because FMA units are always underutilized during softmax, say); if no, the benefit is truly ratio-dependent, and the paper's roofline methodology is validated as the correct way to determine which techniques to apply on a given architecture.

Practical Applications and Downstream Use Cases

Training large language models on Blackwell clusters. Organizations deploying B200 or GB200 clusters for LLM pretraining are the most direct beneficiaries. Attention is the computational bottleneck for long-context training, and the backward pass in particular dominates runtime due to its higher arithmetic intensity (2.5× the FLOPs of forward). FA4's backward pass achieves up to 1.3× the throughput of cuDNN 9.13 at sequence lengths 4K–32K (Figure 6), and this multiplies across the thousands of GPUs in a typical training cluster. For a training run processing sequences of length 32K with head dimension 128, a 1.3× attention throughput improvement translates to a meaningful reduction in wall-clock training time—exact savings depend on the attention-to-FFN compute ratio for the specific model architecture, but for long-context models where attention dominates, the savings approach the full 1.3× factor. The deterministic backward pass (Section 3.2.4) additionally benefits reinforcement learning training pipelines (RLHF, GRPO) where gradient reproducibility is important for debugging and policy evaluation.

High-throughput inference serving with variable sequence lengths. Production inference systems (vLLM, TensorRT-LLM) handle mixed batches of requests with different sequence lengths—some short (simple questions), some long (document analysis). FA4's LPT scheduling for variable-length sequences (Section 3.3) sorts batches by processing time to minimize load imbalance across SMs, and the forward pass achieves up to 2.7× the throughput of Triton at 16K–32K sequence lengths (Figure 4). For inference serving systems where tail latency determines user experience, the LPT scheduler's ability to keep all SMs busy until the final request completes reduces the variance in per-batch processing time. The conditional softmax rescaling (Section 3.1.4) provides additional benefit in the common inference case where the KV cache is precomputed and the query is a single token (decode phase)—the running maximum is already known from the prefill phase, so rescaling operations are largely eliminated.

Custom attention variant development on the CuTe-DSL framework. The paper's CuTe-DSL implementation (Section 4) is positioned as a composable framework where new attention variants—block-sparse patterns, FlexAttention, multi-query, grouped-query—can be built by combining primitive operations without rewriting the core kernel. For research teams developing novel attention mechanisms (e.g., sliding window, dilated attention, learnable sparsity patterns), building on FA4's framework means inheriting the hardware-specific optimizations (exponential emulation, 2-CTA mode, LPT scheduling) automatically, without needing to re-engineer them for each new variant. The 20–30× compile time reduction (Table 4) enables rapid experimentation—a researcher can modify the attention pattern, recompile in minutes, and benchmark, iterating at a speed that would be prohibitively slow with C++ template-based approaches. The paper notes that developers have already built FlexAttention and block-sparse attention variants on top of FA4 without modifying the core framework, demonstrating this composability in practice.

Cost-efficient long-context model deployment. For applications requiring processing of very long contexts—analyzing entire codebases (inputs of 100K+ tokens), multi-document question answering, long-form video understanding—the quadratic scaling of attention means that attention kernel throughput directly determines whether the application is economically feasible. FA4's 1613 TFLOPs/s (71% utilization) on B200 GPUs at sequence length 16K (Figure 4) represents state-of-the-art efficiency for exact attention—for a service processing millions of long-context queries per day, the difference between FA4's throughput and a Triton-based implementation (roughly 2.5× slower) directly determines GPU-hour costs. The head dimension (192, 128) configuration tested in Figure 5 is specifically from DeepSeek V3, a production model, indicating that FA4's techniques target real-world model architectures, not just synthetic benchmarks.