TL;DR — In this post you’ll build a FlashAttention‑style tiled attention kernel from scratch using only NumPy, learning how online softmax reduces memory pressure and enables training on modest hardware. The step‑by‑step code is runnable, vectorised, and ready to extend with performance tweaks or integration into larger pipelines. By the end you’ll have a portfolio‑piece that signals genuine systems‑engineering skill.

Building a tiled attention kernel from the ground up is a rare hands‑on project that lets you demonstrate low‑level optimization, memory‑aware design, and reproducibility—all traits hiring managers look for in ML‑focused engineering roles. You’ll walk away with a working NumPy implementation that mirrors the core ideas of FlashAttention, complete with performance benchmarks and a clear roadmap for productionisation.

Why This Project Stands Out on a CV

  • Memory‑aware algorithm design: You implement online (tiling) softmax that never materialises the full (QK^\top) matrix, a technique directly applicable to large‑scale training.
  • NumPy‑level performance tuning: You’ll learn how to fuse reductions, use float16/ bfloat16, and leverage strided slicing for speed.
  • Systems thinking: Choosing tile sizes, balancing compute vs. memory bandwidth, and writing reproducible test harnesses shows you can ship reliable code.
  • Signal for ML‑systems roles: Roles such as ML engineer, deep‑learning systems researcher, and kernel developer value candidates who can bridge research papers and runnable prototypes.
  • Portfolio‑ready artifact: A self‑contained script, benchmark results, and a clear extension path make for a compelling GitHub project.

Architecture Overview

The kernel breaks the classic attention computation into three logical layers:

  1. Input matrices – (Q \in \mathbb{R}^{n \times d}), (K \in \mathbb{R}^{m \times d}), (V \in \mathbb{R}^{m \times d}).
  2. Tiling strategy – Split the (n) query rows into tiles of size (T_q) and the (m) key/value rows into tiles of size (T_k). Each tile pair produces a partial attention output that is accumulated incrementally.
  3. Online softmax – For each tile pair we maintain running max and sum of exponentials, allowing the softmax to be computed without storing the full score matrix.
┌─────────────────────┐
│   Q (n×d)           │
│   K (m×d)           │
│   V (m×d)           │
└───────┬─────────────┘
        │
   ┌────▼─────┐
   │ Tile Q   │
   │ Tile K/V │
   │ Online   │
   │ Softmax  │
   │ Accum    │
   └────▲─────┘
        │
   ┌────▼─────┐
   │ Partial │ ← accumulate over tiles
   │ Output  │
   └─────────┘

Building It Step by Step

Below are five numbered steps that produce a minimal, correct tiled attention kernel. Each step includes a fenced NumPy code snippet you can paste into a file (flash_attention_numpy.py) and run immediately.

Step 1 – Initialise random Q, K, V and set tile size

import numpy as np

d = 64               # head dimension
n = 256              # number of queries
m = 512              # number of keys/values
T_q = 64             # query tile size
T_k = 64             # key/value tile size

np.random.seed(0)
Q = np.random.randn(n, d).astype(np.float32)
K = np.random.randn(m, d).astype(np.float32)
V = np.random.randn(m, d).astype(np.float32)

Step 2 – Scaled dot‑product scores for a single tile pair

def scores_tile(Q_tile, K_tile, d):
    """Q_tile: (T_q, d), K_tile: (T_k, d) → (T_q, T_k)"""
    return (Q_tile @ K_tile.T) / np.sqrt(d)

# Example usage with the first tiles:
Q_tile = Q[:T_q]
K_tile = K[:T_k]
S = scores_tile(Q_tile, K_tile, d)          # shape (T_q, T_k)

Step 3 – Online softmax (max‑then‑sum) on the tile

def online_softmax(S):
    """S: (T_q, T_k) – returns scaled softmax weights."""
    # 1) max per row (keepdim for broadcasting)
    max_S = S.max(axis=1, keepdims=True)               # (T_q, 1)
    # 2) subtract max, exp, sum
    exp_S = np.exp(S - max_S)                          # numerical stability
    sum_exp = exp_S.sum(axis=1, keepdims=True)        # (T_q, 1)
    # 3) final weights
    return exp_S / sum_exp                            # (T_q, T_k)

W = online_softmax(S)        # shape (T_q, T_k)

Step 4 – Weighted sum of the value tile

def apply_softmax_weights(W, V_tile):
    """W: (T_q, T_k), V_tile: (T_k, d) → (T_q, d)"""
    return W @ V_tile   # matrix multiply

V_tile = V[:T_k]
O_tile = apply_softmax_weights(W, V_tile)   # (T_q, d)

Step 5 – Loop over all tiles and accumulate the output

O = np.zeros((n, d), dtype=np.float32)   # final output

for q_start in range(0, n, T_q):
    q_end = min(q_start + T_q, n)
    Q_tile = Q[q_start:q_end]

    for k_start in range(0, m, T_k):
        k_end = min(k_start + T_k, m)
        K_tile = K[k_start:k_end]
        V_tile = V[k_start:k_end]

        S = scores_tile(Q_tile, K_tile, d)
        W = online_softmax(S)
        O_tile = apply_softmax_weights(W, V_tile)

        # accumulate into the correct row slice
        O[q_start:q_end] += O_tile

# O now holds the tiled attention output (same shape as Q @ K^T @ V would, but with far less memory)

Verification (optional) – Compare the tiled result against a naïve dense implementation:

def naive_attention(Q, K, V):
    S = (Q @ K.T) / np.sqrt(d)               # (n, m)
    W = np.exp(S - S.max(axis=1, keepdims=True))
    W = W / W.sum(axis=1, keepdims=True)
    return W @ V

O_naive = naive_attention(Q, K, V)
print("Max absolute error:", np.max(np.abs(O - O_naive)))

Running the script should print an error well below 1e‑5, confirming correctness.

Running and Testing It

  1. Save the code above to flash_attention_numpy.py.
  2. Execute from the command line:
python flash_attention_numpy.py

You’ll see the “Max absolute error” line printed, followed by no other output (the script is silent on success).

  1. Unit‑test – Add a small pytest suite that asserts the error stays under a tolerance for various d, n, m, and tile sizes. This demonstrates reproducibility and makes the project ready for a CI pipeline.

  2. Profiling – Use python -m cProfile flash_attention_numpy.py or timeit to measure runtime. Compare against the naive version to observe the memory‑footprint reduction (the naive approach builds an ((n \times m)) score matrix, while the tiled version never exceeds (T_q \times T_k)).

Extending It: Your Roadmap to Senior‑Level

  1. Mixed‑precision (float16 / bfloat16) support – Cast intermediate matrices to float16 before the softmax and cast back for the final accumulation. Why it matters: Cuts memory bandwidth by 2× on modern GPUs/TPUs while keeping accuracy acceptable for many models.
  2. Integration with CuPy or JAX – Replace the NumPy kernels with CuPy GPU arrays or JAX‑compiled functions to obtain real‑world speed‑ups. Why it matters: Shows you can move a prototype to production hardware without rewriting the algorithm from scratch.
  3. Benchmark against torch.nn.functional.scaled_dot_product_attention – Measure latency, peak memory, and numerical drift. Why it matters: Provides concrete evidence of where your implementation wins or loses versus established libraries.
  4. Causal (masked) attention support – Add a boolean mask that zeroes out future positions before softmax. Why it matters: Essential for language‑model pre‑training and demonstrates awareness of common transformer constraints.
  5. Persisted tile configuration & autotuning – Store the optimal (T_q, T_k) pair per d and hardware in a JSON config, and auto‑tune at import time. Why it matters: Real‑world services need to adapt tile sizes to varying sequence lengths and hardware without manual tweaks.
  6. Gradient‑through attention for back‑prop – Extend the kernel to compute (\partial L / \partial Q, K, V) using the chain rule through the online softmax. Why it matters: Enables training from scratch or fine‑tuning, turning the kernel into a drop‑in replacement for autograd‑tracked ops.

Key Takeaways

  • Online softmax eliminates the (O(nm)) memory bottleneck by processing attention in tiles.
  • Tile size selection is a trade‑off between compute reuse and memory bandwidth; typical values are 64–128 for head dimension 64–128.
  • NumPy vectorisation (broadcasting, keepdims) is sufficient for a correct prototype; real production kernels often lower to C/CUDA for maximal throughput.
  • Benchmarking against a naïve dense implementation validates correctness and quantifies the memory savings.
  • Extensibility paths (mixed precision, GPU back‑ends, causal masking, gradient flow) turn the toy into a production‑ready component.
  • A well‑documented, test‑covered script serves as a portfolio piece that signals systems‑engineering competence to hiring managers.

Further Reading