---
title: 'FlashAttention-2: Efficient CUDA Kernels'
url: https://www.emergentmind.com/topics/flashattention-2-kernels
type: topic
---

# FlashAttention-2: Efficient CUDA Kernels

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 [2312.11918][2307.08691].

## 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 $K$-tiles, each query tile’s output is finalized as $O_i \leftarrow O_i / \Sigma_i$. This process is executed entirely in on-chip registers/shared memory, eliminating any intermediate global-memory writes of $S$ or $P$ [2312.11918][2307.08691].

## 2. Parallel Work Partitioning and Occupancy

FlashAttention-2 departs from previous implementations by introducing a multi-dimensional tiling and threading scheme:
- The input $Q$ is partitioned into $T_r = \lceil N/B_r \rceil$ row-blocks, each processed independently by a CUDA thread block (CTA).
- Likewise, $K$ and $V$ are partitioned into $T_c = \lceil N/B_c \rceil$ column-blocks per tile iteration.
- Within a CTA, work is distributed across multiple warps ($W=4$ or 8 typical), using a “split-Q” scheme where each warp handles a disjoint subset of $Q_i$'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 $T_r$ [2307.08691].

Comparison with FlashAttention-1 is summarized as follows:

|  Property                  | FlashAttention-1 | FlashAttention-2      |
|----------------------------|------------------|-----------------------|
| Thread block per head      | Yes              | No (per row-tile)     |
| TB count                   | $B \times H$     | $B \times H \times T_r$ |
| 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 $K_j$ and $V_j$ tiles ($[B_c \times d]$), allowing all warps in a CTA rapid access during each tile pass.
- $Q$-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 [2312.11918].
- CUTLASS Layout abstractions explicitly encode the shape and memory stride of GEMM operand and accumulator fragments, e.g.:
  ```cpp
  using CLayout_64x64 = Layout<
    Shape< Shape<4,8,4>, Shape<2,2,8> >,
    Stride< Stride<128,1,16>, Stride<64,8,512> >
  >;
  ```
- For GEMM-II (output accumulation), $V$ is logically treated as transposed via composition of SMEM layouts, eliminating runtime data movement [2312.11918].

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.,
  ```cpp
  cfk::copy(tQgQ(_,0), tQsQ(_,0), tmaLoadQ, tma_load_mbar[0]);
  ```
- 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: $K$ and $V$ tiles are prefetched asynchronously, hiding transfer latency behind computation [2312.11918].
- 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 $64 \times 128$ for head\_dim=64 and $128 \times 64$ 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. $128\times128$ at head\_dim=128) is directly attributed to excessive register spills [2312.11918].

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

FlashAttention-2 improves upon prior approaches through systematic reduction of non-GEMM computation:
\[
\text{FLOPs (forward)} = 2N^2 d + O(N^2)
\]
For attention, the two matmuls $(QK^\top$ and $PV)$ dominate, while softmax and reduction contribute the $O(N^2)$ term. By:
- Deferring the final $1/\ell$ normalization to the end, and,
- Only storing $L = m + \log \ell$ per row,
the non-matmul count is reduced by ~25% over FlashAttention-1 (“from $6N^2$ to $4N^2$”), increasing arithmetic utilization on tensor core hardware.

Global memory traffic per head is reduced from $O(N^2)$ to $O(N)$; only $Q$, $K$, $V$, $O$, and logsumexp ($L$) are stored globally, never the $S$ or $P$ matrices. Letting $\alpha$ be the ratio of non-matmul throughput to matmul (empirically $\alpha \approx 1/16$ for A100), this rebalancing of workloads is a key factor in increased hardware efficiency [2307.08691].

## 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 [2312.11918][2307.08691].

## 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 [2312.11918].

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 [2312.11918][2307.08691].

Source: https://www.emergentmind.com/topics/flashattention-2-kernels