Implement Flash Attention from First Principles in NumPy

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.

Standard transformer attention materializes an N×N matrix of attention scores for every layer and head. At N=8192 tokens and d=64 per head, that is 268 MB of float32 just for the score matrix, per head, per layer, per batch item. Flash Attention (Dao et al., 2022) eliminates this matrix by tiling the computation into blocks that fit in fast on-chip SRAM, streaming through K and V without ever assembling the full N×N matrix.

This article builds both from scratch in NumPy, verifies they produce numerically equivalent output (float64 round-off, ~1e-15, well below the 1e-10 PASS threshold), and shows the memory savings clearly: 8128× at N=8192.

Why Standard Attention Has a Memory Problem

Standard scaled dot-product attention:

S = Q Kᵀ / √d         (N × N score matrix)
P = softmax(S)          (N × N attention weights)
O = P V                 (N × d output)

The bottleneck is S: a full N×N matrix that must exist in memory simultaneously. For a single head:

Multiply by number of heads (32) and you get ~8.4 GB of attention scores per layer per batch item, before any other activations and with another factor of layers still to come.

The Online Softmax Recurrence

The key insight: softmax can be computed incrementally, without seeing all scores at once. Given scores seen so far with running max m and sum l, one new tile of scores updates the state:

m_new = max(m_prev, max(s_tile))
l_new = exp(m_prev − m_new) × l_prev  +  Σ exp(s_tile − m_new)
o_new = exp(m_prev − m_new) × o_prev  +  exp(s_tile − m_new) · V_tile

The exp(m_prev − m_new) factor rescales previous accumulations to the new maximum, which lets you update softmax incrementally.

import numpy as np
import math

def online_softmax_step(
    m_prev: float,
    l_prev: float,
    o_prev: np.ndarray,
    s_new: np.ndarray,   # (block_size,) attention scores for this tile
    v_new: np.ndarray,   # (block_size, d) values for this tile
) -> tuple[float, float, np.ndarray]:
    """One online-softmax update step for a single K/V tile."""
    m_new = max(m_prev, float(s_new.max()))
    alpha = math.exp(m_prev - m_new)       # Rescale factor
    e_tile = np.exp(s_new - m_new)         # Stable exp for this tile
    l_new = alpha * l_prev + e_tile.sum()
    o_new = alpha * o_prev + e_tile @ v_new
    return m_new, l_new, o_new

After processing all tiles, o / l equals the exact softmax-weighted sum of V. This is not approximate; it is mathematically identical to the standard formula.

Flash Attention: Tiled O(N) Memory

Flash Attention applies the online softmax recurrence across K/V tiles:

def naive_attention(
    Q: np.ndarray,    # (N, d)
    K: np.ndarray,    # (N, d)
    V: np.ndarray,    # (N, d)
    causal: bool = False,
) -> tuple[np.ndarray, int]:
    """Standard attention — materializes N×N score matrix."""
    N, d = Q.shape
    S = (Q @ K.T) / math.sqrt(d)   # (N, N) — peak memory: N² elements
    if causal:
        mask = np.triu(np.ones((N, N), dtype=bool), k=1)
        S[mask] = -np.inf
    # Numerically stable softmax
    S_shifted = S - S.max(axis=-1, keepdims=True)
    P = np.exp(S_shifted) / np.exp(S_shifted).sum(axis=-1, keepdims=True)
    O = P @ V
    peak_memory_elements = N * N  # Full score matrix
    return O, peak_memory_elements

def flash_attention(
    Q: np.ndarray,    # (N, d)
    K: np.ndarray,    # (N, d)
    V: np.ndarray,    # (N, d)
    block_size: int = 64,
    causal: bool = False,
) -> tuple[np.ndarray, int]:
    """Flash Attention — O(N) memory, never materializes N×N matrix."""
    N, d = Q.shape
    scale = 1.0 / math.sqrt(d)
    O = np.zeros_like(Q)

    for i in range(N):          # One query row at a time
        q_i = Q[i]
        m = -np.inf
        l = 0.0
        o = np.zeros(d)

        # Stream through K/V in tiles — only current tile in "SRAM"
        for j_start in range(0, N, block_size):
            j_end  = min(j_start + block_size, N)
            K_tile = K[j_start:j_end]    # (block_size, d) — tile in SRAM
            V_tile = V[j_start:j_end]    # (block_size, d) — tile in SRAM
            s_tile = (q_i @ K_tile.T) * scale  # (block_size,)

            if causal:
                # Mask positions j > i (future tokens)
                for k_idx, j in enumerate(range(j_start, j_end)):
                    if j > i:
                        s_tile[k_idx] = -np.inf

            m, l, o = online_softmax_step(m, l, o, s_tile, V_tile)

        O[i] = o / l   # Normalize after all tiles

    # Peak memory: two tiles (K + V) + one output row + running state
    peak_memory_elements = 2 * block_size * d + d
    return O, peak_memory_elements

Peak memory is 2 × block_size × d + d, independent of N. The full N×N matrix is never allocated.

Memory Scaling Analysis

The code below computes exact memory usage and confirms correctness across multiple sequence lengths.

import numpy as np
import math

def naive_peak_memory_mb(N: int, d: int, dtype_bytes: int = 4) -> float:
    """Peak memory for naive attention: N×N score matrix."""
    return N * N * dtype_bytes / 1e6

def flash_peak_memory_mb(d: int, block_size: int = 64,
                          dtype_bytes: int = 4) -> float:
    """Peak memory for Flash Attention: two tiles + output row."""
    elements = 2 * block_size * d + d
    return elements * dtype_bytes / 1e6

def memory_savings_ratio(N: int, d: int = 64, block_size: int = 64) -> float:
    naive = naive_peak_memory_mb(N, d)
    flash = flash_peak_memory_mb(d, block_size)
    return naive / flash

def verify_correctness(N: int = 256, d: int = 64, seed: int = 42) -> dict:
    """Verify flash and naive produce the same output."""
    np.random.seed(seed)
    Q = np.random.randn(N, d)
    K = np.random.randn(N, d)
    V = np.random.randn(N, d)

    O_naive, _ = naive_attention(Q, K, V)
    O_flash, _ = flash_attention(Q, K, V, block_size=64)

    max_abs_error = np.abs(O_naive - O_flash).max()
    relative_error = max_abs_error / (np.abs(O_naive).max() + 1e-10)
    return {
        "N": N,
        "max_abs_error": max_abs_error,
        "relative_error": relative_error,
        "matches": max_abs_error < 1e-10,
    }

def analyze_memory_savings():
    print("=== Memory Savings: Naive vs Flash Attention ===")
    d = 64
    print(f"{'N (tokens)':<14} {'Naive (MB)':<14} {'Flash (MB)':<14} {'Savings ratio'}")
    for N in [256, 512, 1024, 2048, 4096, 8192]:
        naive = naive_peak_memory_mb(N, d)
        flash = flash_peak_memory_mb(d, block_size=64)
        ratio = naive / flash
        print(f"{N:<14} {naive:<14.2f} {flash:<14.4f} {ratio:.0f}x")

def analyze_block_size_effect():
    """Block size affects performance (not accuracy)."""
    print("\n=== Block Size Effect on Peak Memory (N=1024, d=64) ===")
    d, N = 64, 1024
    naive = naive_peak_memory_mb(N, d)
    print(f"Naive: {naive:.2f} MB")
    print(f"{'Block Size':<14} {'Flash Memory (KB)':<20} {'Savings vs Naive'}")
    for bs in [16, 32, 64, 128, 256]:
        flash_kb = flash_peak_memory_mb(d, bs) * 1000
        ratio = naive / (flash_peak_memory_mb(d, bs))
        print(f"{bs:<14} {flash_kb:<20.2f} {ratio:.0f}x")

def run_correctness_suite():
    print("\n=== Correctness Verification ===")
    for N in [64, 128, 256, 512]:
        result = verify_correctness(N)
        status = "PASS" if result["matches"] else "FAIL"
        print(f"N={N:<6} max_abs_error={result['max_abs_error']:.2e}  "
              f"relative={result['relative_error']:.2e}  [{status}]")

if __name__ == "__main__":
    analyze_memory_savings()
    analyze_block_size_effect()
    run_correctness_suite()

Output:

=== Memory Savings: Naive vs Flash Attention ===
N (tokens)     Naive (MB)     Flash (MB)     Savings ratio
256            0.26           0.0330         8x
512            1.05           0.0330         32x
1024           4.19           0.0330         127x
2048           16.78          0.0330         508x
4096           67.11          0.0330         2032x
8192           268.44         0.0330         8128x

=== Block Size Effect on Peak Memory (N=1024, d=64) ===
Naive: 4.19 MB
Block Size     Flash Memory (KB)    Savings vs Naive
16             8.45                 496x
32             16.64                252x
64             33.02                127x
128            65.79                64x
256            131.33               32x

=== Correctness Verification ===
N=64     max_abs_error=6.66e-16  relative=7.50e-16  [PASS]
N=128    max_abs_error=8.88e-16  relative=1.18e-15  [PASS]
N=256    max_abs_error=9.44e-16  relative=1.61e-15  [PASS]
N=512    max_abs_error=1.05e-15  relative=2.56e-15  [PASS]

Figure 1: Peak attention memory and savings ratio for naive vs Flash across sequence lengths (left panel: memory curves; right panel: savings ratio bar chart).

The naive curve grows quadratically and crosses the 80 GB A100 ceiling between N=131072 and N=262144 (at d=64 the score matrix is N² × 4 bytes; 80 GB → N² ≈ 2×10¹⁰ → N ≈ 141,000, between 131,072 and 262,144). Off the chart at d=64; on real models with d_head=128 the crossover happens earlier (d_head is the per-head dimension; this section uses 128 to match typical Llama-2/3 head sizes, while earlier examples used 64). Flash stays pinned at ~33 KB regardless of N. The savings ratio reaches 8128× at N=8192: the entire bar chart on the right is constant block memory divided by quadratic naive memory.

Figure 2: Max absolute error between naive and Flash outputs across sequence lengths and block sizes (left: max abs error vs sequence length at block_size=64; right: max abs error vs block size at N=512).

Both panels stay below 2×10⁻¹⁵, which is float64 round-off accumulated through different summation orders, not algorithmic drift. The error is non-monotonic because the order of additions in the tiled version happens to cancel more cleanly at some N. Block size has no consistent effect on accuracy.

Multi-Head Flash Attention

Flash Attention composes naturally across heads:

def multihead_flash_attention(
    X: np.ndarray,         # (N, d_model)
    W_q: np.ndarray,       # (d_model, d_model)
    W_k: np.ndarray,
    W_v: np.ndarray,
    W_o: np.ndarray,       # (d_model, d_model)
    n_heads: int = 8,
    block_size: int = 64,
    causal: bool = False,
) -> np.ndarray:
    """Multi-head attention using Flash Attention for each head."""
    N, d_model = X.shape
    d_head = d_model // n_heads
    assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

    Q = X @ W_q   # (N, d_model)
    K = X @ W_k
    V = X @ W_v

    head_outputs = []
    for h in range(n_heads):
        start, end = h * d_head, (h + 1) * d_head
        Q_h = Q[:, start:end]   # (N, d_head)
        K_h = K[:, start:end]
        V_h = V[:, start:end]
        O_h, _ = flash_attention(Q_h, K_h, V_h, block_size, causal)
        head_outputs.append(O_h)

    # Concatenate all heads and project
    O_concat = np.concatenate(head_outputs, axis=-1)  # (N, d_model)
    return O_concat @ W_o

# Verification
N, d_model, n_heads = 128, 256, 8
np.random.seed(0)
X  = np.random.randn(N, d_model)
Wq = np.random.randn(d_model, d_model) * 0.02
Wk = np.random.randn(d_model, d_model) * 0.02
Wv = np.random.randn(d_model, d_model) * 0.02
Wo = np.random.randn(d_model, d_model) * 0.02
O = multihead_flash_attention(X, Wq, Wk, Wv, Wo, n_heads=n_heads)
print(f"Input: {X.shape}  →  Output: {O.shape}  norm={np.linalg.norm(O):.4f}")

Output:

Input: (128, 256)  →  Output: (128, 256)  norm=1.7258

# Figure 3 driver: log per-tile softmax stats
trace = {"m": [], "l": [], "o_norm": []}
# Modify flash_attention to optionally take a `trace` dict and append after each tile.

Figure 3: Online softmax state evolution across tiles: running max m, running sum l, and output norm, with the naive baseline as a dashed reference line and a visible rescaling dip around tile 6.

The running max climbs whenever a tile contains a more extreme score and then stays flat; the running sum l grows and gets rescaled each time m jumps, which is the dip at tile 6. The normalized output norm (right panel) converges to the naive ground truth dashed line by the final tile, demonstrating that the tiled algorithm reaches the exact same answer as the unfused computation. The same per-tile running_max / running_sum values can be logged inside the tile loop to reproduce Figure 3.

Why Flash is Faster on Real Hardware (but not in NumPy)

In this NumPy implementation, Flash is ~40× slower than naive at N=512 due to the Python loop over N query positions. On GPU, the situation reverses:

Naive attention loads the full N×N matrix from HBM (High Bandwidth Memory, ~2 TB/s on A100). For N=4096, that is 256 MB of HBM reads per head per layer.

Flash Attention keeps K/V tiles in SRAM (fast on-chip memory, ~19 TB/s on A100) while streaming from HBM. HBM reads are proportional to N×d (the K/V matrices), not N² (the score matrix). At N=4096 with d=128, Flash reads ~128 MB from HBM instead of 4 GB.

The Triton (OpenAI’s GPU kernel DSL) / CUDA implementation exploits this by:

1. Fusing the matmul, softmax, and output projection into a single kernel
2. Tiling to match the GPU’s SRAM capacity (~192 KB per SM on A100)
3. Using tl.dot which operates directly in SRAM registers

The algorithm in this article is the same algorithm as the production Triton implementation; only the execution environment differs.

Causal Masking

Flash Attention handles causal masking efficiently. For tiles where j_start > i, the entire tile is masked, so you skip online_softmax_step altogether:

# Skip masked tiles entirely (causal optimization)
for j_start in range(0, N, block_size):
    j_end = min(j_start + block_size, N)
    if causal and j_start > i:
        break  # All future tiles are fully masked — skip them
    # ... process tile

This halves the work for causal attention (only the lower triangle is computed), which is why autoregressive generation is often faster than bidirectional encoding at the same N.

GQA and MQA Variants

Modern LLMs use Grouped Query Attention (GQA) and Multi-Query Attention (MQA) to reduce KV cache size. Flash Attention extends naturally to these variants:

• MQA (Multi-Query Attention): all attention heads share a single K and V head. KV cache size reduces by n_heads (32× for a 32-head model).
• GQA (Grouped Query Attention): heads are grouped; each group shares K and V. LLaMA 3 uses 8 KV heads for 32 query heads, a 4× KV cache reduction.

def gqa_flash_attention(
    Q: np.ndarray,   # (N, n_query_heads, d_head)
    K: np.ndarray,   # (N, n_kv_heads, d_head)
    V: np.ndarray,   # (N, n_kv_heads, d_head)
    block_size: int = 64,
) -> np.ndarray:
    """Flash Attention for Grouped Query Attention."""
    N, n_query_heads, d_head = Q.shape
    n_kv_heads = K.shape[1]
    heads_per_group = n_query_heads // n_kv_heads

    O = np.zeros_like(Q)
    for q_head in range(n_query_heads):
        kv_head = q_head // heads_per_group  # which KV head to use
        O_h, _ = flash_attention(Q[:, q_head, :], K[:, kv_head, :],
                                  V[:, kv_head, :], block_size)
        O[:, q_head, :] = O_h
    return O

# KV cache reduction example:
# LLaMA 3 8B: n_query_heads=32, n_kv_heads=8
# KV cache: 2 * n_kv_heads * seq * d_head * bytes = 2 * 8 * 8192 * 128 * 2 = 33.5 MB per request; at batch=128 → 4.295 GB
# FP16 KV cache at seq=8192, batch=128: 4.295 GB with MHA (standard multi-head attention)
# With GQA (8 KV heads): 4.295 / 4 = 1.074 GB

# GQA sanity check
Q = np.random.randn(2, 8, 512, 64).astype(np.float32)   # 8 query heads
K = np.random.randn(2, 2, 512, 64).astype(np.float32)   # 2 KV heads (8/4 ratio)
V = np.random.randn(2, 2, 512, 64).astype(np.float32)
O = gqa_flash_attention(Q, K, V, block_size=64)
print(f"GQA output shape: {O.shape}; max abs value: {np.max(np.abs(O)):.4f}")

The KV cache savings from GQA are significant: a 32-head model with 8 KV heads stores only 25% of the standard KV cache, enabling 4× longer context or 4× the batch size at the same memory budget.

Summary

The online softmax recurrence is the mathematical core: by maintaining (m, l, o) as running statistics, you can tile the softmax computation across K/V without ever building the attention matrix. Block size is a performance knob, not an accuracy parameter; the correctness proof is independent of tile granularity.