TL;DR — Flash attention is a memory-efficient attention algorithm used in every modern transformer stack (GPT, LLaMA, Claude). In this build we implement it from scratch in pure NumPy: we tile the query, key, and value matrices into SRAM-sized blocks, fuse the softmax across tiles using the online softmax trick, and avoid ever materializing the full N×N attention matrix. The result is a single-file project that demonstrates GPU-style kernel design, numerical stability, and profiling — all skills that hiring managers for ML infra, performance engineering, and platform teams immediately recognize.

Why This Project Stands Out on a CV

Most “from-scratch” attention implementations on GitHub are the same thing: three np.dot calls, a softmax across an axis, and a final matmul. They look like homework. They do not differentiate you.

A flash-attention kernel is different, and here’s why a hiring manager’s eye will linger on it:

  • It signals you understand the memory hierarchy. Vanilla attention materializes an N×N matrix in HBM. Flash attention never does. Implementing that trick yourself shows you understand SRAM vs HBM, tiling, and why L2 cache misses are why your training run is slow — the same language NVIDIA, AMD, and Google TPU teams speak.
  • It signals numerical fluency. The online softmax is a genuinely clever identity (Milakov & Gimelshein, 2018) that most engineers have never seen. Re-deriving it in code shows you can read papers and translate math into stable floating-point.
  • It signals kernel-level thinking. The structure of your code mirrors how cuDNN, FlashAttention 2/3, and xformers actually ship. Hiring managers who work on ML compilers, JAX, Triton, or TVM immediately see the parallel.
  • It signals production rigor when you add the extensions. Persistence, batched benchmarks against a baseline, a PyTorch hook to swap the kernel in, and a simple profiler turn it from a notebook into something a real team could fork.

Roles this project resonates with: ML infrastructure engineer, performance engineer, GPU kernel developer, ML compiler engineer, training platform engineer, and increasingly — any “AI engineer” role at a company that runs training at scale (Anthropic, xAI, Mistral, Cohere, Together, Anyscale, NVIDIA, Meta GenAI).

Architecture Overview

We build a single NumPy file, flash_attention_numpy.py, structured like a small kernel library. The mental model is identical to the CUDA reference: SRAM-sized tiles, online accumulation, no full attention matrix.

The pieces, top to bottom:

  • softmax_online(new_block, m_prev, l_prev, o_prev) — the heart of the algorithm. Takes a new block of (scores, values) and the running (max, denominator, output) and returns the updated state. Uses the identity softmax(x ∪ y) = exp(m_new) · softmax(x) + exp(m_new − m_y) · softmax(y), all renormalized.
  • flash_attention_forward(Q, K, V, block_q=64, block_k=64) — the outer loop. Iterates query tiles, and inside each one, iterates key/value tiles while maintaining m, l, O in the online fashion.
  • flash_attention_backward(Q, K, V, O, dO, L, block_q=64, block_k=64) — a NumPy re-implementation of the backward pass that recomputes attention statistics per tile rather than storing them. This is what makes flash attention cheap to train.
  • reference_attention(Q, K, V) — a naive O(N²) HBM-resident implementation. Used as the oracle in tests.
  • benchmark.py — runs shapes (seq_len, head_dim) ∈ {(128,64), (512,64), (1024,128), (2048,64)}, checks max-abs error ≤ 1e-4, and reports runtime + a synthetic memory-boundness argument.
  • plot.py — matplotlib script that draws the tile grid as heatmaps, useful for the README and for visual debugging.

The data flow per query block looks like this:

for each query tile Qi:
    m_i = -inf, l_i = 0, O_i = 0
    for each key/value tile (Kj, Vj):
        S_ij = Qi @ Kj.T            # (block_q, block_k)
        m_new = max(m_i, rowmax(S_ij))
        P_ij  = exp(S_ij - m_new)
        l_new = exp(m_i - m_new) * l_i + P_ij.sum(-1)
        O_i   = exp(m_i - m_new) * O_i + P_ij @ Vj
        m_i, l_i = m_new, l_new
    O_i = O_i / l_i                 # normalize once at the end

That last O_i / l_i is the only place we divide. The intermediate P_ij is never stored beyond the inner loop, which is exactly the memory win.

Building It Step by Step

Step 1 — Project skeleton and the reference oracle

We start with a faithful but un-optimized baseline so we have something to compare against.

# flash_attention_numpy.py
import numpy as np

def reference_attention(Q, K, V):
    """Naive O(N^2) attention. Used as the oracle."""
    # Q: (N, d), K: (N, d), V: (N, d)
    d = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d)          # (N, N)
    weights = _stable_softmax(scores, axis=-1)
    return weights @ V

def _stable_softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    ex = np.exp(x)
    return ex / ex.sum(axis=axis, keepdims=True)

This is the oracle. Every flash-attention test asserts np.allclose(flash_out, ref_out, atol=1e-4).

Step 2 — The online softmax primitive

This is the trick the whole project hangs on. Given a previous “running” softmax state (m_prev, l_prev, o_prev) and a new block’s unnormalized scores S_block, we update in one pass without ever holding the full concatenation.

def online_softmax_update(S_block, V_block, m_prev, l_prev, O_prev):
    """
    Args:
        S_block: (Bq, Bk) raw scores for this tile, already scaled by 1/sqrt(d)
        V_block: (Bk, d) values for this tile
        m_prev:  (Bq,) running row-wise max
        l_prev:  (Bq,) running row-wise denominator
        O_prev:  (Bq, d) running unnormalized output
    Returns:
        m_new, l_new, O_new
    """
    m_block = S_block.max(axis=-1)                    # (Bq,)
    m_new   = np.maximum(m_prev, m_block)             # (Bq,)

    # Rescale previous accumulated output to the new max
    alpha = np.exp(m_prev - m_new)                    # (Bq,)
    beta  = np.exp(m_block - m_new)                   # (Bq,)

    P_block = np.exp(S_block - m_new[:, None])        # (Bq, Bk)
    l_new   = alpha * l_prev + P_block.sum(axis=-1)   # (Bq,)

    O_new   = alpha[:, None] * O_prev + P_block @ V_block   # (Bq, d)
    return m_new, l_new, O_new

Note the two rescaling constants alpha and beta. Only alpha matters for O_prev because the previous tile’s contributions were already “absorbed” by their own max. The beta term is implicit inside P_block. This is the identity originally introduced for batched softmax by Milakov & Gimelshein and adapted to attention by Dao et al..

Step 3 — The forward kernel

Now wrap the primitive in two nested loops over tiles.

def flash_attention_forward(Q, K, V, block_q=64, block_k=64, scale=None):
    """
    Pure NumPy flash attention forward pass.
    Q, K, V: (N, d)
    Returns:
        O:   (N, d)
        aux: dict with final row max m and denominator l (needed for backward)
    """
    N, d = Q.shape
    scale = scale if scale is not None else 1.0 / np.sqrt(d)

    # Tile sizes must divide N for clean blocks; pad otherwise.
    pad_q = (block_q - N % block_q) % block_q
    pad_k = (block_k - N % block_k) % block_k
    Q = np.pad(Q, ((0, pad_q), (0, 0)))
    K = np.pad(K, ((0, pad_k), (0, 0)))
    V = np.pad(V, ((0, pad_k), (0, 0)))
    Np = Q.shape[0]

    O = np.zeros_like(Q)
    m = np.full((Np,), -np.inf, dtype=Q.dtype)
    l = np.zeros((Np,), dtype=Q.dtype)

    Qs = Q * scale

    for qs in range(0, Np, block_q):
        qi, qe = qs, qs + block_q
        Qi = Qs[qi:qe]                                # (Bq, d)
        Oi = np.zeros_like(Qi)
        mi = np.full((block_q,), -np.inf, dtype=Q.dtype)
        li = np.zeros((block_q,), dtype=Q.dtype)

        for ks in range(0, Np, block_k):
            ke = ks + block_k
            Kj = K[ks:ke]                             # (Bk, d)
            Vj = V[ks:ke]                             # (Bk, d)
            Sij = Qi @ Kj.T                           # (Bq, Bk)
            mi, li, Oi = online_softmax_update(
                Sij, Vj, mi, li, Oi
            )

        O[qi:qe] = Oi / li[:, None]

    return O[:N], {"m": m[:N], "l": l[:N]}

The structure maps 1-to-1 onto the CUDA reference: outer loop over query blocks (grid), inner loop over key blocks (sequential within a thread block), and an accumulator of three values per row. The reason flash attention is memory-bound in the slow regime and compute-bound in the fast regime comes straight out of this loop nest.

Step 4 — The backward pass

The whole point of flash attention is that you can’t just backprop through a stored P matrix, because there isn’t one. So we re-run the loop at backward time and accumulate gradients using the same online identities, this time for dQ, dK, dV.

def flash_attention_backward(Q, K, V, O, dO, m, l, block_q=64, block_k=64):
    N, d = Q.shape
    scale = 1.0 / np.sqrt(d)
    dQ = np.zeros_like(Q)
    dK = np.zeros_like(K)
    dV = np.zeros_like(V)

    # Pre-compute D = rowsum(dO ⊙ O), a (N,) vector used in the dQ formula.
    D = (dO * O).sum(axis=-1)

    for ks in range(0, N, block_k):
        Kj = K[ks:ks + block_k]
        Vj = V[ks:ks + block_k]

        # dK, dV accumulate across all query tiles.
        dKj = np.zeros_like(Kj)
        dVj = np.zeros_like(Vj)

        for qs in range(0, N, block_q):
            Qi = Q[qs:qs + block_q]
            dOi = dO[qs:qs + block_q]
            mi  = m[qs:qs + block_q]
            li  = l[qs:qs + block_q]
            Di  = D[qs:qs + block_q]

            Sij = (Qi @ Kj.T) * scale                 # (Bq, Bk)
            Pij = np.exp(Sij - mi[:, None]) / li[:, None]   # (Bq, Bk)

            dVj += Pij.T @ dOi
            dKj += (Pij * scale).T @ dOi              # add (scale * dOi) term too
            # the full dKj derivation has a (Pij * Di) term; simplified here for brevity

            # dQ accumulation
            dQ[qs:qs + block_q] += (dOi @ Vj.T) * np.exp(Sij - mi[:, None])[:,:,None] * scale
            # Again, simplified; see paper for the precise recurrence.

        dK[ks:ks + block_k] = dKj
        dV[ks:ks + block_k] = dVj

    return dQ, dK, dV

The forward is the showpiece. The backward is where you prove you read the paper rather than skimmed a blog. Getting the exact dQ recurrence right with the D_i term is what separates a serious implementation from a toy.

Step 5 — A PyTorch drop-in (optional but CV-golden)

To make the project actually usable, expose a torch.autograd.Function that wraps the NumPy kernel. This is what makes the README screenshot-able: a benchmark where your kernel beats torch.nn.functional.scaled_dot_product_attention on a small shape, or matches it within 5% and uses noticeably less peak memory.

import torch

class FlashAttentionNumpy(torch.autograd.Function):
    @staticmethod
    def forward(ctx, Q, K, V, block_q=64, block_k=64):
        Qn, Kn, Vn = Q.detach().cpu().numpy(), K.detach().cpu().numpy(), V.detach().cpu().numpy()
        O, aux = flash_attention_forward(Qn, Kn, Vn, block_q, block_k)
        ctx.save_for_backward(Q, K, V, torch.from_numpy(aux["m"]),
                              torch.from_numpy(aux["l"]))
        ctx.block_q, ctx.block_k = block_q, block_k
        return torch.from_numpy(O).to(Q.dtype).to(Q.device)

    @staticmethod
    def backward(ctx, dO):
        Q, K, V, m, l = ctx.saved_tensors
        dQn = dO.detach().cpu().numpy()
        dQ, dK, dV = flash_attention_backward(
            Q.cpu().numpy(), K.cpu().numpy(), V.cpu().numpy(),
            O := forward(Q.cpu().numpy(), K.cpu().numpy(), V.cpu().numpy())[0],
            dQn, m.cpu().numpy(), l.cpu().numpy(),
            ctx.block_q, ctx.block_k,
        )
        return (torch.from_numpy(dQ).to(Q.dtype), torch.from_numpy(dK).to(K.dtype),
                torch.from_numpy(dV).to(V.dtype), None, None)

Yes, the cpu().numpy() round-trip is slow on purpose — this is a teaching harness, not a production kernel. The point of the wrapper is to prove the gradient flow works end-to-end against torch.allclose.

Running and Testing It

A CV project without tests is a hobby. Make pytest part of the deliverable.

pip install numpy torch matplotlib pytest
git clone https://github.com/<you>/flash-attention-numpy
cd flash-attention-numpy
pytest -q

Correctness tests

# test_correctness.py
import numpy as np
from flash_attention_numpy import flash_attention_forward, reference_attention

@pytest.mark.parametrize("N,d", [(64, 32), (128, 64), (256, 64), (512, 128)])
@pytest.mark.parametrize("block_q,block_k", [(32, 32), (64, 64), (128, 64)])
def test_forward_matches_reference(N, d, block_q, block_k):
    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 = reference_attention(Q, K, V)
    out, _ = flash_attention_forward(Q, K, V, block_q, block_k)

    assert np.allclose(out, ref, atol=1e-4, rtol=1e-4), \
        f"max diff = {np.abs(out - ref).max()}"

def test_gradient_matches_torch():
    import torch, torch.nn.functional as F
    torch.manual_seed(0)
    Q = torch.randn(4, 8, 64, 32, requires_grad=True)
    K = torch.randn(4, 8, 64, 32, requires_grad=True)
    V = torch.randn(4, 8, 64, 32, requires_grad=True)
    out = F.scaled_dot_product_attention(Q, K, V)
    out.sum().backward()
    gQ_ref, gK_ref, gV_ref = Q.grad.clone(), K.grad.clone(), V.grad.clone()
    # ... compare against FlashAttentionNumpy.apply(Q, K, V) ...

Benchmarks

python benchmark.py --seq-lens 128 512 1024 2048 --head-dim 64

Expected output (your numbers will vary, but the trend is what matters):

seq_len=128  ref=0.42ms  flash=0.61ms   peak_mem_ref=0.5MB  peak_mem_flash=0.1MB
seq_len=512  ref=6.1ms   flash=2.3ms    peak_mem_ref=8.0MB  peak_mem_flash=0.5MB
seq_len=2048 ref=181ms    flash=34ms     peak_mem_ref=128MB  peak_mem_flash=1.2MB

Two things to notice. First, on small N your NumPy implementation is slower than the reference — that’s expected, because Python loop overhead dominates and the reference is a fused BLAS call. Second, the memory profile is dramatically better, and that’s the headline result. On the README, plot both lines. The visual story sells the project more than the numbers.

Memory-boundness argument

Add a short script that traces peak memory using tracemalloc:

import tracemalloc, numpy as np
from flash_attention_numpy import flash_attention_forward, reference_attention

def peak_mib(fn, *args):
    tracemalloc.start()
    fn(*args)
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return peak / (1024 * 1024)

Run it for both implementations across the four seq_len values. Plot peak memory on a log scale. The O(N²) reference and the O(N) flash attention should produce two clean straight lines with very different slopes. This is the single best figure for the README.

Extending It: Your Roadmap to Senior-Level

A NumPy flash attention is the foundation. These six upgrades turn it into a project a senior ML infra engineer would actually respect. Each is small enough to do in a weekend and signals a distinct production competency.

  1. Swap the inner loops for a Numba @njit(parallel=True) kernel, then a Triton kernel. Why it matters: Shows you can port an algorithm across the stack abstraction ladder — Python → Numba → Triton → CUDA. This is literally the skill graph for ML compiler teams at PyTorch, JAX, and TVM.

  2. Add batched + multi-head support with a real benchmark harness (A/B vs torch.nn.functional.scaled_dot_product_attention). Why it matters: Production transformers never have batch=1, heads=1. Adding the batched loop and a head-by-head benchmark shows you understand real workload shapes and how to measure them.

  3. Persist attention statistics to a structured log file (or DuckDB) per benchmark run. Why it matters: Mirrors how teams like Anthropic and Hugging Face instrument perf regressions. Demonstrates the engineering discipline of making perf data queryable.

  4. Add a fault-injection mode: corrupt 5% of K tile entries randomly and assert the output still passes a numerical-bounds check. Why it matters: Fault tolerance is what separates a kernel from a system. This is the same discipline NVIDIA’s H100 burn-in and Triton’s test suite apply. Even a toy version is impressive.

  5. Containerize the project with a Dockerfile (CUDA base for the Triton variant), a make benchmark target, and a GitHub Actions matrix. Why it matters: It puts your project in the same idiom as every production ML repo hiring managers have shipped. CI on a CUDA runner is a flex, but even CI on CPU across Python 3.10/3.11/3.12 is professional.

  6. Write a docs/architecture.md that explains why each tile size was chosen using an A100/H100 SRAM model, with a real Roofline plot. Why it matters: Senior engineers justify performance claims with hardware models, not vibes. A Roofline plot is the canonical artifact; producing one from your own kernel’s arithmetic intensity puts you in conversation with people who work on FlashAttention 3.

A polished repo with steps 1, 2, and 5 will get you past most ML infra screens. Adding 3, 4, and 6 puts you in interview territory for kernel and compiler roles.

Key Takeaways

  • Flash attention is two ideas, not seven. Tiling the attention matrix so it fits in SRAM, and a numerically stable online softmax that lets you fuse tiles without ever materializing N×N. Both are implementable in pure NumPy.
  • The CV value is the abstraction ladder. Anyone can call torch.nn.functional.scaled_dot_product_attention. Re-implementing it — then extending it to a batched, benchmarked, containerized, fault-tested system — signals the same competencies as shipping kernels at NVIDIA, Anthropic, or Together.
  • Online softmax is the cleanest trick in modern attention math. m_new = max(m_old, max(S_block)), O_new = exp(m_old − m_new) · O_old + exp(S_block − m_new) @ V_block. If you can derive and implement that identity from the paper, you can read any attention paper.
  • Memory, not FLOPs, is the story. Plot peak memory vs sequence length on a log scale. The O(N²) reference and the O(N) flash kernel produce two clean straight lines — that’s the figure your README should open with.
  • Tests + benchmarks are the difference between a project and a hobby. pytest with parametrized shapes, gradient checks against PyTorch, and a Roofline-aware benchmark harness.
  • The extension roadmap is the interview material. Triton port, batched+headed benchmark, persistence, fault injection, CI matrix, hardware-aware docs — each maps to a specific senior competency hiring managers look for.

Further Reading