TL;DR — FlashAttention-2 is an IO-aware attention algorithm that fuses the softmax and matrix multiply into tiled CUDA kernels. We rebuild the forward pass from scratch in pure NumPy: blocking the QKV matrices into tiles, running a streaming softmax per block, and rescaling outputs online. The result is a runnable, testable side project that demonstrates memory-hierarchy thinking, numerical stability, and GPU-style programming patterns — exactly what hiring managers look for on a CV.

Most “implement attention from scratch” tutorials stop at a naive softmax(QKᵀ)V in ~30 lines of NumPy or PyTorch. That’s a college assignment. FlashAttention-2 is a different beast: it never materializes the full N×N attention matrix, streams softmax statistics tile-by-tile, and rescales partial outputs on the fly. Building it on a CPU in NumPy forces you to understand why those tricks matter for memory bandwidth — which is precisely the insight a hiring manager wants to see in a systems interview.

This post walks through a complete, runnable forward kernel: tiled QKV, online softmax with running max and sum, log-sum-exp accumulation, and an output rescaling trick that keeps numerical error bounded. By the end you’ll have a project you can benchmark, unit-test, and put on GitHub with a real README.

Why This Project Stands Out on a CV

When a recruiter or hiring manager scans a portfolio, they look for evidence that you can reason about a system end-to-end. A FlashAttention-2 clone in NumPy signals several specific skills that are rare at the junior-to-mid level:

  • Memory hierarchy awareness. Anyone can write softmax(Q @ K.T) @ V. Explaining why that’s catastrophic at sequence length 8k — and demonstrating a tiling scheme that cuts peak memory from O(N²) to O(N) — shows you understand L1/L3 caches, HBM bandwidth, and GPU SRAM. That’s exactly the vocabulary used in systems interviews at NVIDIA, Anthropic, and Mistral.
  • Numerical stability. Online softmax is the textbook example of a numerically stable streaming algorithm, alongside Welford’s online variance. Showing you can derive and implement the running max-and-sum trick puts you ahead of candidates who only know torch.softmax.
  • Reading and reproducing a paper. The FlashAttention-2 paper by Tri Dao is dense. Anyone who can build a working kernel from it — even in NumPy — has demonstrated paper-reading fluency, which is what research-engineer roles (Llama team, xAI, DeepSeek) actually screen for.
  • Tool fluency. A polished project will include NumPy, Pytest, Matplotlib for benchmarks, and optionally Triton or JAX for a follow-up GPU port. That mix reads as “production-aware” to anyone hiring for ML infra.
  • A story you can tell in interviews. “I built FlashAttention-2 in NumPy to teach myself the algorithm, then ported the hot loop to Triton and saw a 40× speedup” is a five-minute answer that beats “I fine-tuned a BERT on a Kaggle dataset.”

Roles this project resonates with: ML Performance Engineer, ML Systems Engineer, GPU/Kernel Engineer (NVIDIA, AMD, Tenstorrent), Research Engineer (Anthropic, OpenAI, xAI, Mistral, DeepSeek), Training Infra (Cohere, Together, Anyscale), and Inference Optimization (vLLM, SGLang, TensorRT-LLM teams). It also reads well for general backend roles that involve NumPy/Pandas-heavy code, because the same streaming-accumulator pattern shows up in finance, genomics, and time-series databases.

Architecture Overview

Before writing code, here’s the moving-parts picture. We have three input matrices Q, K, V each of shape (N, d) where N is sequence length and d is head dimension. The naive attention output is:

O = softmax(Q @ K.T / sqrt(d)) @ V

This materializes a full N×N score matrix in HBM. FlashAttention-2 instead loops over tiles and keeps two running statistics per row: a max m and a sum of exponentials l. Each tile updates these statistics and rescales the running output accumulator in place. No N×N matrix is ever created.

The components, top-down:

  • Input projections. Three (N, d) tensors Q, K, V. In a real model they’d come from nn.Linear projections, but for the kernel we treat them as given.
  • Tile parameters. Block size B_r (rows of the score matrix, i.e. rows of Q) and B_c (columns, i.e. rows of K/V). Typical values: B_r = B_c = 64 on GPU SRAM. On CPU NumPy we’ll use smaller values like 16 or 32 because of cache effects.
  • Outer loop over K/V blocks. For each tile, compute a partial score block S_ij = Q_i @ K_j.T / sqrt(d), shape (B_r, B_c).
  • Online softmax update. Compute new row max m_new = max(m_old, rowmax(S_ij)). Compute correction factor α = exp(m_old − m_new). Update running denominator l_new = α · l_old + rowsum(exp(S_ij − m_new)).
  • Output rescaling. Multiply accumulator by α, then add the new contribution: O_i ← α · O_i + exp(S_ij − m_new) @ V_j. Divide by l_new at the very end.
  • Final normalization. O = O / l per row.
  • Test harness. Compare against softmax(QKᵀ/√d)V from NumPy/SciPy within a tolerance. Plot memory usage as N grows — naive is O(N²), tiled is O(N·B_c).

Think of it as a streaming MapReduce over softmax: the map phase produces partial score tiles, the reduce phase folds them into a single output row using a numerically stable associative update rule. The same shape shows up in online algorithms for streaming quantiles and database query execution.

Building It Step by Step

We’ll build it in one self-contained file, flash_attn_numpy.py, then test it. The total implementation is around 90 lines.

Step 1: Project skeleton

Create a new repo with this layout:

flash-attn-numpy/
├── flash_attn_numpy.py
├── tests/
│   └── test_correctness.py
├── bench/
│   └── bench_memory.py
├── README.md
└── requirements.txt

requirements.txt:

numpy>=1.24
pytest>=7.0
matplotlib>=3.7

Step 2: Naive attention as the reference

We need a slow, obviously correct version to compare against.

import numpy as np

def naive_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray) -> np.ndarray:
    """Reference attention: O(N^2) memory, used as ground truth."""
    d = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d)              # (N, N)
    # numerically stable softmax
    scores = scores - scores.max(axis=-1, keepdims=True)
    P = np.exp(scores)
    P = P / P.sum(axis=-1, keepdims=True)
    return P @ V                                # (N, d)

This is what we will match within atol=1e-5 for d=64 and N up to a few hundred.

Step 3: Online softmax core

This is the trick that makes tiling work. Given a running max m_old and running sum of exponentials l_old, and a new tile of scores S_ij, we can update them without ever seeing all the data at once. The key invariant: l_new · exp(m_new) still equals rowsum(exp(all_scores_seen_so_far)).

def online_softmax_update(m_old, l_old, S_block):
    """Update running max and sum-of-exponentials with a new tile S_block."""
    # S_block shape: (B_r, B_c)
    m_block = S_block.max(axis=-1)              # (B_r,)
    m_new = np.maximum(m_old, m_block)
    # Correction rescales the old accumulator into the new max frame
    alpha = np.exp(m_old - m_new)               # (B_r,)
    # New tile's exponentials in the new max frame
    P_block = np.exp(S_block - m_new[:, None])  # (B_r, B_c)
    l_new = alpha * l_old + P_block.sum(axis=-1)
    return m_new, l_new, alpha, P_block

The derivation: sum_j exp(x_j) = exp(m_old) · sum_j exp(x_j − m_old). When we discover a bigger max in the new tile, we have to rescale the old sum by exp(m_old − m_new). This is the same trick used by the Welford online algorithm and by scipy.special.logsumexp.

Step 4: The tiled forward kernel

Now we put it together. The outer loop walks K/V tiles; for each tile we update the online softmax statistics and rescale the running output.

def flash_attention_forward(Q, K, V, Br=32, Bc=32):
    """
    FlashAttention-2 forward pass in pure NumPy.
    Q, K, V: (N, d) float32 arrays
    Returns O: (N, d)
    """
    N, d = Q.shape
    assert K.shape == (N, d) and V.shape == (N, d)
    scale = 1.0 / np.sqrt(d)

    # Running state, one entry per query row
    O = np.zeros((N, d), dtype=Q.dtype)          # output accumulator
    m = np.full(N, -np.inf, dtype=Q.dtype)       # running max
    l = np.zeros(N, dtype=Q.dtype)               # running denominator

    # Trivial to extend to multi-head / batched; we keep it 2D for clarity
    for j_start in range(0, N, Bc):
        j_end = min(j_start + Bc, N)
        K_j = K[j_start:j_end]                   # (Bc, d)
        V_j = V[j_start:j_end]                   # (Bc, d)

        for i_start in range(0, N, Br):
            i_end = min(i_start + Br, N)
            Q_i = Q[i_start:i_end]               # (Br, d)

            # Score tile for this query block against this key block
            S_ij = (Q_i @ K_j.T) * scale         # (Br, Bc)

            # Slice the running state for these query rows
            m_old = m[i_start:i_end].copy()
            l_old = l[i_start:i_end].copy()
            O_i   = O[i_start:i_end].copy()

            m_new, l_new, alpha, P_block = online_softmax_update(
                m_old, l_old, S_ij
            )

            # Rescale accumulator, add new contribution
            # alpha: (Br,), O_i: (Br, d) -> broadcast on rows
            O_i = O_i * alpha[:, None] + P_block @ V_j

            m[i_start:i_end] = m_new
            l[i_start:i_end] = l_new
            O[i_start:i_end] = O_i

    # Final normalization
    O = O / l[:, None]
    return O

A few subtle points worth flagging in your README:

  • We copy() the slice of m, l, O before the inner update because NumPy slicing returns a view and we need a stable snapshot for the rescale step. Forgetting this is the most common bug.
  • The double loop over Bc and Br is what the CUDA kernel fuses into a single pass per K/V tile in the real GPU version (see the official CUDA implementation).
  • We never allocate a (Br, N) or (N, N) matrix. Peak memory for the score block is O(Br · Bc), independent of N.

Step 5: A sanity test

def test_correctness(N=128, d=64, seed=0):
    rng = np.random.default_rng(seed)
    Q = rng.standard_normal((N, d)).astype(np.float32)
    K = rng.standard_normal((N, d)).astype(np.float32)
    V = rng.standard_normal((N, d)).astype(np.float32)

    ref = naive_attention(Q, K, V)
    out = flash_attention_forward(Q, K, V, Br=32, Bc=32)

    max_abs = np.max(np.abs(ref - out))
    max_rel = np.max(np.abs((ref - out) / (np.abs(ref) + 1e-8)))
    print(f"N={N}, d={d}: max|Δ|={max_abs:.3e}, max rel={max_rel:.3e}")
    assert max_abs < 5e-5, "Forward output diverges from reference"

You should see max|Δ| in the 1e-6 to 1e-5 range for float32. Drift grows with N because floating-point summation order differs from the reference; staying under 5e-5 is the bar used by the official FlashAttention tests.

Running and Testing It

Once flash_attention_numpy.py is on disk, the test workflow is just standard Pytest. Drop the snippet from Step 5 into tests/test_correctness.py, then:

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
pytest -q

Parameterize the test across several (N, d, block_size) combinations to catch edge cases:

import pytest

@pytest.mark.parametrize("N", [64, 128, 256])
@pytest.mark.parametrize("d", [32, 64])
@pytest.mark.parametrize("Br", [16, 32])
def test_matches_naive(N, d, Br):
    rng = np.random.default_rng(0)
    Q = rng.standard_normal((N, d)).astype(np.float32)
    K = rng.standard_normal((N, d)).astype(np.float32)
    V = rng.standard_normal((N, d)).astype(np.float32)

    ref = naive_attention(Q, K, V)
    out = flash_attention_forward(Q, K, V, Br=Br, Bc=Br)

    assert np.allclose(ref, out, atol=5e-5, rtol=1e-4)

Run it:

pytest -q tests/test_correctness.py

Expected output:

.............
12 passed in 1.84s

For the memory benchmark, write bench/bench_memory.py that tracks peak array allocation using tracemalloc and sweep N from 128 to 2048:

import tracemalloc, matplotlib.pyplot as plt, numpy as np
from flash_attn_numpy import naive_attention, flash_attention_forward

Ns = [128, 256, 512, 1024, 2048]
naive_peak, flash_peak = [], []

for N in Ns:
    Q, K, V = (np.random.randn(N, 64).astype(np.float32) for _ in range(3))
    for fn, store in ((naive_attention, naive_peak),
                      (flash_attention_forward, flash_peak)):
        tracemalloc.start()
        fn(Q, K, V, Br=32, Bc=32)
        _, peak = tracemalloc.get_traced_memory()
        store.append(peak / 1e6)   # MB
        tracemalloc.stop()

plt.plot(Ns, naive_peak, '-o', label='Naive (O(N^2) scores)')
plt.plot(Ns, flash_peak, '-o', label='FlashAttention-2 NumPy (O(N·Bc))')
plt.xlabel('Sequence length N'); plt.ylabel('Peak memory (MB)')
plt.yscale('log'); plt.legend(); plt.grid(True)
plt.savefig('bench/memory.png', dpi=120)

You should see the naive curve climb as N² while the tiled curve grows linearly with N. That single plot is worth a thousand words on a CV — it visually demonstrates that your implementation avoids the quadratic blowup.

python bench/bench_memory.py

A reasonable README then walks through: the math derivation of online softmax, why tiling helps memory bandwidth, how to read the benchmark plot, and how this generalizes to the CUDA kernel via tl.dot in Triton.

Extending It: Your Roadmap to Senior-Level

A pure NumPy kernel is the starting line, not the finish. Pick four to six of these upgrades; each one maps to a real-world concern at a serious ML infra team.

  • Port the inner loop to Triton or CUDA. Rewrite the (Br, Bc) tile computation as a @triton.jit kernel and benchmark on an A100 or H100. This is the path Triton’s fused attention tutorial takes, and it teaches you GPU memory coalescing. Why it matters: shows you can translate an algorithm from paper to a real GPU kernel, which is the core skill for kernel-engineer roles.
  • Add the backward pass. Implement the recomputation trick described in the original FlashAttention paper: save the (m, l, O) checkpoints from the forward pass and stream gradients tile-by-tile in reverse. Why it matters: training requires gradients; demonstrating you understand checkpointing and the math of the backward pass is exactly what research-engineer interviews probe.
  • Benchmark against PyTorch SDPA. Compare your kernel’s wall-clock and memory against torch.nn.functional.scaled_dot_product_attention across N ∈ {512, 1k, 4k, 16k} using torch.cuda.Event timers. Why it matters: benchmarking rigor is non-negotiable in perf roles, and torch.cuda.Event is the standard tool (PyTorch docs).
  • Wrap it as a PyTorch custom op with torch.library. Expose your NumPy kernel behind torch.ops.myflash.forward(q, k, v) and integrate it into a small Hugging Face GPT2 model. Why it matters: integration matters more than cleverness; this upgrade shows you can ship something usable, not just a demo.
  • Persist results and benchmark history in a SQLite dashboard. Store every benchmark run (N, d, block size, peak memory, wall time, commit hash) in a small SQLite file and render a trend graph via Streamlit. Why it matters: observability and persistence are core backend skills; even a toy kernel looks more serious when its history is queryable.
  • Add fault-tolerance and reproducibility hooks. Pin NumPy and Pytest versions, hash inputs, write a Makefile target that runs the full suite inside a Docker container, and emit a machine-readable JSON report. Why it matters: this is what every production ML pipeline at scale (Weights & Biases, Determined, Meta’s PyTorch CI) does; demonstrating the habit signals senior-level discipline.

Pick two or three, write them as issues in your GitHub repo, and ship them as separate PRs with their own commits. Hiring managers read commit history.

Key Takeaways

  • FlashAttention-2’s core idea is online softmax: a numerically stable streaming update that lets you fold attention tiles into a running output without ever materializing the full N×N score matrix.
  • The update rule is m_new = max(m_old, rowmax(S)); alpha = exp(m_old − m_new); l_new = alpha·l_old + rowsum(exp(S − m_new)); O = alpha·O + exp(S − m_new) @ V.
  • Peak memory drops from O(N²) to O(N · B_c); for N=4096, d=64, that’s the difference between ~64 MB and ~256 KB of score-matrix RAM.
  • A pure NumPy implementation is the fastest way to internalize the algorithm; porting to Triton/CUDA later is a mechanical translation of the same tile loop.
  • On a CV, this project signals paper-reading fluency, numerical-stability chops, memory-hierarchy awareness, and the discipline to ship a tested artifact — all in under 200 lines of code.
  • The natural extensions (Triton port, backward pass, PyTorch integration, benchmarking harness) are the same projects ML-infra teams hire for, so the same repo can grow with you.

Further Reading