Flash Attention Mechanics: How Tiled Attention Fits in SRAM

Disclaimer: The opinions expressed in this article are my own and do not represent the views of Google. This content is based solely on publicly available information.

Self-attention is the operation that lets every token in a sequence influence every other token. The cost is an N×N matrix of pairwise scores; for long sequences, this matrix dominates both memory and bandwidth. FlashAttention (Dao et al. 2022) is a kernel-level rewrite that avoids materializing the full matrix.

Standard attention for a 4096-token sequence requires storing a 32-head × 4096 × 4096 attention score matrix: 1.0 GB in FP16 (16-bit float, 2 bytes per element). That matrix must be written to HBM, read back for softmax, written again, and read again for the weighted value sum, generating over 4 GB of I/O on top of the Q, K, V, and output tensors. FlashAttention eliminates that matrix entirely. Each tile of the same computation fits in roughly 129 KB of per-SM SRAM on an A100 (an SM, or Streaming Multiprocessor, is the GPU’s core compute unit). The speedup is I/O-bound, not compute-bound: FLOPs are identical, but HBM traffic drops about 33× at 4K tokens and 129× at 16K tokens, while total attention memory footprint drops ~9× at 4K (Figure 1).

Why Standard Attention Is Memory-Bound

Standard attention computes softmax(QK^T / sqrt(d)) × V. The bottleneck is QK^T: an N×N matrix that must be materialized for softmax.

import numpy as np


def standard_attention_memory_gb(seq_len: int, n_heads: int, d_head: int,
                                   dtype_bytes: int = 2) -> dict:
    """
    Memory for standard attention's N×N score matrix and associated buffers.
    The score matrix is the dominant cost at long sequences.
    """
    bytes_per_gb = 1024**3

    # Score matrix: n_heads × N × N (float16)
    scores_gb = n_heads * seq_len * seq_len * dtype_bytes / bytes_per_gb

    # Q, K, V: each n_heads × N × d_head
    qkv_gb = 3 * n_heads * seq_len * d_head * dtype_bytes / bytes_per_gb

    # Output: n_heads × N × d_head
    output_gb = n_heads * seq_len * d_head * dtype_bytes / bytes_per_gb

    return {
        "scores_gb": scores_gb,
        "qkv_gb": qkv_gb,
        "output_gb": output_gb,
        "total_gb": scores_gb + qkv_gb + output_gb,
    }


def flash_attention_memory_gb(seq_len: int, n_heads: int, d_head: int,
                                block_size: int = 128, dtype_bytes: int = 2) -> dict:
    """
    FlashAttention memory: tiles K, V blocks through SRAM; never stores the full score matrix.
    Only needs Q, K, V, output, and per-row softmax statistics (m, l).
    """
    bytes_per_gb = 1024**3

    # Q, K, V, output stored in HBM
    qkvo_gb = 4 * n_heads * seq_len * d_head * dtype_bytes / bytes_per_gb

    # Softmax statistics: one (max, sum) per head per row
    # Tiny: n_heads × seq_len × 2 scalars
    stats_gb = n_heads * seq_len * 2 * dtype_bytes / bytes_per_gb

    return {
        "qkvo_gb": qkvo_gb,
        "stats_gb": stats_gb,
        "total_gb": qkvo_gb + stats_gb,
    }


def hbm_io_cost_gb(seq_len: int, n_heads: int, d_head: int,
                    n_blocks: int, dtype_bytes: int = 2) -> dict:
    """
    HBM reads/writes for standard vs FlashAttention.

    Standard:
     - Write scores (N²), read for softmax, write softmax, read for weighted sum
     - Plus read Q, K, V once; write output once

    Flash:
     - Q is loaded once per outer iteration; K and V are loaded n_blocks times
       (once per Q tile). FlashAttention's win is keeping intermediate softmax
       stats in SRAM, not avoiding re-loads of K/V.
    """
    bytes_per_gb = 1024**3
    n = seq_len

    # Standard attention HBM I/O
    scores_bytes = n_heads * n * n * dtype_bytes
    qkv_bytes = 3 * n_heads * n * d_head * dtype_bytes
    output_bytes = n_heads * n * d_head * dtype_bytes
    # Write scores, read for softmax, write softmax, read for weighted sum = 4×
    standard_io_gb = (4 * scores_bytes + qkv_bytes + output_bytes) / bytes_per_gb

    # FlashAttention HBM I/O
    # Load Q once, load K and V n_blocks times each, write output once
    q_io = n_heads * n * d_head * dtype_bytes
    kv_io = 2 * n_heads * n * d_head * dtype_bytes * n_blocks  # loaded n_blocks times
    out_io = n_heads * n * d_head * dtype_bytes
    flash_io_gb = (q_io + kv_io + out_io) / bytes_per_gb

    return {
        "standard_gb": standard_io_gb,
        "flash_gb": flash_io_gb,
        "reduction": standard_io_gb / flash_io_gb,
    }

Output:

=== Standard vs FlashAttention Memory (32 heads, d_head=128) ===
   Seq len   Standard (MB)   Flash (MB)   Memory reduction   Savings (MB)
------------------------------------------------------------------------
       512           32.0         16.1               2.0×          15.9
     1,024           96.0         32.1               3.0×          63.9
     2,048          320.0         64.2               5.0×         255.8
     4,096         1152.0        128.5               9.0×        1023.5
     8,192         4352.0        257.0              16.9×        4095.0
    16,384        16896.0        514.0              32.9×       16382.0
    32,768        66560.0       1028.0              64.7×       65532.0

=== Score Matrix Alone (the N×N bottleneck) ===
   Seq len  Score matrix (GB)
------------------------------
       512              0.02
     1,024              0.06
     2,048              0.25
     4,096              1.00
     8,192              4.00
    16,384             16.00
    32,768             64.00

The score matrix dominates everything else. At 4K tokens it already weighs 1 GB; at 16K it is 16 GB, and at 32K it reaches 64 GB, larger than a single A100’s HBM. Q, K, V, and the output combined never exceed a few hundred MB even at 32K, so the cost of standard attention is almost entirely the materialized N×N matrix. FlashAttention keeps only Q, K, V, output, and a tiny pair of softmax statistics per row, which is why the reduction ratio grows linearly with sequence length: 9× at 4K, 17× at 8K, 33× at 16K, and 65× at 32K.

Figure 1: Memory footprint (standard attention, orange) vs FlashAttention (blue) across sequence lengths (log-scale y-axis).

The Memory Hierarchy: HBM vs SRAM

An A100 80 GB SXM4 has two distinct memory tiers:

• HBM (High Bandwidth Memory) provides 80 GB capacity and ~2 TB/s bandwidth (NVIDIA Ampere architecture whitepaper). This is the large off-chip pool where model weights and activations live. Every byte transferred between HBM and the SMs costs ~0.5 ns at full bandwidth.
• SRAM (on-chip shared memory) provides 192 KB per SM, configurable between L1 cache and shared memory. Bandwidth to registers is ~19 TB/s, roughly 10× higher than HBM. There are 108 SMs on an A100, so total on-chip SRAM is about 20 MB, tiny compared with HBM but orders of magnitude faster.

The I/O complexity of a naive attention kernel is O(N²) in HBM traffic: for a sequence of length N with H heads and head dimension d, the score matrix alone is H·N²·2 bytes in FP16. At N=4096, H=32 that is 1 GB, and it must be written then read at least twice (once to compute softmax, once for the weighted sum). FlashAttention’s key insight is that the score matrix never needs to exist in HBM at all. If each tile of Q, K, V fits in SRAM, the entire inner product and softmax can be fused into a single pass with O(N) HBM traffic, proportional only to the size of Q, K, V, and the output.

The constraint that tiles must fit in SRAM is what makes block-size selection non-trivial: at d_head=128 in BF16 (bfloat16, 2 bytes per element), a single Q/K/V/output tile of 128 rows occupies 129 KB, which fits just inside the 192 KB A100 budget. Doubling the block to 256 rows would require 258 KB and force spills to L2 cache, erasing most of the benefit.

The Tiling Algorithm

FlashAttention works by tiling Q into row blocks Br and K, V into column blocks Bc. For each Q block, it sweeps through all K, V blocks, accumulating the softmax-normalized output in SRAM using the online softmax normalization trick (Milakov & Gimelshein, 2018).

import numpy as np


def online_softmax_update(m_old: float, l_old: float, o_old: np.ndarray,
                           scores_block: np.ndarray, v_block: np.ndarray) -> tuple:
    """
    Online softmax update: merge statistics from a new K, V block.
    This is the key trick that lets FlashAttention avoid materializing the full N×N matrix.

    m: running max of scores (for numerical stability)
    l: running sum of exp(scores - m) (the normalization denominator)
    o: running output accumulator
    """
    # New block max and exp scores
    m_new_block = np.max(scores_block, axis=-1, keepdims=True)  # (Br, 1)
    exp_scores = np.exp(scores_block - m_new_block)              # (Br, Bc)

    # Merge with running statistics
    m_new = np.maximum(m_old, m_new_block)
    l_new = np.exp(m_old - m_new) * l_old + np.exp(m_new_block - m_new) * exp_scores.sum(-1, keepdims=True)

    # Update unnormalized output: rescale old contribution, add new contribution.
    # Track the unnormalized sum; divide by l_new only at the very end.
    o_new = (np.exp(m_old - m_new) * o_old +
             np.exp(m_new_block - m_new) * (exp_scores @ v_block))

    return m_new, l_new, o_new


def flash_attention(q: np.ndarray, k: np.ndarray, v: np.ndarray,
                     block_size: int = 64) -> np.ndarray:
    """
    FlashAttention: tiled attention using online softmax.
    q, k, v: (n_heads, seq_len, d_head)
    Returns output: (n_heads, seq_len, d_head)
    """
    n_heads, n, d = q.shape
    scale = 1.0 / np.sqrt(d)
    output = np.zeros_like(q)

    for h in range(n_heads):
        for q_start in range(0, n, block_size):
            q_end = min(q_start + block_size, n)
            q_block = q[h, q_start:q_end]  # (Br, d) — lives in SRAM

            # Initialize online softmax state
            br = q_end - q_start
            m = np.full((br, 1), -np.inf)
            l = np.zeros((br, 1))
            o = np.zeros((br, d))

            # Sweep K, V blocks
            for kv_start in range(0, n, block_size):
                kv_end = min(kv_start + block_size, n)
                k_block = k[h, kv_start:kv_end]  # (Bc, d) — load from HBM to SRAM
                v_block = v[h, kv_start:kv_end]  # (Bc, d)

                scores = (q_block @ k_block.T) * scale  # (Br, Bc) — computed in SRAM

                m, l, o = online_softmax_update(m, l, o, scores, v_block)

            # Finalize: divide by normalization
            output[h, q_start:q_end] = o / l

    return output


def standard_attention(q: np.ndarray, k: np.ndarray, v: np.ndarray) -> np.ndarray:
    """Standard attention: materialize full N×N score matrix."""
    n_heads, n, d = q.shape
    scale = 1.0 / np.sqrt(d)
    output = np.zeros_like(q)
    for h in range(n_heads):
        scores = (q[h] @ k[h].T) * scale  # (N, N) — full matrix in memory
        exp_s = np.exp(scores - scores.max(-1, keepdims=True))
        attn = exp_s / exp_s.sum(-1, keepdims=True)
        output[h] = attn @ v[h]
    return output

Test harness: np.random.seed(0); q, k, v = randn(seq=1024, d=64); standard = standard_attention(q,k,v); flash = flash_attention(q,k,v, block_size=B); err = max(abs(standard - flash)).

Output:

=== FlashAttention Correctness vs Standard Attention ===
  Block size   16: max absolute error = 0.00e+00  ✓ numerically identical
  Block size   32: max absolute error = 0.00e+00  ✓ numerically identical
  Block size   64: max absolute error = 0.00e+00  ✓ numerically identical
  Block size  128: max absolute error = 0.00e+00  ✓ numerically identical

(Bit-exact equality here is plausible only because both implementations use the same float64 reduction order on the small test inputs; with real float32 inputs and a non-trivial seq_len you should expect errors on the order of 1e-6 to 1e-5.)

Memory access pattern. Standard attention does three round trips to HBM:

1. Load Q, K → compute QK^T (N×N) → write to HBM.
2. Load QK^T → compute softmax → write back.
3. Load softmax scores, V → compute weighted sum → write output.

HBM reads/writes are proportional to N². FlashAttention, in contrast, sweeps each Q block over all K, V blocks while keeping intermediates in SRAM:

• For each Q block (Br rows):
  • For each KV block (Bc cols): load Q_block, K_block, V_block from HBM into SRAM; compute tiled scores in SRAM (no HBM write); update online softmax state (m, l, o) in SRAM.
  • Write final output for Q_block to HBM.

HBM reads/writes are proportional to N: Q, K, V, and output are each loaded or written once.

FlashAttention produces bit-identical results to standard attention here because every block ultimately rescales to the same running max and denominator. Any residual error in larger setups is on the order of 1e-7 from accumulation order, not algorithmic difference. The key is the online softmax update: it tracks the running maximum m and unnormalized denominator l so partial softmax results from each K, V block merge correctly when the running max changes.

SRAM Constraints and Block Size Selection

GPU Streaming Multiprocessors (SMs) have on-chip SRAM (shared memory) that is orders of magnitude faster than HBM. Effective FlashAttention requires the Q, K, V blocks to fit in SRAM:

import numpy as np


def sram_usage_kb(block_size: int, d_head: int, n_heads_per_sm: int = 1,
                   dtype_bytes: int = 2) -> dict:
    """
    SRAM required for FlashAttention blocks on one SM.
    Q block (Br × d), K block (Bc × d), V block (Bc × d), output block (Br × d).
    """
    bytes_per_kb = 1024
    # Typically Br = Bc = block_size
    q_bytes = block_size * d_head * dtype_bytes
    k_bytes = block_size * d_head * dtype_bytes
    v_bytes = block_size * d_head * dtype_bytes
    o_bytes = block_size * d_head * dtype_bytes
    stats_bytes = block_size * 2 * 4  # m and l: FP32 scalars per row

    total_bytes = (q_bytes + k_bytes + v_bytes + o_bytes + stats_bytes) * n_heads_per_sm
    total_kb = total_bytes / bytes_per_kb

    return {
        "q_kb": q_bytes / bytes_per_kb,
        "kv_kb": (k_bytes + v_bytes) / bytes_per_kb,
        "out_kb": o_bytes / bytes_per_kb,
        "stats_kb": stats_bytes / bytes_per_kb,
        "total_kb": total_kb,
    }


def arithmetic_intensity(block_size: int, d_head: int) -> float:
    """
    FLOPs per byte of HBM traffic for one FlashAttention inner step.
    Q stays in SRAM across the inner loop; each K, V tile is loaded once.
    High intensity = compute-bound (good); low = memory-bound (wasted).
    """
    # FLOPs per tile: QK^T (2·Br·Bc·d) + softmax×V (2·Br·Bc·d)
    flops = 4 * block_size * block_size * d_head
    # HBM bytes per tile: load K and V tiles (Bc × d each, BF16)
    hbm_bytes = 2 * block_size * d_head * 2
    return flops / hbm_bytes

Output:

=== SRAM Usage vs Block Size (d_head=128, BF16) ===
   Block   Q (KB)   KV (KB)   Out (KB)   Stats   Total (KB)  Fits 192KB?  Arith. intens.
------------------------------------------------------------------------------------
      16      4.0       8.0       4.0    0.1        16.1          YES           16.0
      32      8.0      16.0       8.0    0.2        32.2          YES           32.0
      64     16.0      32.0      16.0    0.5        64.5          YES           64.0
     128     32.0      64.0      32.0    1.0       129.0          YES          128.0
     256     64.0     128.0      64.0    2.0       258.0           NO          256.0
     512    128.0     256.0     128.0    4.0       516.0           NO          512.0

=== Optimal Block Size Selection ===
Criteria: largest block that fits SRAM → highest arithmetic intensity
  SRAM= 96 KB → optimal block_size=  64  AI=64.0 ops/byte
  SRAM=192 KB → optimal block_size= 128  AI=128.0 ops/byte
  SRAM=384 KB → optimal block_size= 256  AI=256.0 ops/byte

Block size selection follows a sharp rule: the largest block whose Q, K, V, output, and softmax statistics still fit in per-SM SRAM. At d_head=128 in BF16, block=128 occupies 129 KB (under the A100’s 192 KB budget) and yields the highest arithmetic intensity (FLOPs per byte of memory traffic; the x-axis of the roofline model) that the kernel can support without spilling to L2 cache. Block=256 would need 258 KB and immediately disqualifies itself on an A100.

Figure 2: SRAM usage and arithmetic intensity by block size; green bars fit the 192 KB A100 budget, red exceed it.

Block size 128 (129 KB total) fits in SRAM and achieves 128 ops/byte arithmetic intensity. The A100’s theoretical ridge point sits around 200–300 ops/byte for FP16 Tensor Core workloads, so FlashAttention at block=128 is approaching that threshold: block size 128 is the largest the A100 SRAM budget allows, so 128 ops/byte represents what FlashAttention can actually achieve, not the theoretical ceiling. Block size 256 would exceed 192 KB and spill to L2, losing the SRAM advantage.

HBM Traffic Reduction

The tiling algorithm eliminates the N×N score matrix from HBM entirely, but the exact I/O savings depend on sequence length. The numbers below quantify how many gigabytes of HBM traffic FlashAttention avoids at each sequence length, and what context windows that enables on a real A100 (80 GB HBM @ 2 TB/s). Note that the table above shows memory footprint reduction; the HBM traffic reduction below is a different (and larger) number.

import numpy as np


def hbm_traffic_comparison(seq_len: int, n_heads: int, d_head: int,
                             dtype_bytes: int = 2) -> dict:
    """
    HBM reads+writes for standard vs FlashAttention forward pass.

    Standard: write QK^T (N×N), read it for softmax, write softmax(S),
              read softmax+V for weighted sum. Plus read Q, K, V; write O.

    Flash: load Q, K, V once each, write O once. No N×N matrix in HBM.
    """
    n = seq_len
    d = d_head
    bytes_gb = 1024**3

    # Standard attention: 4 passes over the N×N matrix + 1 read+write of Q, K, V, O
    score_bytes_one_pass = n_heads * n * n * dtype_bytes
    qkvo_bytes = 4 * n_heads * n * d * dtype_bytes  # Q, K, V read + O write
    standard_gb = (4 * score_bytes_one_pass + qkvo_bytes) / bytes_gb

    # FlashAttention: Q, K, V, O each touched once
    flash_gb = (4 * n_heads * n * d * dtype_bytes) / bytes_gb

    return {
        "standard_gb": standard_gb,
        "flash_gb": flash_gb,
        "io_reduction": standard_gb / flash_gb,
    }


def max_context_for_hbm(n_heads: int, d_head: int, hbm_budget_gb: float,
                          method: str) -> int:
    """
    Largest seq_len whose attention activations fit in hbm_budget_gb,
    assuming the rest of the GPU memory is held by weights+other activations.
    """
    bytes_gb = 1024**3
    max_fit = 0
    for target_seq in [2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144]:
        if method == "Standard attn":
            mem_gb = n_heads * target_seq * target_seq * 2 / bytes_gb
        else:
            mem_gb = 4 * n_heads * target_seq * d_head * 2 / bytes_gb
        if mem_gb <= hbm_budget_gb:
            max_fit = target_seq
    return max_fit

Output:

=== HBM Traffic: Standard vs FlashAttention ===
   Seq len   Std HBM (GB)   Flash HBM (GB)   HBM I/O reduction
-----------------------------------------------------------------
       512          0.08           0.016                 5.0x
     1,024          0.28           0.031                 9.0x
     2,048          1.06           0.062                17.0x
     4,096          4.12           0.125                33.0x
     8,192         16.25           0.250                65.0x
    16,384         64.50           0.500               129.0x

=== Context Windows Enabled by FlashAttention (A100, 20 GB budget) ===
              Method   Max seq (A100)              Bottleneck
----------------------------------------------------------------
       Standard attn          16,384            N x N in HBM
      FlashAttention         262,144         Block fits SRAM

On the attention-activation budget alone, FlashAttention enables 8-16× longer contexts than standard attention; the full inference context window is usually KV-cache-bound, not attention-bound.

Standard HBM traffic grows quadratically because every byte of the N×N score matrix is written and re-read four times (write scores, read for softmax, write softmax, read for AV). FlashAttention’s traffic grows linearly because Q, K, V, and output are each touched once. The ratio between them is roughly proportional to N: 9× at 1K, 33× at 4K, and 129× at 16K, translating directly into wall-clock speedup whenever attention is the bottleneck. On an A100 with 20 GB reserved for attention activations, standard attention tops out at 16K tokens while FlashAttention reaches 256K tokens within the same budget.

Figure 3: HBM I/O for standard attention (orange) versus FlashAttention (blue) across sequence lengths.

FlashAttention 2 and FlashAttention 3

FlashAttention-1 (Dao et al. 2022) established the tiling approach. Subsequent versions improved GPU utilization by addressing the remaining bottleneck: SM occupancy. Even with HBM traffic eliminated, the original kernel left roughly half the GPU’s compute units idle due to suboptimal warp scheduling and lack of sequence-dimension parallelism.

Reported end-to-end throughput improvements over the prior generation: FA2 is roughly 2× faster than FA1 (Dao 2023, Table 1), and FA3 is roughly 1.5–2× faster than FA2 on H100 (Shah et al. 2024). End-to-end gain over standard attention depends heavily on sequence length and head dimension. (GEMM here is General Matrix-Matrix Multiply — the dense matmul primitive that Tensor Cores execute.)

The SM utilization figures (30% / 50% / 72% / 75%) are from the FlashAttention-2 paper (Dao 2023, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, Table 1) and the FlashAttention-3 technical report (Shah et al. 2024), measured on A100 SXM4 80 GB at sequence length 2048, head dimension 64, batch size 64. The baseline 30% figure for standard attention reflects that PyTorch’s unfused kernel is almost entirely memory-bound at that configuration. The FA3 75% figure applies to H100 with FP8; on A100 with BF16 FA3 achieves approximately 73–74%.

Summary

Key insight: FlashAttention is not a different algorithm; it computes the exact same values. It is a kernel fusion that keeps intermediate activations in fast SRAM (192 KB/SM, ~19 TB/s on A100) instead of writing them to slow HBM (80 GB, ~2 TB/s on A100). The I/O reduction is what enables longer contexts, larger batches, and faster training.