DL S013 · Vectorization and Speed

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 13 / 80 · Module M02 · ~2 hours

The story

Every ML engineer has a moment where they realise their code is 100× too slow because they wrote a Python loop where a matmul would do. Mine was in 2019 — a "batch" processor that iterated over each sample in a mini-batch and applied a function one at a time. It ran at 5 examples/second. After I replaced the loop with X @ W, it ran at 1000 examples/second. Same math. Two lines shorter.

Vectorization is not about being clever with math. It's about giving your work to library code that has been carefully hand-tuned in C, Fortran, and assembly for the last 40 years. BLAS. LAPACK. MKL. Intel's optimized numpy backend. Apple's Accelerate. NVIDIA's cuBLAS. These libraries take a matmul and turn it into a symphony of SIMD instructions, cache-blocked memory accesses, and multi-threaded execution. A tight Python for loop is a single-threaded interpreter jog. There is no contest.

Today's session takes the MNIST MLP from S011 and makes it fast. Same accuracy. Same math. Different speed. Along the way you'll learn einsum (the one syntax to rule them all), the shape-broadcasting patterns that show up in attention, and enough performance intuition to catch slow code before you commit it.

Prerequisites

• Session 002 — matmul.
• Session 003 — broadcasting is 50% of vectorization.
• Session 011 — the code we'll speed up.

1 · The one-paragraph explanation

Python is slow because every operation goes through the interpreter: bytecode dispatch, reference-counted object allocation, dynamic type lookup. A Python for loop over a million floats takes ~1 second.

NumPy operations, in contrast, take one Python call (which does bytecode overhead once), then descend into C code that operates on a contiguous block of memory: no reference counting, no dynamic dispatch, and it uses SIMD instructions to process 4-8 floats per clock cycle. A million-element np.dot takes ~0.5 milliseconds. ~2000× faster.

Additionally, numpy's matmul is dispatched to BLAS (typically Intel MKL or OpenBLAS), which uses multi-threading, cache-blocking, and hand-tuned kernels. On modern CPUs, matmul reaches ~90% of theoretical peak FLOPs. Python loops reach ~0.01%.

Rule: Any inner loop over data should be a numpy expression. Any inner-inner loop can only exist inside a numpy expression.

2 · A concrete timing experiment

import numpy as np
import time
 
def dot_python(a, b):
    s = 0.0
    for i in range(len(a)):
        s += a[i] * b[i]
    return s
 
n = 1_000_000
a = np.random.randn(n).astype(np.float32)
b = np.random.randn(n).astype(np.float32)
 
t = time.time(); dot_python(a, b);          print(f"python loop: {time.time()-t:.3f} s")
t = time.time(); (a * b).sum();             print(f"numpy       : {time.time()-t:.3f} s")
t = time.time(); a @ b;                     print(f"np.dot / @  : {time.time()-t:.6f} s")

Typical output on a modern laptop:

python loop: 0.842 s
numpy      : 0.002 s
np.dot / @ : 0.000567 s

Python loop → numpy expression: ~400× faster. Numpy expression → BLAS matmul: another ~4× faster. Combined: ~1500× faster for the same computation.

Every Python loop you catch and replace with numpy is a similar win.

3 · Einsum — the notation to rule them all

np.einsum is a domain-specific language for tensor operations. It looks weird for 30 seconds, then it clicks and you'll write everything in einsum.

The syntax: np.einsum("subscripts -> output", *operands). Each subscript is a dimension label. Repeated labels get contracted (summed over). Non-repeated labels stay.

3.1 Six examples to memorize

Read those six until they're obvious. Every op in every transformer you'll ever write can be expressed in einsum. Session 036 (self-attention from scratch) leans on this hard.

3.2 Why bother with einsum if @ works?

Two reasons:

1. Reads like the math. "bqd,bkd->bqk" is more explicit than Q @ K.transpose(-1, -2) about what's happening.
2. Handles arbitrary tensor ranks. For 4-way tensor contractions (attention with multiple heads), einsum is the clean way.

Downside: einsum's default plan can be slower than a hand-optimized sequence for very complex contractions. Use np.einsum_path(..., optimize=True) (or torch.einsum with opt_einsum installed) for those.

4 · Vectorizing the S011 MNIST script — before and after

The S011 script is already vectorized (we use X @ W1, not a Python loop over examples). But there's a common naive version people write first. Let's compare.

4.1 The naive version — slow

def forward_naive(X, W1, b1, W2, b2):
    N = X.shape[0]
    H = np.zeros((N, 128))
    for i in range(N):                    # PYTHON LOOP over examples
        for j in range(128):              # PYTHON LOOP over hidden units
            H[i, j] = max(0, X[i] @ W1[:, j] + b1[j])
    Y_pre = np.zeros((N, 10))
    for i in range(N):
        for j in range(10):
            Y_pre[i, j] = H[i] @ W2[:, j] + b2[j]
    return Y_pre

Correct math. Two nested Python loops. On a batch of 64 with 128 hidden units, this is ~8000 Python iterations per forward pass. Slow.

4.2 The vectorized version — fast

def forward_fast(X, W1, b1, W2, b2):
    H = np.maximum(0, X @ W1 + b1)        # ONE numpy call
    Y_pre = H @ W2 + b2                    # ONE numpy call
    return Y_pre

Same math. Zero Python-level loops. Two numpy expressions.

4.3 Timing

import time
X = np.random.randn(64, 784).astype(np.float32)
W1, b1, W2, b2 = init_params()
 
t = time.time()
for _ in range(100): forward_naive(X, W1, b1, W2, b2)
print(f"naive:      {time.time()-t:.3f} s")
 
t = time.time()
for _ in range(100): forward_fast(X, W1, b1, W2, b2)
print(f"vectorized: {time.time()-t:.3f} s")

Output on my laptop:

naive:      2.94 s
vectorized: 0.019 s

155× speedup. Same result, exactly. If you took the naive version of S011's whole training loop, it would take ~4 hours per epoch instead of ~9 seconds. That's why you never write Python loops over neural network operations.

5 · Batching — the ONE OPERATION that makes GPUs matter

A GPU's compute peak is ~100× a CPU's. But you only see that speedup if you feed it enough parallel work. A single-example matmul on a GPU is often SLOWER than on a CPU because you spend more time launching the kernel than computing.

The solution: bigger batches. Instead of one example at a time, process 256 or 1024 examples simultaneously. Each example runs on its own SIMD lane / thread. GPU utilization goes from 2% to 90%.

Batching is why:

• Modern training uses batch sizes 64+ (bare minimum) up to 4096 (LLM pretraining).
• Inference optimization is HARD: users send one query at a time; you have to batch them together (continuous batching, Session 070).
• Every framework op is designed to take batched inputs first, single-example second.

Rule: the first dim of any tensor is always batch. Every op takes advantage of it. Session 016 (DataLoader), Session 018 (GPU), and Session 070 (batching for inference) all revolve around this.

6 · FLOP counting — know what you're paying for

For a matmul (M, K) @ (K, N) → (M, N):

FLOPs = 2 * M * K * N       (each output entry: K mults + K-1 adds ≈ 2K flops)

MNIST layer 1 forward on batch 64:

• (64, 784) @ (784, 128) = 2 * 64 * 784 * 128 = 12.8 MFLOPs.

Modern CPU peak: ~100 GFLOPs. So this matmul takes ~0.13 ms if BLAS gets full peak. Actual observed: ~0.5 ms (BLAS overhead + fp32 vs int).

MNIST layer 2 forward on batch 64:

• (64, 128) @ (128, 10) = 2 * 64 * 128 * 10 = 0.16 MFLOPs. 80× fewer.

Rule: in most networks, ~90% of compute is in the largest few matmuls. Optimize those first.

6.1 For a transformer forward pass

Roughly 12 * n_layers * d_model^2 * seq_len per token (attention + FFN). For a 7B model, ~14 GFLOPs per token per forward pass. Multiply by billions of tokens for training. That's why you need clusters.

7 · Cases where you CAN'T vectorize

Some things genuinely can't be vectorized because they're inherently sequential:

• RNN forward pass (S032) — each timestep's output depends on the previous, so you can't do all timesteps in parallel. You CAN batch across the batch dim though; that's where the parallelism lives.
• Autoregressive generation (S041) — you sample one token, feed it back in, sample the next. Each token depends on all prior. This is why LLM inference is memory-bandwidth-bound (S070).
• Rejection sampling / beam search with early termination — some paths terminate early, some don't. Hard to vectorize without wasted work.

For these, you focus on (a) batching across whatever IS independent, and (b) minimizing per-step overhead (KV cache, S042).

8 · Vectorized backward pass — same story

The S011 backward pass:

dZ2 = P.copy();  dZ2[np.arange(N), y] -= 1;  dZ2 /= N
dW2 = H.T @ dZ2                              # one matmul, ~0.16 MFLOP
db2 = dZ2.sum(axis=0)                        # numpy reduction
dH  = dZ2 @ W2.T                             # one matmul
dZ1 = dH * (Z1 > 0)                          # elementwise
dW1 = X.T @ dZ1                              # one matmul, ~12.8 MFLOP
db1 = dZ1.sum(axis=0)

Six lines. Three matmuls, two reductions, one elementwise mask. All vectorized. Backward is at most 2× the FLOPs of forward. That's why deep learning is affordable.

9 · The einsum version of MLP forward — for the joy of it

def forward_einsum(X, W1, b1, W2, b2):
    Z1 = np.einsum("nd,dh->nh", X, W1) + b1      # (N, hidden)
    H  = np.maximum(0, Z1)
    Z2 = np.einsum("nh,hc->nc", H, W2) + b2      # (N, classes)
    return Z2

Identical performance to the @ version (numpy dispatches both to BLAS). But when we get to multi-head attention (Session 037), einsum becomes indispensable:

# Multi-head attention forward — a taste of Session 037
# Q, K, V: (batch, seq, n_heads, d_head)
scores = np.einsum("bshd,bthd->bhst", Q, K) / np.sqrt(d_head)   # (b, h, s, t)
attn   = softmax(scores)
out    = np.einsum("bhst,bthd->bshd", attn, V)                   # (b, s, h, d)

Try writing that with @ and transposes. You'll appreciate einsum.

10 · War stories

11 · Retention scaffold

Recall — from memory

Stretch — for one extra hour

Rewrite the S012 micrograd training loop for the 4-point moons dataset in vectorized numpy (no Value class, no per-scalar ops). Time both. Expect ~100× speedup. This is the transition from "clarity" (micrograd) to "speed" (numpy/PyTorch) that motivates all of M03.

In your own words

Why is vectorization the biggest performance lever in ML?

Spaced-review pointer

• From S002 — matmul shape rules are the atoms of every vectorized op.
• From S003 — broadcasting is the OTHER 50% of vectorization; einsum handles the "reduce" case, broadcasting handles the "match" case.
• From S011 — the fast forward pass you saw there is what vectorization looks like in practice.

Next-session teaser

Module M02 is DONE. You've built neural networks from first principles: forward, backward, autograd, vectorization. Session 014 opens Module M03 with PyTorch. Same MLP you built in S011, rewritten in PyTorch — half the code, GPU-ready, autograd for free. Everything you learned about the mathematics still applies; PyTorch just handles the engineering. From here on, we build fast.

What to bring back tomorrow — sticky note

• Rule: never write a Python loop over ML data. Ever.
• Formula: matmul FLOPs = 2*M*K*N.
• Notation: np.einsum("mk,kn->mn", A, B) = A @ B. Learn einsum. It'll save you many transformer implementations.

Previous: ← DL S012 · Micrograd — Build a Tiny Autograd Engine · Next: DL S014 · PyTorch Tensors and Autograd →