FlashAttention-2: Efficient CUDA Kernels

• FlashAttention-2 is a highly optimized CUDA implementation of exact scaled dot-product attention that fuses GEMM operations with online softmax reduction.
• The technique employs multi-dimensional tiling and advanced memory staging via shared memory and CUTLASS layouts to achieve near-GEMM-level throughput.
• Empirical benchmarks show 20–50% speedups on A100 and H100 GPUs by reducing non-matmul overhead and global memory traffic.

FlashAttention-2 kernels are highly optimized CUDA implementations of memory-efficient, exact scaled dot-product attention for transformer models. Leveraging both algorithmic and microarchitectural advances, these kernels fuse matrix multiply-accumulate (GEMM) operations with online softmax reduction, maximize occupancy via parallel tiling, and exploit hardware primitives such as the NVIDIA Hopper Tensor Memory Accelerator (TMA) and Warpgroup Matrix-Multiply-Accumulate (WGMMA) instructions through the CUTLASS library. The approach achieves near-GEMM-level throughput and drastically reduces global memory traffic and non-matmul computational overhead, sustaining a substantial performance lead over prior state-of-the-art attention kernels on both A100 and H100 architectures (Bikshandi et al., 2023, Dao, 2023).

1. Mathematical Structure and Online Softmax Fusion

FlashAttention-2 implements per-head scaled dot-product attention: $$\mathrm{Attention}(Q, K, V) = \mathrm{Softmax}\!\Bigl(\tfrac{1}{\sqrt d}\,Q\,K^\top\Bigr)V$$ To avoid materializing the intermediate $S = QK^\top / \sqrt d$ matrix in GPU global memory, the algorithm tiles the $Q$, $K$, and $V$ matrices and fuses the entire attention computation with “online softmax”, a streaming reduction applied to each tile. For each query tile $Q_i \in \mathbb{R}^{b_M \times d}$ and over key tiles $K_j \in \mathbb{R}^{b_N \times d}$, the kernel maintains per-row running maxima $m_i$ and sum-exponentials $\Sigma_i$: $$\begin{aligned} S_{ij} &= (1/\sqrt d) Q_i K_j^\top \ m^{\mathrm{new}}_i &= \max\Bigl(m^{\mathrm{old}}_i, \max_\ell S_{ij}[\ell]\Bigr) \ \widetilde P_{ij} &= \exp\bigl(S_{ij} - m^{\mathrm{new}}_i\bigr) \ \Sigma^{\mathrm{new}}_i &= \exp(m^{\mathrm{old}}_i - m^{\mathrm{new}}_i)\Sigma^{\mathrm{old}}_i + \sum_\ell \widetilde P_{ij}[\ell] \ O_i &= \exp(m^{\mathrm{old}}_i - m^{\mathrm{new}}_i) O_i + \widetilde P_{ij} V_j \end{aligned}$$ After processing all $S = QK^\top / \sqrt d$0-tiles, each query tile’s output is finalized as $S = QK^\top / \sqrt d$1. This process is executed entirely in on-chip registers/shared memory, eliminating any intermediate global-memory writes of $S = QK^\top / \sqrt d$2 or $S = QK^\top / \sqrt d$3 (Bikshandi et al., 2023, Dao, 2023).

2. Parallel Work Partitioning and Occupancy

FlashAttention-2 departs from previous implementations by introducing a multi-dimensional tiling and threading scheme:

• The input $S = QK^\top / \sqrt d$4 is partitioned into $S = QK^\top / \sqrt d$5 row-blocks, each processed independently by a CUDA thread block (CTA).
• Likewise, $S = QK^\top / \sqrt d$6 and $S = QK^\top / \sqrt d$7 are partitioned into $S = QK^\top / \sqrt d$8 column-blocks per tile iteration.
• Within a CTA, work is distributed across multiple warps ($S = QK^\top / \sqrt d$9 or 8 typical), using a “split-Q” scheme where each warp handles a disjoint subset of $Q$0's rows.
• Empirically, this partitioning raises kernel occupancy, especially when total batch size or number of attention heads is small, as it multiplies the total thread-block count by $Q$1 (Dao, 2023).

Comparison with FlashAttention-1 is summarized as follows:

Property	FlashAttention-1	FlashAttention-2
Thread block per head	Yes	No (per row-tile)
TB count	$Q$2	$Q$3
Warp partition	Split-K	Split-Q
Intra-TB reduction passes	2 per tile	0

This approach reduces the requirement for shared memory reduction, increasing concurrent operations and raising arithmetic intensity.

3. Shared Memory, CUTLASS Layouts, and Data Movement

The kernel exploits architectural features and custom layouts for efficiency:

• Shared memory is used for staging $Q$4 and $Q$5 tiles ($Q$6), allowing all warps in a CTA rapid access during each tile pass.
• $Q$7-row fragments are loaded directly into registers for immediate matmul, never staged in shared memory.
• To avoid SMEM bank conflicts, swizzled layouts (e.g., K-major SW128) are used for Tensor Memory Accelerator (TMA) loads (Bikshandi et al., 2023).
• CUTLASS Layout abstractions explicitly encode the shape and memory stride of GEMM operand and accumulator fragments, e.g.: $Q_i \in \mathbb{R}^{b_M \times d}$3
• For GEMM-II (output accumulation), $Q$8 is logically treated as transposed via composition of SMEM layouts, eliminating runtime data movement (Bikshandi et al., 2023).

Within tiles, intra-warp communication via CUDA shuffles (__shfl_xor) enables register-local reductions for max/logsumexp, removing the need for global or atomic operations.

4. Hopper-Specific Kernel Fusion, Pipelining, and Tuning

On NVIDIA Hopper (SM90), several hardware primitives are tightly integrated:

• TMA (Tensor Memory Accelerator) is used for asynchronous, single-thread-initiated memory copies from global to shared memory, e.g., $Q_i \in \mathbb{R}^{b_M \times d}$4
• WGMMA (Warpgroup Matrix Multiply Accumulate) enables 128-thread warpgroup matmul. The two major GEMMs (Q-K and P-V) are distinct WGMMA calls, interleaved with softmax update logic in between.
• Coarse-grained software pipelining is achieved by overlapping GEMM computation with memory copy: $Q$9 and $K$0 tiles are prefetched asynchronously, hiding transfer latency behind computation (Bikshandi et al., 2023).
• Tile-size tuning is governed by register pressure and shared memory budget. Empirical experiments on H100 show, for FP16 input and FP32 accumulation, optimal throughput at $K$1 for head_dim=64 and $K$2 for 128, with large tiles penalized by register spills.

head_dim	64×64	64×128	128×64	128×128
64	230.1	259.5	247.9	251.4
128	292.6	289.3	295.7	208.7
256	308.1	276.1	39.3	36.7

Significant performance drop for larger tiles (e.g. $K$3 at head_dim=128) is directly attributed to excessive register spills (Bikshandi et al., 2023).

5. Complexity, Bottlenecks, and Reduction of Non-Matmul FLOPs

FlashAttention-2 improves upon prior approaches through systematic reduction of non-GEMM computation: $K$4 For attention, the two matmuls $K$5 and $K$6 dominate, while softmax and reduction contribute the $K$7 term. By:

• Deferring the final $K$8 normalization to the end, and,
• Only storing $K$9 per row, the non-matmul count is reduced by ~25% over FlashAttention-1 (“from $V$0 to $V$1”), increasing arithmetic utilization on tensor core hardware.

Global memory traffic per head is reduced from $V$2 to $V$3; only $V$4, $V$5, $V$6, $V$7, and logsumexp ($V$8) are stored globally, never the $V$9 or $Q_i \in \mathbb{R}^{b_M \times d}$0 matrices. Letting $Q_i \in \mathbb{R}^{b_M \times d}$1 be the ratio of non-matmul throughput to matmul (empirically $Q_i \in \mathbb{R}^{b_M \times d}$2 for A100), this rebalancing of workloads is a key factor in increased hardware efficiency (Dao, 2023).

6. Empirical Performance and Benchmarks

Benchmarking on a single H100 PCIe (CUTLASS 3.3 + CUDA 12.2, FP16 input / FP32 accumulation):

• For sequence length=4096, batch=4, head_dim=64/128/256:
  • H100 (COLFAX) achieves 259, 296, 308 TFLOPs/s for dims 64/128/256, respectively.
  • Speedup over previous (Ampere) FLASH-2 kernels: 20–50%.
  • Up to 3× throughput over default CUTLASS FMHA kernels on SM80.
  • On A100, forward-pass peak is 230 TFLOPs/s (≈73% of FP16 tensor core peak); backward pass peak 205 TFLOPs/s (≈65%).
  • End-to-end GPT training with FlashAttention-2 delivers up to 225 TFLOPs/s per A100 GPU, a 1.3× improvement vs. FlashAttention-1, and a 2.8× gain over conventional non-flash baselines (Bikshandi et al., 2023, Dao, 2023).

7. Implementation Practices and Future Directions

Key lessons for obtaining maximal performance:

• Mastery of CUTLASS/CuTe Layout and Tensor abstractions is fundamental, as kernel fusion is dominated by correct indexing and layout mapping.
• Swizzled SMEM layouts (SW128) for TMA mitigate bank conflicts; pre- and post-composing layouts allows efficient treatment of transpositions without actual memory copy.
• Max/sum reductions in the online softmax are best implemented by warp-shuffle, rather than atomics, enabling all computation to remain on-chip.
• Overlapping copy and compute using two back-to-back GEMMs provides effective pipelining in the presence of tight shared-memory constraints.
• Tile-size tuning should prioritize avoiding register spills, as indicated by NVCC occupancy reports.
• Extensions under active investigation include employing multiple warpgroups per CTA and deeper pipeline stages for scaling to even larger hardware (Bikshandi et al., 2023).

The comprehensive design of FlashAttention-2 kernels, from fused online softmax to pipeline-coupled GEMMs, PUSHes memory-bound attention layers toward the throughput ceiling established by pure matrix multiply, representing a fundamental advance in both transformer training and deployment efficiency (Bikshandi et al., 2023, Dao, 2023).