04 — The Math (and Numerics) You Actually Need

“Not a math degree. Specifically: matmul as the atom, softmax numerics as the heart of attention, bit layouts as the price of speed, and enough probability to sample intelligently.”

This file is the densest. Work through it with paper next to the screen. Every arithmetic exercise here re-appears in Phases 1–5 with harder numbers.

Part 1 — Matrix Multiplication as the Atomic Unit

The definition, blocked

For C = A @ B with A: M×K, B: K×N, C: M×N:

C[i,j] = sum_{k=0..K-1} A[i,k] * B[k,j]

Naive triple loop:

for (int i = 0; i < M; ++i)
    for (int j = 0; j < N; ++j)
        for (int k = 0; k < K; ++k)
            C[i][j] += A[i][k] * B[k][j];

On a large matrix (say 1024×1024 fp32, 4 MB per matrix), this is catastrophic on the cache. B[k][j] strides through columns; every access misses. You will measure ~1 GFLOPS on a machine capable of 100+.

Tiled (blocked) matmul

Block C into tiles of size BM × BN, A into BM × BK, B into BK × BN. For each output tile, stream through K:

for (int i0 = 0; i0 < M; i0 += BM)
  for (int j0 = 0; j0 < N; j0 += BN)
    for (int k0 = 0; k0 < K; k0 += BK)
      // Multiply-accumulate the BM×BK block of A into the BK×BN block of B into the BM×BN block of C.
      for (int i = i0; i < i0+BM; ++i)
        for (int j = j0; j < j0+BN; ++j)
          for (int k = k0; k < k0+BK; ++k)
            C[i][j] += A[i][k] * B[k][j];

Choose BM, BN, BK so that one tile of A + one tile of B + one tile of C fits in L1 (typically ~32 KB on x86). For fp32, BM = BN = BK = 32 gives 3 × 32×32×4 = 12 KB — comfortable. Result: ~10× speedup on 1024² fp32 matmul on a modern CPU.

The arithmetic intensity story:

  • Naive: reads M*N*K + M*N*K = 2 M N K elements from memory (assuming cache misses on B).

  • Tiled: for each BM×BN output tile, reads M K elements of A once, K N elements of B once (approximately, with M/BM tiles per row of C).

  • Arithmetic intensity climbs from ~O(1) FLOPs/byte to O(BM) FLOPs/byte. That is the entire game.

The GPU generalization (preview for Phase 2)

On a GPU the story is identical but with three levels: HBM → shared memory (per-block L1-equivalent) → registers. You block twice: once from HBM to shared memory, once from shared memory to registers. Simon Boehm’s SGEMM ladder (Phase 2) is exactly this idea, applied with love.

Do this now

Write matmul_naive and matmul_tiled in C++ for fp32, 1024×1024. Measure with perf stat. Predict cache-miss ratio before measuring. Predicted-vs-measured discipline starts here.

Part 2 — Softmax Numerics: The Heart of FlashAttention

Naive softmax

def softmax_naive(x):
    return np.exp(x) / np.sum(np.exp(x))

This breaks when any x_i > ~89 (in fp32, exp(89.4) overflows to inf). In fp16, threshold is ~11. Attention logits, before softmax, easily exceed these.

The max-subtraction trick (numerically stable)

def softmax_stable(x):
    m = np.max(x)
    e = np.exp(x - m)
    return e / np.sum(e)

Mathematically identical (multiply top and bottom by exp(-m)), numerically bounded (all exp arguments are ≤ 0, so exp results are ≤ 1, no overflow).

Why this trick is the whole game: it makes softmax a two-pass algorithm (pass 1: find max; pass 2: compute exponents and sum). For attention, that’s two passes over the K matrix per row of Q. On HBM this is expensive. The next section fixes it.

Online (streaming) softmax — Milakov & Gimelshein, 2018

One-pass softmax by maintaining running m (max) and running l (sum-of-exp-relative-to-current-max). For each new element x_i:

m_new = max(m_old, x_i)
l_new = l_old * exp(m_old - m_new) + exp(x_i - m_new)

After the pass, softmax denominator is l_new and the running m_new is the true max. To emit softmax(x_j) for a stored x_j, use exp(x_j - m_new) / l_new.

Derive this from scratch until you can rebuild it in 30 seconds on a whiteboard. This identity is what FlashAttention exploits to tile Q/K/V into SRAM without ever materializing the N×N attention matrix. If you internalize nothing else from Phase 0, internalize this recurrence.

Reading:

The consequence: FlashAttention in one sentence

“Because softmax admits an online algorithm, we can tile attention along the sequence dimension, compute S = QK^T tile by tile into shared memory, rescale as we go, and never write the N×N matrix to HBM.”

If that sentence isn’t obvious yet, read the paper again. It will be by Phase 3.

Part 3 — Floating-Point Formats: Bit Layouts to Memorize

Inference performance is a story about how few bits you can afford to keep. Draw these from memory.

The universal layout

[sign : 1 bit] [exponent : e bits] [mantissa : m bits]

Value = (-1)^sign * (1 + mantissa_as_fraction) * 2^(exponent - bias) for normal numbers. Bias = 2^(e-1) - 1 (so exponents span roughly -bias to +bias).

FP32 (IEEE 754 single)

[1][8 exp][23 mantissa]  — total 32 bits, 4 bytes
Bias = 127. Range ~1.2e-38 to ~3.4e38. Precision ~7 decimal digits.

Baseline reference format. All comparisons made against this.

FP16 (IEEE 754 half)

[1][5 exp][10 mantissa]  — 16 bits, 2 bytes
Bias = 15. Range ~6e-5 to ~65504. Precision ~3 decimal digits.

The problem: 65504 max means training gradients underflow/overflow. Loss scaling was invented to work around this.

BF16 (Brain Float 16) — Google’s fix

[1][8 exp][7 mantissa]   — 16 bits, 2 bytes
Bias = 127 (same as fp32). Range ~1.2e-38 to ~3.4e38. Precision ~2 decimal digits.

The key insight: bf16 keeps fp32’s exponent (so it has fp32’s dynamic range) and truncates the mantissa. Training becomes numerically boring — no loss scaling needed. Precision drops but neural nets don’t care much. This is why bf16 became the training default on A100/H100 and TPUs.

FP8-E4M3 (inference forward pass)

[1][4 exp][3 mantissa]   — 8 bits, 1 byte
Bias = 7. Range: ±448 (does NOT represent Inf; single NaN bit-pattern). Precision ~1 decimal digit.

Higher precision (mantissa), less range. Weights and activations forward. Per-tensor or per-block scaling required to keep values in range. Reference: NVIDIA/Arm/Intel FP8 spec, Micikevicius et al., 2022: https://arxiv.org/abs/2209.05433.

FP8-E5M2 (gradients backward)

[1][5 exp][2 mantissa]   — 8 bits, 1 byte
Bias = 15. Range ±57344. IEEE 754-like (has Inf, NaNs). Precision <1 decimal digit.

More range for gradients (which span many magnitudes), less precision (which gradients tolerate). Not commonly used at inference time — inference is dominated by e4m3.

INT8

[8 bits signed] — no exponent, no mantissa. Just -128..127.
Or unsigned uint8: 0..255.

Quantization: x_int8 = round(x_fp / scale) + zero_point, clipped. Symmetric (zero_point = 0) simpler; asymmetric (zero_point ≠ 0) fits unsigned distributions better. Per-channel or per-group scaling to preserve dynamic range.

INT4

[4 bits] — either signed (-8..7) or unsigned (0..15). Packed two-per-byte.

Used in GPTQ, AWQ, GGUF Q4_K/IQ4 quantization. Weight-only quantization: weights stored as int4 with per-group scale + zero_point (~32 fp16 scales per 128 weights → ~4.5 bits effective). Dequantized in registers just before the tensor-core MMA.

The Blackwell additions (awareness only for Phase 0)

The horror-story exercise

Compute: at fp32, a 70B model needs 280 GB of weight memory. At bf16, 140 GB. At fp8-e4m3, 70 GB. At int4, ~35 GB (plus scales). This is why quantization exists.

The bit-layout drill

On a blank sheet, draw fp32, bf16, fp8-e4m3, fp8-e5m2, int8, int4. Label sign / exponent bits / mantissa bits, bias, min/max representable normal, and precision. Do this cold once a week for a month until it’s mechanical. This exact drill has appeared in ML systems studies at NVIDIA, DeepSeek, Anthropic, Meta.

Part 4 — Sampling Math

Generative inference ends in a probability distribution over the vocabulary. The sampling step turns that into a chosen token. Own the math.

Greedy

token = argmax(logits)

Deterministic. Repetitive (models revisit high-probability loops). Rarely optimal.

Temperature

probs = softmax(logits / T)
token = sample(probs)

  • T → 0: approaches greedy.

  • T = 1: use raw distribution.

  • T > 1: flattens (more diverse, more error).

  • T < 1: sharpens.

Top-k

top_k_logits, top_k_indices = top_k(logits, k)
probs = softmax(top_k_logits / T)
token = top_k_indices[sample(probs)]

Keep only the k highest-probability tokens. Common k: 40, 50, 100.

Top-p (nucleus sampling; Holtzman et al. 2019)

Sort probs descending; keep the smallest set whose cumulative probability ≥ p. Renormalize; sample.

sorted_probs, sorted_indices = sort(softmax(logits), descending=True)
cumsum = cumulative_sum(sorted_probs)
cutoff = argmax(cumsum >= p)  # smallest cutoff satisfying threshold
top_p_probs = sorted_probs[:cutoff+1] / cumsum[cutoff]
token = sorted_indices[sample(top_p_probs)]

Common p: 0.9, 0.95. Adaptive: fewer candidates when the distribution is peaked, more when it’s flat.

Min-p (2024, gaining adoption)

Keep tokens whose probability ≥ min_p * max_prob. More robust to varying distribution shape than top-p. Paper: “Min-P Sampling” (Nguyen et al., 2024): https://arxiv.org/abs/2407.01082.

Repetition penalty

Divide (or subtract from) logits of already-generated tokens by a factor > 1. Cheap; effective against loops.

The logit-processor abstraction

Every sampler above can be expressed as a chain of transformations on the logit vector before the final softmax-and-sample step:

def sample(logits, processors):
    for p in processors:      # temperature, top-k mask, top-p mask, penalty, grammar mask, ...
        logits = p(logits)
    probs = softmax(logits)
    return categorical_sample(probs)

This is the interface HuggingFace LogitsProcessor, vLLM’s sampler, and structured-output frameworks (Outlines, XGrammar) all use. Implement it yourself in Phase 1 — it maps directly onto guided decoding in Phase 4.

Part 5 — Backprop Mechanics at Karpathy Depth

You probably know this already, but the review is worth 90 minutes. Karpathy’s micrograd (https://github.com/karpathy/micrograd, ~13k stars) is a scalar-valued autograd engine in ~150 lines of Python. Read it. Understand:

  • Every op stores its inputs and a local gradient rule.

  • backward() topological-sorts the computation graph and calls each op’s _backward in reverse.

  • Chain rule = accumulate gradients.

The extension from scalars to tensors is bookkeeping. You don’t need to write backward passes for kernels in this roadmap (inference-focused), but you need to understand how PyTorch does it — because CUDA graph capture, torch.compile, and DDP grad-sync all live on top of this mental model.

Companion video (~2h): Karpathy, “The spelled-out intro to neural networks and backpropagation”: https://www.youtube.com/watch?v=VMj-3S1tku0

Part 6 — Roofline Model (Williams et al. 2009)

The rule that governs every optimization. Given a device with peak FLOP/s P and memory bandwidth B:

  • Ridge point: P/B FLOPs per byte moved.

  • Kernel with arithmetic intensity I (FLOPs/byte) achieves at most min(P, I * B) FLOP/s.

  • Below ridge: memory-bound. Optimize by reducing bytes moved (fusion, quantization, tiling).

  • Above ridge: compute-bound. Optimize by using more of the tensor cores (better GEMM, more parallelism).

H100 SXM numbers to memorize:

  • BF16 dense TC: ~989 TFLOP/s

  • HBM3: ~3.35 TB/s

  • Ridge point: ~989e12 / 3.35e12 ≈ 295 FLOPs/byte (bf16)

  • FP8 dense TC: ~1979 TFLOP/s → ridge ≈ 591 FLOPs/byte (fp8, with fp8 bytes)

Decode arithmetic intensity, roughly: for batch B, one forward token, weights of size W bytes:

  • Bytes moved ≈ W (weights) + KV footprint + activations. For a 7B fp16 model at batch 1, ≈ 14 GB.

  • FLOPs ≈ 2 * params * B ≈ 14 GFLOP at B=1.

  • Intensity ≈ 14e9 / 14e9 = 1 FLOP/byte. Deep in memory-bound territory.

  • Increasing batch to 32: FLOPs scale linearly, bytes-of-weights stays constant. Intensity ≈ 32. Still below ridge, but heading toward it. This is why batching matters so much.

Paper: Williams, Waterman, Patterson, “Roofline: an insightful visual performance model for multicore architectures” (2009): https://dl.acm.org/doi/10.1145/1498765.1498785

Draw the H100 roofline on paper. Plot: decode at batch 1, decode at batch 32, prefill at 4k tokens, a fused RMSNorm kernel. This diagram is your mental model for Phases 2–5.

Exit Deliverable

  1. matmul.cpp: naive and tiled, with a benchmark harness reporting GFLOPS and perf cache-miss deltas. Aim for ≥8× speedup; document it.

  2. softmax_derivation.md: the online softmax recurrence, derived from scratch, with a Python check that it matches scipy.special.softmax.

  3. bit_layouts.md: the six formats (fp32, bf16, fp16, fp8-e4m3, fp8-e5m2, int4) drawn as ASCII diagrams with example encodings of 1.0, -0.5, and the smallest positive normal.

  4. roofline.md: the H100 roofline plot with your annotated operating points for decode-B1, decode-B32, prefill-4k, RMSNorm.

  5. sampling.py: greedy / temperature / top-k / top-p / min-p / repetition penalty implemented as a LogitsProcessor chain, with unit tests.