Frontier Lab Engineer

Chapter 1: Roofline Analysis

Your colleague submits a Pallas kernel for grouped-query attention. It runs at 180 TFLOPS on an H100. Is that good? Is there room for improvement? Without a framework for answering this question, you are guessing. The roofline model gives you a principled upper bound on performance for any operation, in under 60 seconds of mental math.

Reiner Pope (DeepMind) popularized a style of back-of-the-envelope analysis where every architecture decision — attention variant, parallelism strategy, batch size — is evaluated through the roofline lens. At a frontier lab, this is the first tool you reach for. Master it, and you will never be surprised by a kernel's performance again.

The Two Resources: Compute and Bandwidth

A GPU has two fundamental resources that limit performance. Every operation you run is bottlenecked by one or the other — never both simultaneously (at the ridge point, both bind equally, but this is rare in practice).

Compute is measured in FLOPS (floating-point operations per second). An H100 SXM5 delivers:

Memory bandwidth is measured in bytes per second. The H100 has HBM3 at 3.35 TB/s (that is 3,350 GB/s). This is the rate at which data can flow between the GPU's main memory (80 GB of HBM) and its compute units.

Think of it this way: the GPU is a factory. Compute is the number of workers on the factory floor. Bandwidth is the width of the loading dock. If you have 1,000 workers but a narrow loading dock, the workers stand idle waiting for materials. If you have a massive loading dock but only 10 workers, materials pile up unprocessed.

Arithmetic Intensity: The Single Number That Matters

Arithmetic intensity (AI) is the ratio of computation to data movement for a given operation:

The units are FLOPs/byte. This single number tells you whether your operation is compute-bound (limited by how fast the GPU can crunch numbers) or memory-bound (limited by how fast data can flow to the compute units).

Deriving the Roofline from First Principles

Let us derive the roofline model step by step. We want to find the maximum achievable FLOPS for an operation with arithmetic intensity AI.

This is the entire model. Below the ridge point (AI < 295), you are memory-bound and performance scales linearly with arithmetic intensity. Above it (AI > 295), you are compute-bound and performance is flat at the peak compute rate. Plot this on a log-log graph and you get the characteristic "roofline" shape: a diagonal line (memory-bound region) that hits a flat ceiling (compute-bound region).

Worked Example 1: Matrix Multiply on H100

Matrix multiplication is the most important operation in deep learning. Let us analyze it with the roofline model.

The takeaway: large matrix multiplications are compute-bound. This is why Tensor Cores matter — they are the fast lane for matrix multiplication. Small matrix multiplications (batch size 1, short sequences) can become memory-bound, which is exactly why LLM inference at batch=1 is often bandwidth-limited.

Worked Example 2: Self-Attention Roofline

Self-attention in a transformer is more nuanced because it involves multiple operations with very different arithmetic intensities. Let us dissect it.

The Three Bounds: Compute, Bandwidth, Communication

For multi-GPU training and inference, there is a third bound that the single-GPU roofline ignores: communication bandwidth. When your model is sharded across 8 GPUs with tensor parallelism, every attention layer requires an AllReduce of the output. The AllReduce transfers data proportional to the tensor size over NVLink or InfiniBand.

The three-way roofline analysis is what separates a staff engineer from a senior. A senior knows that attention is memory-bound. A staff engineer can tell you whether the AllReduce or the attention kernel itself is the bottleneck for a specific sharding configuration, and can redesign the parallelism strategy accordingly.

The Roofline Calculator

python
import dataclasses

# ── GPU specifications ──
@dataclasses.dataclass
class GPUSpec:
    name: str
    peak_tflops_fp16: float    # Tensor Core peak
    hbm_bandwidth_tbs: float  # HBM bandwidth in TB/s
    sram_per_sm_kb: float     # Shared memory per SM in KB
    num_sms: int
    nvlink_bw_gbs: float     # NVLink bandwidth per direction in GB/s

    @property
    def ridge_point(self) -> float:
        """Arithmetic intensity where compute = bandwidth bound."""
        return self.peak_tflops_fp16 * 1e12 / (self.hbm_bandwidth_tbs * 1e12)

    def achieved_tflops(self, arithmetic_intensity: float) -> float:
        """Roofline model: min(peak, AI * bandwidth)."""
        bandwidth_limited = arithmetic_intensity * self.hbm_bandwidth_tbs
        return min(self.peak_tflops_fp16, bandwidth_limited)

    def efficiency(self, arithmetic_intensity: float) -> float:
        """What fraction of peak compute we achieve."""
        return self.achieved_tflops(arithmetic_intensity) / self.peak_tflops_fp16

H100 = GPUSpec(
    name="H100 SXM5",
    peak_tflops_fp16=989,
    hbm_bandwidth_tbs=3.35,
    sram_per_sm_kb=228,
    num_sms=132,
    nvlink_bw_gbs=450,
)

A100 = GPUSpec(
    name="A100 SXM4",
    peak_tflops_fp16=312,
    hbm_bandwidth_tbs=2.0,
    sram_per_sm_kb=164,
    num_sms=108,
    nvlink_bw_gbs=300,
)

# ── Roofline analysis for common operations ──

def matmul_roofline(M, K, N, bytes_per_elem=2, gpu=H100):
    """Roofline analysis for C[M,N] = A[M,K] @ B[K,N]."""
    flops = 2 * M * K * N
    bytes_moved = bytes_per_elem * (M*K + K*N + M*N)
    ai = flops / bytes_moved
    tflops = gpu.achieved_tflops(ai)
    time_ms = flops / (tflops * 1e12) * 1e3
    bound = "COMPUTE" if ai > gpu.ridge_point else "MEMORY"
    print(f"Matmul [{M},{K}]x[{K},{N}] on {gpu.name}:")
    print(f"  FLOPs:      {flops:.2e}")
    print(f"  Bytes:      {bytes_moved:.2e}")
    print(f"  AI:         {ai:.1f} FLOPs/byte")
    print(f"  Ridge:      {gpu.ridge_point:.1f} FLOPs/byte")
    print(f"  Bound:      {bound}")
    print(f"  Achieved:   {tflops:.1f} TFLOPS ({gpu.efficiency(ai)*100:.1f}%)")
    print(f"  Time:       {time_ms:.3f} ms")
    return ai, tflops, time_ms

def attention_roofline(seq_len, d_head, n_heads, gpu=H100):
    """Roofline for standard (non-Flash) multi-head attention."""
    S, d, h = seq_len, d_head, n_heads
    # Per head: 2 matmuls (QK^T, PV) + softmax
    flops_per_head = 2 * (2 * S * S * d) + 5 * S * S
    total_flops = h * flops_per_head
    # Without FlashAttention: must materialize S×S attn matrix per head
    bytes_per_head = 2 * (2*S*d + S*S + S*S + S*d + S*d)  # Q,K,P,P_softmax,V,O
    total_bytes = h * bytes_per_head
    ai = total_flops / total_bytes
    tflops = gpu.achieved_tflops(ai)
    bound = "COMPUTE" if ai > gpu.ridge_point else "MEMORY"
    print(f"\nAttention S={S}, d={d}, h={h} on {gpu.name}:")
    print(f"  AI:       {ai:.1f} FLOPs/byte  ({bound})")
    print(f"  Achieved: {tflops:.1f} TFLOPS ({gpu.efficiency(ai)*100:.1f}%)")
    return ai

# ── Run examples ──
matmul_roofline(4096, 4096, 4096)      # Square matmul
matmul_roofline(1, 4096, 4096)         # Batch=1 inference
matmul_roofline(4096, 4096, 11008)    # LLaMA FFN up-projection

attention_roofline(2048, 128, 32)     # Standard LLM attention
attention_roofline(32768, 128, 32)    # Long-context attention

Interactive Roofline Plot

Predicting Real-World Throughput

The roofline gives you an upper bound. Real kernels fall below the roofline for several reasons:

A well-optimized kernel achieves 70-85% of the roofline prediction. If you are below 50%, look for one of the above issues. If you are above 85%, you are doing excellent work. Above 90% is exceptional and usually requires deep hardware-specific tuning.

Roofline for Different Operations

Here is a table of common deep learning operations and their typical arithmetic intensity. Memorize these — they come up constantly in interviews.

Chapter 2: GPU Memory Hierarchy

The roofline model told us that most operations in a transformer are memory-bound. The natural follow-up question: which memory? A GPU does not have one flat memory space. It has a hierarchy of four levels, each with vastly different bandwidth, latency, and capacity. Understanding this hierarchy is the foundation of all kernel optimization.

Think of the GPU memory hierarchy as a city's transportation system. Registers are the desk you are sitting at — instant access, but tiny (a few papers). Shared memory (SRAM) is the filing cabinet in your office — a few steps away, fast to access, holds a decent amount. L2 cache is the building's mailroom — takes a walk, much larger. HBM is the warehouse across town — takes a truck to deliver, but stores everything.

The Four Levels: Real Numbers on H100

Look at the bandwidth column. Registers are 23x faster than HBM. Shared memory is 5.7x faster than HBM. These are not small differences — they are order-of-magnitude gaps that completely determine kernel performance.

Why This Hierarchy Exists: Physics

The hierarchy is not an arbitrary engineering choice. It is dictated by physics.

Speed vs. capacity trade-off. Flip-flops (registers) toggle in a single clock cycle but require 6 transistors each. One 32-bit register = 192 transistors. To store 80 GB at register density would require trillions of transistors and consume absurd power. SRAM (shared memory) uses 6 transistors per bit but with a different layout — it is denser than registers but slower because the read circuitry adds latency. DRAM (the technology behind HBM) uses just 1 transistor + 1 capacitor per bit, making it ~6x denser than SRAM, but the capacitor must be refreshed periodically and the read is destructive (requiring precharge), adding ~400 cycles of latency.

Speed of light limits. An electrical signal travels about 15 cm per nanosecond on a chip. At 2 GHz clock speed, each cycle is 0.5 ns, so a signal can travel at most 7.5 cm in one cycle. The HBM stacks sit on the edge of the GPU die, about 2-3 cm from the SM array. Even at the speed of light, a round trip takes ~0.4 ns = ~1 cycle just for the physical propagation. Add the DRAM access protocol (row activate, column select, data transfer, precharge) and you get 400+ cycles.

The wire delay problem. On a modern chip, transistor switching is no longer the bottleneck — wire delay is. Moving a bit across a 1 mm wire takes longer than a transistor can switch. This is why local memory (registers, shared memory on the same SM) is fast and global memory (HBM on the edge of the die) is slow. Every level of the hierarchy represents a different physical distance from the compute units.

Registers: The Fastest Memory You Never Think About

Each thread on an H100 SM has access to up to 255 32-bit registers (the maximum register file allocation per thread). The total register file per SM is 256 KB. Registers are the only memory that can be accessed at full speed — zero additional latency, and the read happens in the same cycle as the instruction that uses the value.

Shared Memory (SRAM): The Programmer-Controlled Cache

Shared memory is the most important memory level for kernel developers because it is the one you explicitly control. Unlike L1/L2 caches (which are hardware-managed), you decide exactly what data goes into shared memory and when.

On H100, each SM has 228 KB of combined shared memory and L1 cache. You can configure the split: up to 228 KB all shared memory with 0 L1, or various splits in between. For most kernels, maximizing shared memory is the right choice because you get to control the access pattern perfectly.

Bank Conflicts: The Shared Memory Performance Trap

Shared memory's 32 banks are the source of its high bandwidth — but also its most common performance pitfall. When two threads in the same warp access different addresses that map to the same bank, the accesses serialize. This is called a bank conflict.

HBM: The Vast, Slow Main Memory

HBM (High Bandwidth Memory) is a stack of DRAM dies connected to the GPU die via silicon interposers. The H100 has six 16-GB HBM3 stacks, totaling 80 GB at 3.35 TB/s aggregate bandwidth.

Despite its name ("High Bandwidth Memory"), HBM is the slowest memory level. Its bandwidth is "high" relative to traditional DDR5 (which peaks at ~0.1 TB/s), but low relative to the GPU's compute capability.

How Tensor Cores Fit In

Tensor Cores are specialized matrix multiply units on each SM. They are the reason H100 achieves 989 TFLOPS for FP16 while standard CUDA cores top out at 67 TFLOPS. A single Tensor Core operation on H100 computes a 16×8×16 matrix multiply-accumulate in one cycle.

TMA: The Tensor Memory Accelerator (Hopper/Blackwell)

H100 introduced a dedicated hardware unit called the Tensor Memory Accelerator (TMA). TMA offloads data movement from the SM's compute pipeline to a separate engine. Instead of threads computing addresses and issuing loads, a single thread programs the TMA with a tensor descriptor (base pointer, shape, strides, data type), and the TMA fetches the entire tile asynchronously.

Interactive: GPU Memory Hierarchy

Worked Example: Where Does a Matmul Spend Its Time?

Let us trace a single matmul through the memory hierarchy to understand where time is spent.

Blackwell and Beyond: What Changes

Chapter 3: Transformer Math

In a scaling law meeting, someone asks: "How many FLOPs to train a 70B model on 2T tokens?" You need to answer in under 10 seconds. In a system-design interview, the question is: "How much memory does a 400B MoE model need at inference time with batch=64, seq_len=8192?" You need to derive it on the whiteboard in 2 minutes. This chapter gives you the formulas and worked examples to do both.

All of these are derived from first principles — no memorization required. Once you understand where each parameter lives and how each FLOP is spent, you can re-derive any formula on the spot.

Parameter Count: Deriving It Layer by Layer

A standard transformer decoder (GPT-style) has L identical layers, each containing a multi-head attention block and a feed-forward network (FFN). Let us count every parameter.

FLOPs Per Token: The 2P and 6P Rules

For every parameter, the forward pass performs approximately 2 FLOPs (one multiply and one add in a matrix multiplication). This gives us a beautifully simple rule:

Memory Footprint: Where Every Byte Goes

Training a model requires storing four categories of data in GPU memory. Let us account for each one.

KV Cache: The Inference Memory Bottleneck

During inference, the model generates one token at a time. For each new token, it needs to attend to all previous tokens. Re-computing K and V for all previous tokens at every step would be wasteful, so we cache them. This KV cache is often the largest memory consumer at inference time.

MFU: How to Measure Training Efficiency

Model FLOPs Utilization (MFU) is the fraction of the GPU's peak compute that is actually used for model computation:

Interactive: Transformer Parameter Calculator

Chinchilla Scaling Laws: How Much Data for How Many Parameters?

In 2022, the Chinchilla paper (Hoffmann et al.) established that for a given compute budget C, there is an optimal split between model size N and training data D. Training a model that is too large on too little data wastes compute; training a small model on excessive data also wastes compute.

Putting It All Together: LLaMA 3 70B Worksheet

Chapter 4: JAX Fundamentals [full lesson →]

At a frontier lab, the training infrastructure is almost certainly built on JAX. Not PyTorch — JAX. Why? Three reasons: (1) JAX compiles your Python code through XLA, which generates fused, optimized kernels for TPUs and GPUs automatically. (2) JAX's functional design (no hidden state, no in-place mutation) makes distributed training with model parallelism radically simpler — the compiler can shard pure functions across devices. (3) JAX's composable transforms (jit, grad, vmap) let you write research code that is simultaneously readable AND production-fast.

If you are coming from PyTorch, JAX will feel strange at first. PyTorch is imperative: you build a computation graph by executing Python statements, and the graph is traced eagerly. JAX is functional: you write pure functions, and JAX transforms (traces, compiles, differentiates) them. The payoff for this discipline is enormous — but the learning curve is real.

jax.numpy: Tracing, Not Executing

The first surprise: jax.numpy operations do not execute immediately (unless you are in eager mode). When you call a function decorated with @jax.jit, JAX traces the function: it runs your Python code with abstract "tracer" values instead of real data. The tracers record the operations, building an XLA computation graph. Then XLA compiles this graph into optimized machine code.

python
import jax
import jax.numpy as jnp

# ── Eager mode (no jit) ──
# Operations execute immediately, like NumPy.
# Useful for debugging, but slow (no fusion, no optimization).
x = jnp.ones((3, 3))
y = x + 1       # Dispatches immediately to GPU/TPU
z = y * 2       # Another dispatch — TWO kernel launches
print(z)        # [[4. 4. 4.], [4. 4. 4.], [4. 4. 4.]]

# ── JIT mode ──
# JAX traces the function, compiles it, then runs the compiled version.
@jax.jit
def f(x):
    y = x + 1   # NOT executed — recorded as a trace operation
    z = y * 2   # NOT executed — recorded
    return z    # The trace is: z = (x + 1) * 2

# First call: trace + compile + execute (slow)
result = f(x)   # XLA fuses add+mul into ONE kernel launch!

# Second call with same shape/dtype: just execute (fast, no recompile)
result2 = f(jnp.zeros((3, 3)))

# Third call with DIFFERENT shape: retrace + recompile!
result3 = f(jnp.zeros((4, 4)))  # New compilation

Pure Functions: Why No Side Effects Matter

JAX requires your functions to be pure: given the same inputs, they must always return the same outputs, and they must not modify any external state. This is not just a style preference — it is a fundamental requirement for correctness when JAX transforms your code.

python
# ── BAD: impure function ──
counter = [0]  # Mutable external state

@jax.jit
def bad_f(x):
    counter[0] += 1   # Side effect! Modifies external state
    print("traced")   # Side effect! Only prints during tracing
    return x + counter[0]

# Call 1: traces, prints "traced", counter becomes 1, returns x+1
# Call 2: DOES NOT retrace (same shape), DOES NOT print,
#         counter is STILL 1 (the value was captured at trace time),
#         returns x+1 (not x+2!)
# This is a silent correctness bug.

# ── GOOD: pure function ──
@jax.jit
def good_f(x, counter):
    return x + counter, counter + 1  # Return new counter, don't mutate

# State is threaded through as an explicit argument and return value.
# No hidden mutation, no trace-time capture bugs.
result, new_counter = good_f(x, jnp.array(0))
result2, new_counter2 = good_f(x, new_counter)

Why does purity matter beyond avoiding bugs? Because the XLA compiler assumes pure functions to perform optimizations. It can reorder operations, fuse kernels, and even eliminate dead code — all of which are only correct if the function has no side effects. The compiler also needs purity to shard the function across multiple devices: if a function modifies global state, you cannot safely run it on 2048 TPUs simultaneously.

jax.grad: Automatic Differentiation

jax.grad computes the gradient of a scalar-valued function with respect to its first argument (by default). It uses reverse-mode automatic differentiation (backpropagation), which computes all parameter gradients in a single backward pass.

python
import jax
import jax.numpy as jnp

# ── Basic gradient ──
def loss_fn(w, x, y):
    """MSE loss: L = mean((w @ x - y)^2)"""
    pred = w @ x
    return jnp.mean((pred - y) ** 2)

# grad takes the gradient w.r.t. the FIRST argument (w)
grad_fn = jax.grad(loss_fn)

w = jnp.ones((3, 4))
x = jnp.ones((4, 2))
y = jnp.zeros((3, 2))

dL_dw = grad_fn(w, x, y)  # Shape: (3, 4) — same as w
print(dL_dw.shape)  # (3, 4)

# ── Gradient w.r.t. multiple arguments ──
grad_fn_multi = jax.grad(loss_fn, argnums=(0, 1))  # grad w.r.t. w AND x
dL_dw, dL_dx = grad_fn_multi(w, x, y)

# ── Value and gradient together ──
# In training, you need both the loss value and its gradient.
# Computing them separately wastes work (shared forward pass).
loss_val, dL_dw = jax.value_and_grad(loss_fn)(w, x, y)
print(f"Loss: {loss_val:.4f}, Grad norm: {jnp.linalg.norm(dL_dw):.4f}")

Forward Mode vs Reverse Mode: When Each Shines

python
# ── Forward mode: Jacobian-Vector Product (JVP) ──
# Computes f(x) and J @ v simultaneously (J = Jacobian, v = tangent vector)
def f(x):
    return jnp.array([x[0]**2 + x[1], x[0] * x[1]])

x = jnp.array([3.0, 4.0])
v = jnp.array([1.0, 0.0])  # Tangent vector: derivative in x[0] direction

primals, tangents = jax.jvp(f, (x,), (v,))
# primals = f(x) = [9+4, 12] = [13, 12]
# tangents = J @ v = [2*x[0]*1 + 0, x[1]*1 + 0] = [6, 4]

# ── Reverse mode: Vector-Jacobian Product (VJP) ──
# Computes f(x) and a function that maps cotangent to gradient
primals, vjp_fn = jax.vjp(f, x)
# primals = f(x) = [13, 12]

# vjp_fn takes a cotangent vector (same shape as output)
# and returns the gradient w.r.t. the input
cotangent = jnp.array([1.0, 0.0])  # grad of output[0] only
(grad_x,) = vjp_fn(cotangent)
# grad_x = [2*x[0], 1] = [6, 1]  (df[0]/dx[0], df[0]/dx[1])

jax.vmap: Auto-Vectorization

jax.vmap automatically vectorizes a function that operates on a single example to work on a batch of examples. Instead of writing explicit batch loops or reshaping tensors, you write your function for one example and vmap handles batching.

python
# ── Without vmap: manual batching ──
def predict_single(params, x):
    """Forward pass for a single input x of shape (d,)."""
    for W, b in params:
        x = jax.nn.relu(W @ x + b)
    return x

# To handle a batch, you'd either:
# (a) Loop: [predict_single(params, x_i) for x_i in batch]  ← SLOW
# (b) Rewrite: change W @ x to x_batch @ W.T + b  ← manual work

# ── With vmap: automatic batching ──
predict_batch = jax.vmap(predict_single, in_axes=(None, 0))
# in_axes=(None, 0) means:
#   params: NOT batched (same params for all examples)
#   x: batched along axis 0 (first dim is batch)

# Now predict_batch(params, x_batch) works on a batch
# where x_batch has shape (batch_size, d)
# The output has shape (batch_size, output_d)

# ── vmap + grad: per-example gradients ──
# This is one of JAX's superpowers: get the gradient for EACH
# example in the batch, not just the average gradient.
# Useful for: differential privacy, per-example clipping, debugging.

def loss_single(params, x, y):
    pred = predict_single(params, x)
    return jnp.sum((pred - y) ** 2)

# Per-example gradient function:
per_example_grad = jax.vmap(jax.grad(loss_single), in_axes=(None, 0, 0))
# per_example_grad(params, x_batch, y_batch) returns a pytree
# with the same structure as params, but each leaf has an extra
# batch dimension: shape (batch_size, *param_shape)

Composing Transforms: The Power Move

JAX's transforms (jit, grad, vmap) compose freely. You can stack them in any order, and JAX will trace through all of them to produce a single optimized program.

python
# ── The classic training step ──
# This is the pattern you'll see in every JAX training loop.

@jax.jit                    # 3. Compile the whole thing
def train_step(params, x_batch, y_batch):
    def loss_fn(params):
        # vmap handles the batch dimension automatically
        preds = jax.vmap(predict_single, in_axes=(None, 0))(params, x_batch)  # 1. vmap for batching
        return jnp.mean((preds - y_batch) ** 2)

    loss, grads = jax.value_and_grad(loss_fn)(params)  # 2. grad for differentiation
    # Simple SGD update
    new_params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)
    return new_params, loss

# What happens under the hood:
# 1. jit traces train_step
# 2. Inside, it encounters vmap → adds batch tracing
# 3. Inside, it encounters grad → adds differentiation tracing
# 4. The entire thing (forward + backward + update) compiles
#    to a SINGLE XLA program with fused kernels
# 5. On subsequent calls: no Python overhead, just the compiled kernel

PyTrees: The Tree-of-Arrays Abstraction

Neural network parameters are not a single array — they are a nested structure of arrays (a dict of dicts of arrays, a list of tuples, etc.). JAX calls these nested structures PyTrees and provides utilities to work with them uniformly.

python
# ── Model parameters as a PyTree ──
params = {
    'layer1': {'W': jnp.ones((128, 64)), 'b': jnp.zeros(64)},
    'layer2': {'W': jnp.ones((64, 32)),  'b': jnp.zeros(32)},
    'head':   {'W': jnp.ones((32, 10)),  'b': jnp.zeros(10)},
}
# This is a PyTree: a nested dict of jax arrays.

# ── jax.tree.map: apply a function to every leaf ──
# SGD update: params = params - lr * grads
# grads has the SAME structure as params (JAX guarantees this)
new_params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)
# This walks both trees in parallel, applying the lambda to each pair.
# No manual iteration over layers!

# ── Counting parameters ──
n_params = sum(p.size for p in jax.tree.leaves(params))
print(f"Total params: {n_params:,}")  # 10,474

# ── Flatten and unflatten ──
flat_params, tree_def = jax.tree.flatten(params)
# flat_params is a list of arrays: [W1, b1, W2, b2, W3, b3]
# tree_def remembers the structure for reconstruction
restored = jax.tree.unflatten(tree_def, flat_params)
# restored == params (same structure and values)

# ── Custom PyTree nodes ──
# You can register your own classes as PyTree nodes.
# This is how Flax/Equinox/Haiku make their module objects
# work seamlessly with jax.grad, jit, vmap, etc.
import dataclasses

@dataclasses.dataclass
class Linear:
    W: jax.Array
    b: jax.Array

    def __call__(self, x):
        return self.W @ x + self.b

# Register with JAX so jit/grad/vmap know how to traverse it:
jax.tree_util.register_dataclass(Linear, data_fields=['W', 'b'], meta_fields=[])

Random Numbers: The Split Model

JAX's random number generation is different from NumPy's, and for good reason. NumPy uses a global, stateful PRNG: np.random.seed(42) followed by np.random.randn() advances a global counter. This is fundamentally incompatible with JAX's functional paradigm (no hidden state) and with parallel execution (what happens when two TPUs both try to advance the same counter?).

python
# ── JAX PRNG: explicit, functional, splittable ──

# Create an initial key (this is the ONLY global state)
key = jax.random.PRNGKey(42)
# key is a pair of uint32 values, not a global state
print(key)  # [ 0 42]

# WRONG: reusing the same key gives the SAME random numbers!
a = jax.random.normal(key, (3,))
b = jax.random.normal(key, (3,))
# a == b! This is intentional: same key = same result (purity)

# RIGHT: split the key to get new, independent subkeys
key, subkey1, subkey2 = jax.random.split(key, 3)
a = jax.random.normal(subkey1, (3,))
b = jax.random.normal(subkey2, (3,))
# a != b, and both are deterministic given the original key

# ── Pattern for training loops ──
def train_step(params, batch, key):
    # Split key for this step's randomness
    key, dropout_key, noise_key = jax.random.split(key, 3)

    def loss_fn(params):
        # Use dropout_key for dropout mask
        logits = model_forward(params, batch, dropout_key)
        return cross_entropy_loss(logits, batch['labels'])

    loss, grads = jax.value_and_grad(loss_fn)(params)
    params = update_params(params, grads)
    return params, loss, key  # Return the updated key!

# The key is threaded through the loop as explicit state.
# Every step gets a unique, deterministic key.
# The entire training run is perfectly reproducible.

# ── Why the split model matters for parallelism ──
# On 2048 TPUs, each device needs independent random numbers.
# With the split model, you split the key into 2048 subkeys
# at the start, one per device. No coordination needed.
keys = jax.random.split(jax.random.PRNGKey(42), 2048)
# Each device gets keys[device_id] — deterministic and independent.

Interactive: JAX Computation Graph Tracer

A Complete JAX Training Loop

Let us put it all together with a complete, production-style training loop for a small MLP. This is the pattern you will see (with more complexity) in frontier lab training codebases.

python
import jax
import jax.numpy as jnp
from typing import NamedTuple

# ── Model definition (pure functions + parameter PyTree) ──

class MLPParams(NamedTuple):
    """Parameters for a 2-layer MLP."""
    W1: jax.Array   # [d_in, d_hidden]
    b1: jax.Array   # [d_hidden]
    W2: jax.Array   # [d_hidden, d_out]
    b2: jax.Array   # [d_out]

def init_params(key, d_in, d_hidden, d_out):
    """Xavier initialization. Key is split for each layer."""
    k1, k2 = jax.random.split(key)
    scale1 = jnp.sqrt(2.0 / (d_in + d_hidden))
    scale2 = jnp.sqrt(2.0 / (d_hidden + d_out))
    return MLPParams(
        W1=scale1 * jax.random.normal(k1, (d_in, d_hidden)),
        b1=jnp.zeros(d_hidden),
        W2=scale2 * jax.random.normal(k2, (d_hidden, d_out)),
        b2=jnp.zeros(d_out),
    )

def forward(params: MLPParams, x):
    """Forward pass for a single example x of shape (d_in,)."""
    h = jax.nn.silu(x @ params.W1 + params.b1)  # SiLU activation (like LLaMA)
    return h @ params.W2 + params.b2

# Batched forward: vmap over the data dimension, not over params
forward_batch = jax.vmap(forward, in_axes=(None, 0))

# ── Loss function ──
def loss_fn(params, x_batch, y_batch):
    """Cross-entropy loss for classification."""
    logits = forward_batch(params, x_batch)  # [B, n_classes]
    # Numerically stable log-softmax
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    # Select the log-prob of the correct class
    loss = -jnp.mean(log_probs[jnp.arange(y_batch.shape[0]), y_batch])
    return loss

# ── Optimizer state (AdamW, simplified) ──
class AdamState(NamedTuple):
    m: MLPParams    # First moment (same tree structure as params)
    v: MLPParams    # Second moment
    step: jax.Array # Step counter for bias correction

def init_adam(params, lr=1e-3, beta1=0.9, beta2=0.999, wd=0.01):
    m = jax.tree.map(jnp.zeros_like, params)
    v = jax.tree.map(jnp.zeros_like, params)
    return AdamState(m=m, v=v, step=jnp.array(0))

def adam_update(params, grads, state, lr=1e-3, beta1=0.9, beta2=0.999, wd=0.01, eps=1e-8):
    step = state.step + 1
    # Update moments
    new_m = jax.tree.map(lambda m, g: beta1 * m + (1 - beta1) * g, state.m, grads)
    new_v = jax.tree.map(lambda v, g: beta2 * v + (1 - beta2) * g**2, state.v, grads)
    # Bias correction
    bc1 = 1 - beta1 ** step
    bc2 = 1 - beta2 ** step
    m_hat = jax.tree.map(lambda m: m / bc1, new_m)
    v_hat = jax.tree.map(lambda v: v / bc2, new_v)
    # AdamW: weight decay applied to params directly, not through gradient
    new_params = jax.tree.map(
        lambda p, m, v: p * (1 - lr * wd) - lr * m / (jnp.sqrt(v) + eps),
        params, m_hat, v_hat
    )
    return new_params, AdamState(m=new_m, v=new_v, step=step)

# ── Training step (compiled) ──
@jax.jit
def train_step(params, opt_state, x_batch, y_batch):
    loss, grads = jax.value_and_grad(loss_fn)(params, x_batch, y_batch)
    # Gradient clipping (max norm = 1.0)
    grad_norm = jnp.sqrt(sum(
        jnp.sum(g**2) for g in jax.tree.leaves(grads)
    ))
    clip_coeff = jnp.minimum(1.0, 1.0 / (grad_norm + 1e-6))
    grads = jax.tree.map(lambda g: g * clip_coeff, grads)
    # Update params
    new_params, new_opt_state = adam_update(params, grads, opt_state)
    return new_params, new_opt_state, loss, grad_norm

# ── Training loop ──
key = jax.random.PRNGKey(42)
key, init_key = jax.random.split(key)
params = init_params(init_key, d_in=784, d_hidden=256, d_out=10)
opt_state = init_adam(params)

for step in range(1000):
    key, data_key = jax.random.split(key)
    # (In real code, load from data pipeline)
    x_batch = jax.random.normal(data_key, (64, 784))
    y_batch = jax.random.randint(data_key, (64,), 0, 10)

    params, opt_state, loss, grad_norm = train_step(
        params, opt_state, x_batch, y_batch
    )
    if step % 100 == 0:
        print(f"Step {step}: loss={loss:.4f}, grad_norm={grad_norm:.4f}")

Common JAX Gotchas

Chapter 5: Quantization — The Full Landscape

It is 10:15 AM. The inference team Slack lights up: the new 70B model fits on a single 80 GB A100 in FP16 — barely. At 140 GB of weights, it overflows. Two options: buy a second GPU (double the cost, double the latency from tensor parallelism), or quantize. Shrink the weights from 16 bits to 4 bits, and the model fits with room to spare for the KV cache. But will the quality hold?

This is not a theoretical exercise. Every frontier lab ships quantized models. The question is how many bits and which method. Get it wrong and your chatbot starts hallucinating. Get it right and you serve 4x more users on the same hardware.

Why Memory Bandwidth, Not Compute

Here is the single most important insight in LLM inference: autoregressive decoding is memory-bandwidth-bound, not compute-bound. During generation, each token requires reading the entire weight matrix once (to compute one output vector), but the arithmetic intensity is tiny — you multiply a matrix by a single vector, not a matrix by a matrix.

Let us make this concrete. An A100 GPU has:

For a matrix-vector multiply W · x where W is (d, d), x is (d, 1): you read d² · 2 bytes of weights (FP16) and perform 2d² FLOPs. The arithmetic intensity is:

One. One FLOP per byte. The A100 needs 156 to be compute-bound. You are 156x below the threshold. The GPU spends virtually all its time waiting for weights to arrive from HBM. This means: if you halve the weight size, you nearly double generation speed. Not because you do less math — because you move less data.

The Basics: How Quantization Works

Quantization maps continuous (or high-precision) values to a smaller set of discrete levels. A FP16 weight can represent ~65,536 distinct values. An INT4 weight can represent only 16. The art is choosing which 16 values best represent the original distribution.

Symmetric quantization maps the range [-α, α] to integer range [-2^b-1, 2^b-1 - 1], using a single scale factor:

Dequantization is simply x̂ = x_q · s. This is fast because there is no zero-point offset, but wasteful when the distribution is asymmetric (e.g., mostly positive activations after ReLU).

Asymmetric quantization adds a zero-point z to handle shifted distributions:

Dequantization: x̂ = (x_q - z) · s. More accurate for skewed distributions, but the extra subtraction costs a few cycles per element.

Granularity: Per-Tensor vs Per-Channel vs Per-Group

The choice of how many weights share a single (s, z) pair dramatically affects accuracy:

Per-group quantization with group size g = 128 means that for a weight matrix of shape (4096, 4096), you store 4096 · (4096/128) = 131,072 scale/zero-point pairs. At FP16 each, that is 0.5 MB of overhead versus the 32 MB of quantized weights. Tiny price for dramatically better accuracy.

LLM.int8(): The Outlier Discovery [full paper →]

In 2022, Tim Dettmers made a pivotal observation: transformer weights have emergent outlier features. A tiny fraction of hidden dimensions (as few as 6 out of 4096) have magnitudes 10-100x larger than the rest. These outliers appear in every layer, in the same feature dimensions, and they are critical for model quality. Quantize them to INT8 and the model breaks. Leave the rest in INT8 and the model is fine.

LLM.int8() implements mixed-precision decomposition:

python
import torch

def llm_int8_matmul(X, W, threshold=6.0):
    """Mixed-precision decomposition from Dettmers et al. 2022."""
    # Step 1: find outlier feature dimensions
    outlier_mask = X.abs().max(dim=0).values > threshold  # shape: (d_in,)
    normal_mask = ~outlier_mask

    # Step 2: decompose into outlier and normal parts
    X_outlier = X[:, outlier_mask].half()   # (batch, d_outlier) in FP16
    W_outlier = W[outlier_mask, :].half()   # (d_outlier, d_out)  in FP16

    X_normal = X[:, normal_mask]            # (batch, d_normal)
    W_normal = W[normal_mask, :]            # (d_normal, d_out)

    # Step 3: quantize the normal part to INT8
    sx = X_normal.abs().max(dim=1, keepdim=True).values / 127
    sw = W_normal.abs().max(dim=0, keepdim=True).values / 127
    X_int8 = (X_normal / sx).round().to(torch.int8)
    W_int8 = (W_normal / sw).round().to(torch.int8)

    # Step 4: compute both parts
    out_normal = (X_int8.float() @ W_int8.float()) * (sx * sw)  # dequantize
    out_outlier = X_outlier @ W_outlier                          # FP16 matmul

    return (out_normal + out_outlier).half()

The Chris De Sa Quantization Series: QuiP → QuiP# → QTIP

While LLM.int8() proved mixed-precision works at 8 bits, a group at Cornell led by Chris De Sa asked a bolder question: can we go to 2 bits per weight without destroying the model? The answer required ideas from information theory and lattice coding.

QuiP (2023) introduced two key ideas:

1. Incoherence processing. The problem with naive quantization is that some weight columns have much larger norms than others. If you apply the same quantization grid to all columns, the high-norm columns suffer. QuiP multiplies the weight matrix by random orthogonal matrices U and V: W̃ = U · W · V^T. This "spreads" the information evenly across all entries (makes the matrix incoherent), so uniform quantization works better. You store U and V as random seeds (cheap) and undo the rotation during inference.

2. LDLQ (Lattice-based Data-aware Quantization). Instead of rounding each weight independently (which ignores correlations), LDLQ uses the inverse Hessian H^-1 = (X^TX)^-1 to quantize weights in an order that accounts for their mutual information. The Hessian tells you which weight errors are cheap (dimensions the model rarely uses) and which are expensive (dimensions on which loss is sensitive). LDLQ performs LDL^T decomposition of the Hessian and uses it to guide a greedy per-column rounding scheme similar to GPTQ.

This is a weighted quantization error where H = 2X^TX is the Hessian of the layer's squared error. LDLQ solves it approximately via the LDL^T factorization of H, processing one row at a time with adaptive rounding.

QuiP# (2024) made two improvements that turned QuiP from a proof of concept into a practical method:

1. Hadamard rotations instead of random orthogonal matrices. A random orthogonal matrix multiply costs O(d²). A Hadamard transform costs O(d log d) and achieves the same incoherence property. This is because Hadamard matrices have the special property that all entries have the same magnitude (1/√d), which is exactly the definition of maximal incoherence. QuiP# replaces the expensive random rotations with fast Walsh-Hadamard transforms — a huge speedup at inference time.

2. E8 lattice codebook. Instead of rounding to a uniform integer grid, QuiP# rounds to the nearest point in the E8 lattice — the densest known sphere packing in 8 dimensions. Why does this matter? The quantization error of a codebook depends on how efficiently it tiles space. A uniform grid in 8D wastes space in the corners of each hypercube. The E8 lattice packs more representable points into the same volume, reducing the average distance from any real weight vector to its nearest codebook entry. At 2 bits per weight, QuiP# with E8 lattice achieves perplexity close to what older methods achieve at 3 bits.

python
# Conceptual: E8 lattice quantization in 8D
# The E8 lattice consists of all vectors in R^8 whose coordinates
# are either all integers or all half-integers, and whose sum is even.

import numpy as np

def nearest_e8(x):
    """Find the nearest E8 lattice point to vector x in R^8."""
    # Two candidates: round to integers, or round to half-integers
    c1 = np.round(x)                           # nearest integers
    if np.sum(c1) % 2 != 0:                     # enforce even sum
        idx = np.argmin(np.abs(x - c1))         # flip the closest-to-boundary entry
        c1[idx] += 1 if x[idx] > c1[idx] else -1

    c2 = np.round(x - 0.5) + 0.5              # nearest half-integers
    if np.sum(c2) % 2 != 0:
        idx = np.argmin(np.abs(x - c2))
        c2[idx] += 1 if x[idx] > c2[idx] else -1

    # Pick the closer candidate
    d1 = np.sum((x - c1)**2)
    d2 = np.sum((x - c2)**2)
    return c1 if d1 <= d2 else c2

QTIP (2024) pushed even further. Instead of the E8 lattice (which is optimal in 8D but doesn't generalize easily), QTIP uses trellis codes — a technique borrowed from telecommunications. A trellis code defines a state machine where each transition corresponds to a quantized symbol. The Viterbi algorithm finds the sequence of quantized weights that minimizes total error while respecting the code constraints. This is analogous to how your cell phone decodes noisy signals, but applied to weight quantization.

The key advantage of trellis codes: they naturally capture sequential dependencies between adjacent weights. If weight[i] was rounded up, the trellis can compensate by rounding weight[i+1] down. This error diffusion is automatic and optimal (in the Viterbi sense).

AQLM: Additive Codebook Quantization [full paper →]

AQLM (2024) takes a completely different approach: multi-codebook additive quantization. Instead of quantizing each weight independently, AQLM groups weights into vectors of size g (typically 8) and represents each vector as the sum of entries from M learned codebooks:

Each codebook has 2^B entries (e.g., B=8 gives 256 entries per codebook). With M=2 codebooks, the total number of representable vectors is 256 × 256 = 65,536 — far more than the 16 levels of straight INT4. The indices (i₁, i₂) are stored as integers, and the codebooks are learned via the beam search variant of product quantization with calibration data.

The effective bit rate is (M · B) / g bits per weight. With M=2, B=8, g=8: (2 · 8) / 8 = 2 bits per weight. But the representational power is vastly greater than naive 2-bit quantization because the codebook entries are learned from the actual weight distribution.

python
# AQLM dequantization (simplified)
def aqlm_dequant(indices, codebooks, group_size=8):
    """
    indices: (n_groups, M) tensor of codebook indices
    codebooks: list of M tensors, each (2^B, group_size)
    Returns: (n_groups * group_size,) dequantized weights
    """
    M = len(codebooks)
    n_groups = indices.shape[0]
    result = torch.zeros(n_groups, group_size)
    for m in range(M):
        result += codebooks[m][indices[:, m]]  # additive: sum across codebooks
    return result.reshape(-1)

The Progression: 8-bit to 2-bit

Each paper in the quantization lineage solved a specific limitation of its predecessor:

Practical Decision Tree

When a new model lands on your desk and you need to deploy it, here is the decision process a frontier lab engineer follows:

Implementation: Quantizing a Weight Matrix

python
import torch
import torch.nn.functional as F

def quantize_per_group(W, bits=4, group_size=128):
    """
    Per-group asymmetric quantization.
    W: (out_features, in_features) weight tensor
    Returns: quantized weights, scales, zero_points
    """
    out_f, in_f = W.shape
    assert in_f % group_size == 0, "in_features must divide group_size"
    n_groups = in_f // group_size

    # Reshape to (out_features * n_groups, group_size)
    W_grouped = W.reshape(out_f * n_groups, group_size)

    # Compute per-group min/max
    w_min = W_grouped.min(dim=1).values  # (out_f * n_groups,)
    w_max = W_grouped.max(dim=1).values

    # Compute scale and zero-point
    q_max = 2 ** bits - 1
    scale = (w_max - w_min) / q_max
    scale = torch.clamp(scale, min=1e-8)  # avoid division by zero
    zero_point = torch.round(-w_min / scale).clamp(0, q_max).to(torch.int8)

    # Quantize
    W_q = torch.round(W_grouped / scale.unsqueeze(1)) + zero_point.unsqueeze(1)
    W_q = W_q.clamp(0, q_max).to(torch.uint8)

    # Verify: dequantize and check error
    W_deq = (W_q.float() - zero_point.unsqueeze(1).float()) * scale.unsqueeze(1)
    error = (W_grouped - W_deq).pow(2).mean().sqrt()
    print(f"RMSE: {error:.6f}, relative: {error / W_grouped.abs().mean():.4f}")

    return W_q.reshape(out_f, in_f), scale.reshape(out_f, n_groups), zero_point.reshape(out_f, n_groups)

# Example: quantize a 4096x4096 matrix to 4-bit with group_size=128
W = torch.randn(4096, 4096) * 0.02  # typical init scale
# Inject outlier features (simulating the LLM.int8() discovery)
W[:, 42] *= 15   # feature 42 is an outlier
W[:, 1337] *= 12 # feature 1337 is an outlier
W_q, scales, zps = quantize_per_group(W, bits=4, group_size=128)
# RMSE: ~0.0029, relative: ~0.18
# Memory: 4096*4096*0.5 bytes (4-bit) + scales + zps ≈ 8.4 MB vs 33.6 MB (FP16)

Chapter 6: Attention Kernels — FlashAttention 1 → 4 [attention lesson]

It is 11:30 AM. You are reviewing a profiling trace from the training cluster. Attention layers consume 42% of wall-clock time and, more importantly, 68% of peak HBM usage. The model processes sequences of length 8192, which means the attention matrix is 8192 × 8192 = 67 million entries per head, per layer. With 32 heads across 80 layers, that is 171 billion intermediate values materialized in memory. Your hardware budget says you need 32K context. Quadrupling the sequence length quadruples memory but sixteen-tuples the attention matrix. Something has to give.

This is the problem FlashAttention solves. Not by changing the math — the output is bit-for-bit identical to standard attention — but by changing how the math is executed on the hardware.

Standard Attention: The Memory Materialization Problem

The textbook attention computation proceeds in three steps:

Where Q, K, V have shape (N, d) — N tokens, d head dimension (typically 64 or 128). The problem is the intermediate matrix S (and P), which has shape (N, N). Let us trace the memory access pattern:

python
# Standard attention — what actually happens on the GPU

# Step 1: S = Q @ K.T / sqrt(d)
# Read Q (N×d) from HBM, read K (N×d) from HBM
# Compute matmul (fast on tensor cores)
# Write S (N×N) to HBM  ← THE BOTTLENECK
S = Q @ K.T / math.sqrt(d)  # O(N²d) FLOPs, O(N²) HBM writes

# Step 2: P = softmax(S, dim=-1)
# Read S (N×N) from HBM
# Compute softmax (element-wise, cheap)
# Write P (N×N) to HBM
P = torch.softmax(S, dim=-1)  # O(N²) FLOPs, O(N²) HBM reads + writes

# Step 3: O = P @ V
# Read P (N×N) from HBM, read V (N×d) from HBM
# Compute matmul
# Write O (N×d) to HBM
O = P @ V  # O(N²d) FLOPs, O(N²) HBM reads

Total HBM accesses: O(N²) reads and writes for S and P. For N = 8192, d = 128: the compute is 8192² × 128 × 2 ≈ 17.2 billion FLOPs (trivial for a GPU). But reading and writing S and P requires 8192² × 2 bytes × 2 (read + write) × 2 (for S and P) ≈ 1 GB of memory traffic. Per head. Per layer. This is the problem.

FlashAttention-1 (2022): Tiling + Online Softmax

Tri Dao's key insight: never materialize the N × N matrix. Instead, compute attention in tiles that fit in SRAM (the GPU's fast on-chip memory, ~20 MB on A100 vs 80 GB HBM). Each tile computes a partial result and accumulates it, without ever writing the full S or P to HBM.

The challenge: softmax is a non-decomposable operation. To compute softmax(S[i, :]), you need max(S[i, :]) and sum(exp(S[i, :] - max)), which requires the entire row. If you are processing the row in blocks, you do not have the full row available. You need the online softmax trick.

Deriving Online Softmax

Standard softmax of row s = [s₁, ..., s_N]:

Now suppose you have processed the first block of B columns and computed:

Then you process the second block and find a new maximum m⁽²⁾ = max(m⁽¹⁾, max(s_B+1, ..., s_2B)). The old partial sum ℓ⁽¹⁾ was computed with respect to m⁽¹⁾, not m⁽²⁾. We need to rescale:

And the partial output O⁽¹⁾ (which was already multiplied by the old softmax weights) must also be rescaled:

This is the online softmax algorithm: maintain a running (m, ℓ, O) triple, and after each new block, rescale the accumulated result by exp(m_old - m_new). The final result is mathematically identical to computing softmax on the full row.

python
def flash_attention_forward(Q, K, V, block_size=64):
    """
    Simplified FlashAttention-1 forward pass.
    Q, K, V: (N, d) tensors
    Returns O: (N, d) — same as standard attention
    """
    N, d = Q.shape
    O = torch.zeros(N, d)
    m = torch.full((N,), float('-inf'))  # running max per query
    l = torch.zeros(N)                    # running sum per query

    # Iterate over KV blocks
    for j in range(0, N, block_size):
        j_end = min(j + block_size, N)
        K_block = K[j:j_end]  # (B, d)
        V_block = V[j:j_end]  # (B, d)

        # Iterate over Q blocks (in practice, both loops are tiled)
        for i in range(0, N, block_size):
            i_end = min(i + block_size, N)
            Q_block = Q[i:i_end]  # (B, d)

            # Compute tile of S = Q_block @ K_block.T / sqrt(d)
            S_tile = Q_block @ K_block.T / (d ** 0.5)  # (B, B) — fits in SRAM!

            # Online softmax: update running stats
            m_new = torch.max(m[i:i_end], S_tile.max(dim=-1).values)
            correction = torch.exp(m[i:i_end] - m_new)
            P_tile = torch.exp(S_tile - m_new.unsqueeze(-1))
            l_new = l[i:i_end] * correction + P_tile.sum(dim=-1)

            # Rescale accumulated output and add new contribution
            O[i:i_end] = O[i:i_end] * correction.unsqueeze(-1) + P_tile @ V_block

            # Update running stats
            m[i:i_end] = m_new
            l[i:i_end] = l_new

    # Final normalization
    O = O / l.unsqueeze(-1)
    return O

IO Complexity: Why FlashAttention Is Faster

Let us be precise about the IO complexity. Define M = SRAM size in elements.

Same FLOPs. Dramatically less memory traffic. And the intermediate memory drops from O(N²) to O(N) — this is what enables 32K, 64K, 128K context lengths.

FlashAttention-2 (2023): Better Parallelism

FlashAttention-1 was IO-optimal but left GPU SMs underutilized. FA-2 made two key changes:

1. Swap the loop order. FA-1 iterates over KV blocks in the outer loop and Q blocks in the inner loop. FA-2 reverses this: iterate over Q blocks in the outer loop. Why? Each Q block now runs on a separate thread block (SM), and the inner KV loop is serial within that block. This means the number of thread blocks = N / B_Q, which is much larger than the number of SMs, giving the GPU scheduler plenty to work with.

2. Reduce non-matmul FLOPs. On modern GPUs (A100, H100), matmul operations run on specialized tensor cores at 10-16x the throughput of general-purpose FP32 ops. The rescaling in online softmax uses exp() and division, which run on the slower CUDA cores. FA-2 carefully restructures the computation to minimize these non-matmul operations, achieving ~70% of theoretical peak matmul throughput (vs ~35% for FA-1).

3. Work partitioning within warps. FA-2 splits Q among warps but shares K/V across warps. This minimizes communication between warps (no cross-warp reduction needed for the softmax statistics) while maximizing the reuse of K/V from shared memory.

FlashAttention-3 (2024): Hopper Architecture

The H100 Hopper GPU introduced three hardware features that FA-3 exploits:

Warp specialization. FA-3 assigns different roles to different warp groups: "producer" warps issue TMA loads, "consumer" warps issue WGMMA instructions. They communicate through shared memory barriers. This is a software pipeline: while consumers compute on data that producers loaded in the previous step, producers are simultaneously loading the next batch. The result: near-zero memory latency hiding.

pseudocode
// FA-3 warp-specialized pipeline (conceptual)
// Two warp groups: PRODUCER and CONSUMER

// PRODUCER warp group:
for tile in kv_tiles:
    TMA_async_load(K_tile[next], smem_K[next_buf])   // hardware-accelerated
    TMA_async_load(V_tile[next], smem_V[next_buf])
    signal(barrier[next_buf])                          // tell consumer data is ready
    wait(barrier[curr_buf])                            // wait for consumer to finish
    swap(curr_buf, next_buf)

// CONSUMER warp group:
for tile in kv_tiles:
    wait(barrier[curr_buf])                            // wait for producer to load
    S = WGMMA(Q_reg, smem_K[curr_buf])                // 128×128×64 matmul
    P = online_softmax(S, &m, &ℓ)                   // on CUDA cores
    O += WGMMA(P, smem_V[curr_buf])                   // accumulate output
    signal(barrier[curr_buf])                          // tell producer buffer is free
    swap(curr_buf, next_buf)

FlashAttention-4 (2025): Blackwell Architecture [full paper →]

The B200 Blackwell GPU doubles down on the trends from Hopper, and FA-4 exploits new capabilities:

CuTe DSL. FA-4 is written using NVIDIA's CuTe (Cute Templates) abstraction, which represents tiles as first-class objects with layout algebra. Instead of manually indexing into shared memory with pointer arithmetic, you declare a tile's layout (e.g., "128 rows × 64 cols, row-major, swizzled by 128 bytes") and CuTe generates the correct indexing code. This eliminates an entire class of shared memory bank conflict bugs.

Pingpong scheduling. Blackwell has even larger shared memory (228 KB per SM vs 228 KB on Hopper, but with a new partitioned structure). FA-4 uses a "pingpong" pattern where two warp groups alternate between computing and loading, with triple-buffering to keep both the compute pipeline and the memory pipeline fully saturated at all times.

FP4 support. Blackwell's tensor cores natively support FP4 matmul. FA-4 can run the Q×K^T multiply in FP4 (with FP32 accumulation), achieving 4x the throughput of FP16 at the cost of some precision. For inference with quantized models, this is a game-changer.

The Evolution in One Table

Backward Pass: Why It's Harder

The forward pass avoids materializing S and P, but the backward pass needs them for the gradient computation. FA-1 and FA-2 solve this by recomputing S and P from Q, K, V during the backward pass, which trades extra FLOPs for HBM savings. Since the forward pass is memory-bound (extra FLOPs are free), this recomputation is essentially costless. FA-3 further optimizes the backward by keeping a small auxiliary tensor (the log-sum-exp per row, size N) from the forward pass to avoid some redundant work.

Practical Integration

python
# Using FlashAttention in practice (PyTorch 2.0+)
import torch
import torch.nn.functional as F

# Option 1: PyTorch native SDPA (auto-selects FlashAttention when possible)
with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

# Option 2: Direct flash-attn library
from flash_attn import flash_attn_func

# Q, K, V shape: (batch, seqlen, nheads, headdim)
out = flash_attn_func(Q, K, V, causal=True)

# Option 3: FA-3 with FP8 on H100
out = flash_attn_func(
    Q.to(torch.float8_e4m3fn),   # FP8 E4M3 format
    K.to(torch.float8_e4m3fn),
    V.to(torch.float8_e4m3fn),
    causal=True,
    softmax_scale=1.0 / (d ** 0.5),
)

# Memory comparison for sequence length 32768, 32 heads, dim 128:
# Standard:  32768² × 32 × 2 bytes = 64 GB (!) — doesn't fit on any GPU
# FlashAttn: 32768 × 32 × 4 bytes  = 4 MB  (just the L vector)

Chapter 7: Kernel DSLs — The 2025 Explosion [full lesson →]

It is 1:15 PM, back from lunch. A Slack thread is debating which framework to use for a custom fused kernel: someone advocates Triton, another insists on CUTLASS, a third says ThunderKittens is the only way to get 80%+ utilization on Hopper. The team lead posts a screenshot of a timeline from an AI infrastructure newsletter showing eleven new kernel DSLs released in the past twelve months. "We need a decision by Friday," she writes.

This explosion is not arbitrary. It reflects a fundamental tension in GPU programming that every frontier lab grapples with: productivity vs performance. Writing CUDA by hand gives you total control over shared memory layout, warp scheduling, and instruction selection. But a single fused attention kernel can take 2,000+ lines of CUDA and weeks to debug. Triton gives you 100 lines of Python for the same kernel, but the compiler might leave 30% of the hardware idle because it made suboptimal scheduling choices. Every DSL in the 2025 explosion stakes out a different position on this tradeoff.

The Productivity-Performance Spectrum

Let us map the landscape. At one end: raw CUDA/PTX, where you control every register and every instruction. At the other: PyTorch's autograd, where you write math and hope the compiler figures out the rest.

CUDA: The Baseline

Every kernel ultimately compiles to CUDA (or rather, to PTX and then SASS). Understanding the CUDA programming model is non-negotiable for a frontier lab engineer, even if you rarely write raw CUDA yourself.

The key abstractions: threads execute scalar operations. Warps (32 threads) execute in lockstep (SIMT). Thread blocks (up to 1024 threads) share SRAM. Grids of thread blocks fill the GPU. The programmer must manually manage:

cuda
// CUDA GEMM kernel — simplified, no tiling optimizations
// This is what you're trying to AVOID writing by using a DSL
__global__ void matmul_naive(float* A, float* B, float* C,
                             int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}
// Performance: ~5% of peak on A100. Why? No shared memory,
// no tiling, no tensor core usage, terrible memory access pattern.
// A production GEMM kernel is 50-100x more complex.

Triton: Python-Level Kernel Writing

Triton (by OpenAI, first released 2021, mature by 2023) lets you write GPU kernels in Python. The key abstraction is the block program: your code operates on blocks (tiles) of data, and the Triton compiler handles the mapping to warps, shared memory, and tensor cores.

python
import triton
import triton.language as tl

@triton.jit
def matmul_kernel(
    A_ptr, B_ptr, C_ptr,
    M, N, K,
    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    # Which tile of C does this program compute?
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    # Compute offsets for this tile
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    # Pointers to first tile of A and B
    a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = B_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn

    # Accumulate in FP32 for numerical stability
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # Loop over K dimension in blocks
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k)
        acc += tl.dot(a, b)  # Triton automatically uses tensor cores!
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    # Write result
    c_ptrs = C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

What the compiler does for you: shared memory allocation, double-buffering (loading the next tile while computing the current one), warp scheduling, tensor core instruction selection, predicated loads for boundary tiles. What it does NOT do well: cross-SM communication, complex warp-level synchronization, fine-grained instruction scheduling (which is why hand-tuned CUDA can still beat Triton by 10-20% on specific kernels).

CUTLASS / CuTe: C++ Template Metaprogramming

CUTLASS (NVIDIA) is a C++ template library for writing high-performance GPU kernels. Its modern incarnation, CuTe, introduces layout algebra — a type system for describing how data is arranged in memory. A layout is a function from logical coordinates to physical memory addresses.

c++
// CuTe layout example: a 128×64 tile in shared memory
// with 128-byte swizzle to avoid bank conflicts
using SmemLayout = Layout,
                          Stride<_64, _1>>;

// A swizzled version: XOR the row index with bytes-per-row
// This eliminates bank conflicts for column-major access
using SmemLayoutSwizzled = composition(
    Swizzle<3, 3, 3>{},  // 8-byte swizzle pattern
    SmemLayout{}
);

// The layout algebra handles indexing automatically:
// smem(row, col) returns the correct physical address
// even with the swizzle — no manual index math needed.

CUTLASS is what NVIDIA uses internally for cuBLAS. If you need absolute peak performance on NVIDIA hardware and are willing to invest in C++ template wizardry, it is the tool.

ThunderKittens: Register Tiles as First-Class Citizens [full paper →]

ThunderKittens (Stanford, 2024) takes a radical position: the right abstraction level is the register tile — a 16×16 matrix that lives in a warp's registers and maps directly to a single tensor core MMA instruction.

c++
// ThunderKittens: the entire abstraction in a few types
using namespace kittens;

// Register tiles: 16×16 matrices that live in warp registers
rt_fl<16, 16> a_tile, b_tile, c_tile;     // float register tiles
rt_bf<16, 16> a_bf, b_bf;                  // bfloat16 register tiles

// Shared memory tiles
st_bf<64, 64> smem_a, smem_b;              // shared memory tiles

// Operations map directly to hardware
mma_ABt(c_tile, a_bf, b_bf, c_tile);      // tensor core MMA: C += A @ B^T
load(a_bf, smem_a, row_offset, col_offset);// load from shared to register
store(smem_a, a_bf, row_offset, col_offset);// store from register to shared

The philosophy: "If it's not a 16×16 tile, you're not thinking at the right level." ThunderKittens forces you to decompose your algorithm into register-tile operations, which guarantees that the tensor cores are always engaged. The downside: you still manage shared memory layout and warp scheduling manually.

The 2025-2026 Explosion

Starting in mid-2024, a wave of new DSLs appeared, each targeting a specific niche in the productivity-performance space:

When to Use What: The Decision Tree

The Same Matmul, Three Ways

To make the tradeoffs concrete, here is the conceptual core of a tiled GEMM in three DSLs. All three compute C += A × B in tiles.

python — triton
# Triton: ~50 lines, compiler handles shared memory + scheduling
@triton.jit
def gemm(A, B, C, M, N, K, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m, pid_n = tl.program_id(0), tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for k in range(0, K, BK):
        rk = k + tl.arange(0, BK)
        a = tl.load(A + rm[:, None] * K + rk[None, :])
        b = tl.load(B + rk[:, None] * N + rn[None, :])
        acc += tl.dot(a, b)
    tl.store(C + rm[:, None] * N + rn[None, :], acc)

c++ — cutlass cute
// CuTe: ~200 lines, explicit layout algebra, explicit pipeline
// (heavily abbreviated — real code has copy engines, pipeline barriers)
auto tiled_mma = make_tiled_mma(SM80_16x8x16_F32BF16BF16F32{},
                                Layout>{});
auto smem_a = make_tensor(smem_ptr, SmemLayoutA{});
auto smem_b = make_tensor(smem_ptr + smem_a_size, SmemLayoutB{});
for (int k = 0; k < K; k += BK) {
    copy(gmem_a_slice, smem_a);  // global → shared (via TMA on Hopper)
    copy(gmem_b_slice, smem_b);
    __syncthreads();
    gemm(tiled_mma, smem_a, smem_b, acc);  // tensor core MMA
    __syncthreads();
}

c++ — thunderkittens
// ThunderKittens: ~100 lines, register-tile-centric
using namespace kittens;
st_bf smem_a; st_bf smem_b;
rt_fl<16, 16> acc[BM/16][BN/16] = {};  // register tile grid
for (int k = 0; k < K; k += BK) {
    load(smem_a, A + k, K);    // HBM → shared
    load(smem_b, B + k * N, N);
    __syncthreads();
    for (int i = 0; i < BM/16; i++)
        for (int j = 0; j < BN/16; j++)
            mma_ABt(acc[i][j], smem_a.get_tile(i), smem_b.get_tile(j), acc[i][j]);
    __syncthreads();
}

Same algorithm. Three levels of abstraction. Triton is shortest but gives you least control over the schedule. CuTe is longest but lets you specify exactly how data flows through the memory hierarchy. ThunderKittens is in between — you think in register tiles, but the library handles tile loads and MMA instruction selection.

The Meta-Trend: Convergence

Despite the explosion, the DSLs are converging on a shared set of abstractions: tiles as first-class objects, pipeline stages for memory latency hiding, layout descriptors for swizzled shared memory, and auto-tuning for block-size selection. The eventual outcome is likely one or two dominant frameworks (Triton + CUTLASS/CuTe) with specialized tools for niche use cases.

Chapter 8: KV Cache Optimization — The Memory Crisis

It is 2:30 PM. The product team wants to launch a feature that requires 128K context: users upload a 200-page document and ask questions about it. Your current serving stack handles 8K context comfortably. At 128K, the KV cache — the stored key and value tensors from previous tokens — consumes more memory than the model weights themselves. Your 70B model at FP16 uses 140 GB for weights. The KV cache at 128K tokens adds another 160 GB. You need a plan.

This chapter covers every major technique for taming the KV cache: architectural changes (MQA, GQA), system-level tricks (PagedAttention), intelligent eviction (SnapKV, H2O, StreamingLLM), and quantization. By the end, you will know exactly which technique to apply for which use case.

The KV Cache Memory Formula

Let us derive the exact memory cost. For a transformer with:

Each layer stores K and V tensors of shape (n_kv, S, d_h). The total KV cache memory:

The factor of 2 is for K and V. Let us plug in numbers:

python
def kv_cache_memory_gb(L, n_kv, d_h, S, bytes_per_elem=2):
    """Compute KV cache memory in GB."""
    total_bytes = 2 * L * n_kv * S * d_h * bytes_per_elem
    return total_bytes / (1024**3)

# Llama-2 70B with MHA (64 KV heads)
print(kv_cache_memory_gb(80, 64, 128, 8192))      # → 20 GB at 8K context
print(kv_cache_memory_gb(80, 64, 128, 32768))     # → 80 GB at 32K context
print(kv_cache_memory_gb(80, 64, 128, 131072))    # → 320 GB at 128K context!

# Llama-2 70B with GQA (8 KV heads) — 8x reduction
print(kv_cache_memory_gb(80, 8, 128, 8192))       # → 2.5 GB
print(kv_cache_memory_gb(80, 8, 128, 131072))     # → 40 GB — manageable!

# With INT4 KV quantization on top of GQA:
print(kv_cache_memory_gb(80, 8, 128, 131072, 0.5)) # → 10 GB — fits easily

Architectural Solutions: MHA → MQA → GQA

The most effective KV cache reduction is architectural — baked into the model during training:

Multi-Head Attention (MHA): standard transformer. Each of the n_h query heads has its own K and V head. KV cache stores n_h × S × d_h per layer. This is what GPT-3, Llama-1, and Llama-2 7B/13B use.

Multi-Query Attention (MQA): all query heads share ONE K head and ONE V head. KV cache drops to 1 × S × d_h per layer — an n_h-fold reduction (typically 32-64x). Introduced by Noam Shazeer in 2019. Used by PaLM, Falcon, StarCoder. The downside: some quality degradation because all heads are forced to attend to the same keys/values.

Grouped-Query Attention (GQA): the compromise. Group every g query heads to share one KV head. With n_h = 64 query heads and g = 8 groups: n_kv = 8 KV heads. An 8x reduction over MHA with minimal quality loss. Used by Llama-2 70B, Llama-3, Mistral, Gemma. GQA is now the de facto standard for large models.

PagedAttention: Virtual Memory for KV Caches

Even with GQA, serving multiple concurrent requests requires careful memory management. The naive approach: pre-allocate a contiguous buffer of max_seq_len × n_kv × d_h for each request. If max_seq_len = 128K but the average request uses 4K tokens, you waste 97% of the allocated memory. With 32 concurrent requests, the waste is hundreds of GB.

PagedAttention (vLLM, 2023) borrows the virtual memory concept from operating systems. Instead of contiguous buffers, KV cache entries are stored in fixed-size pages (blocks) of, say, 16 tokens. Pages are allocated on demand as the sequence grows, and reclaimed when the sequence ends. Pages need not be contiguous in physical GPU memory — a page table maps logical token positions to physical page addresses.

python
# PagedAttention: conceptual implementation
class PagedKVCache:
    def __init__(self, page_size=16, n_layers=80, n_kv_heads=8, d_head=128):
        self.page_size = page_size
        self.n_layers = n_layers
        self.n_kv_heads = n_kv_heads
        self.d_head = d_head

        # Physical page pool: pre-allocate all GPU memory as pages
        self.total_pages = 65536  # e.g., 65K pages
        # Each page: (n_layers, 2, n_kv_heads, page_size, d_head) — K and V
        self.page_pool = torch.zeros(
            self.total_pages, n_layers, 2, n_kv_heads, page_size, d_head,
            dtype=torch.float16, device="cuda"
        )
        self.free_pages = list(range(self.total_pages))  # free list

        # Per-sequence page tables: seq_id → list of page indices
        self.page_tables = {}  # Dict[int, List[int]]

    def allocate_page(self, seq_id):
        """Allocate a new page for a sequence (called when it grows)."""
        if not self.free_pages:
            raise MemoryError("No free KV cache pages")
        page_idx = self.free_pages.pop()
        if seq_id not in self.page_tables:
            self.page_tables[seq_id] = []
        self.page_tables[seq_id].append(page_idx)
        return page_idx

    def free_sequence(self, seq_id):
        """Release all pages for a completed sequence."""
        for page_idx in self.page_tables.pop(seq_id, []):
            self.free_pages.append(page_idx)

    def get_kv(self, seq_id, layer, token_pos):
        """Read K/V for a specific token position."""
        page_num = token_pos // self.page_size
        offset = token_pos % self.page_size
        page_idx = self.page_tables[seq_id][page_num]
        k = self.page_pool[page_idx, layer, 0, :, offset, :]  # (n_kv_heads, d_head)
        v = self.page_pool[page_idx, layer, 1, :, offset, :]
        return k, v

    def utilization(self):
        used = self.total_pages - len(self.free_pages)
        return used / self.total_pages

# Without paging: 32 requests × 128K tokens × ~0.31 GB each = ~10 GB
# With paging: only allocate pages for actual tokens used
# If average length is 4K: 32 × 4K × 0.31/128K × GB ≈ 0.31 GB — 32x saving!

SnapKV: Attention-Guided KV Compression [full paper →]

SnapKV (2024) made a key observation: in long-context attention, most tokens receive negligible attention weight. A small set of "important" tokens dominate the attention pattern, and this set can be identified cheaply by looking at the attention pattern in a few observation layers near the end of the prefix.

The algorithm:

python
def snapkv_compress(K, V, attn_scores, k=1024, window=64):
    """
    SnapKV: attention-guided KV cache compression.
    K, V: (n_heads, seq_len, d_head)
    attn_scores: (n_heads, window, seq_len) — attention from last `window` tokens
    """
    n_heads, seq_len, d_head = K.shape
    prefix_len = seq_len - window

    # Step 1: attention scores from observation window to prefix tokens
    prefix_scores = attn_scores[:, :, :prefix_len]  # (n_heads, window, prefix_len)

    # Step 2: vote — sum attention across observation window positions
    votes = prefix_scores.sum(dim=1)  # (n_heads, prefix_len)

    # Step 3: top-k selection per head
    _, top_indices = votes.topk(k, dim=-1)  # (n_heads, k)
    top_indices, _ = top_indices.sort(dim=-1)  # maintain order

    # Step 4: gather selected K/V + keep observation window
    K_prefix = K.gather(1, top_indices.unsqueeze(-1).expand(-1, -1, d_head))
    V_prefix = V.gather(1, top_indices.unsqueeze(-1).expand(-1, -1, d_head))

    K_window = K[:, prefix_len:, :]  # last `window` tokens — always kept
    V_window = V[:, prefix_len:, :]

    K_compressed = torch.cat([K_prefix, K_window], dim=1)  # (n_heads, k+window, d_head)
    V_compressed = torch.cat([V_prefix, V_window], dim=1)

    return K_compressed, V_compressed  # seq_len reduced from S to k+window

H2O: Heavy Hitter Oracle

H2O (2023) takes a different approach to KV eviction. Instead of a one-shot compression at prompt time, H2O maintains a running eviction policy during generation. It tracks the cumulative attention score for each token across all generation steps. Tokens that consistently receive high attention ("heavy hitters") are kept. Tokens with low cumulative scores are evicted.

python
class H2OCache:
    """Heavy Hitter Oracle: cumulative attention-based eviction."""
    def __init__(self, max_size=2048, window_size=256):
        self.max_size = max_size
        self.window_size = window_size  # recent tokens always kept
        self.cumulative_scores = {}     # token_pos → cumulative attention

    def update_and_evict(self, attn_weights, K, V):
        """
        attn_weights: (n_heads, 1, current_cache_size) — attention from new token
        K, V: current KV cache tensors
        """
        # Update cumulative scores
        scores = attn_weights.mean(dim=0).squeeze()  # average across heads
        for pos in range(len(scores)):
            self.cumulative_scores[pos] = self.cumulative_scores.get(pos, 0) + scores[pos].item()

        # Check if eviction needed
        current_size = K.shape[1]
        if current_size <= self.max_size:
            return K, V

        # Protect recent tokens (sliding window)
        n_evictable = current_size - self.window_size
        evictable_scores = {pos: self.cumulative_scores[pos]
                           for pos in range(n_evictable)}

        # Evict the token with lowest cumulative attention
        victim = min(evictable_scores, key=evictable_scores.get)
        keep_mask = torch.ones(current_size, dtype=torch.bool)
        keep_mask[victim] = False

        return K[:, keep_mask], V[:, keep_mask]

The key insight: in practice, attention patterns follow a power law. A small set of tokens (typically the first few tokens, punctuation, and key content words) receive the lion's share of attention. H2O exploits this by keeping the heavy hitters and evicting the long tail.

StreamingLLM: Attention Sinks + Sliding Window

StreamingLLM (2023) discovered a remarkable phenomenon: the first few tokens in any sequence always receive disproportionately high attention, regardless of their content. These are called attention sinks. The hypothesis is that softmax needs somewhere to dump attention mass when no token is particularly relevant, and position 0 becomes the default "sink."

StreamingLLM's algorithm is dead simple:

StreamingLLM does NOT help with tasks that require accessing information from the evicted middle (e.g., "what was mentioned in paragraph 5 of a 100-paragraph document?"). It is designed for streaming use cases: chat, real-time transcription, monitoring dashboards.

Quantized KV Caches

Orthogonal to eviction strategies, you can quantize the KV cache itself from FP16 to INT8 or INT4. This is different from weight quantization because:

1. KV values change every token. You cannot pre-compute quantization parameters. You need an online quantization scheme that is fast enough to run per-token.

2. The error budget is different. Weight quantization errors are fixed — the same error on every inference. KV quantization errors accumulate: an error in the cached K at position 100 affects every subsequent attention computation forever. This makes KV quantization more sensitive.

3. The distribution is friendlier. KV cache values tend to be more uniformly distributed (post-RoPE, K values are roughly Gaussian) compared to weights (which have those pesky outlier features). INT8 KV quantization with per-head scales typically causes <0.1% perplexity increase. INT4 requires per-token-group scales and causes 0.3-1% increase.

python
def quantize_kv_cache(K, V, bits=8):
    """
    Per-head quantization of KV cache.
    K, V: (n_heads, seq_len, d_head) in FP16
    Returns quantized tensors + per-head scales
    """
    q_max = 2 ** (bits - 1) - 1  # symmetric: -127..127 for INT8

    # Per-head scale (one scale per head)
    K_absmax = K.abs().amax(dim=(1, 2), keepdim=True)  # (n_heads, 1, 1)
    V_absmax = V.abs().amax(dim=(1, 2), keepdim=True)

    K_scale = K_absmax / q_max
    V_scale = V_absmax / q_max

    K_q = (K / K_scale).round().clamp(-q_max, q_max).to(torch.int8)
    V_q = (V / V_scale).round().clamp(-q_max, q_max).to(torch.int8)

    # Memory: INT8 = 1 byte vs FP16 = 2 bytes → 2x compression
    # INT4 (packed) = 0.5 bytes → 4x compression
    return K_q, V_q, K_scale, V_scale

def dequantize_for_attention(K_q, K_scale, Q):
    """Dequantize K on-the-fly during attention computation."""
    K_deq = K_q.float() * K_scale  # dequant before Q @ K^T
    attn = Q @ K_deq.transpose(-1, -2)
    return attn

Combining Techniques: The Full Strategy

In practice, frontier labs combine multiple techniques. Here is the typical stack for long-context serving:

The "Barbarians at the Gate" Perspective [full paper →]

The phrase "Barbarians at the Gate" in the context of KV caches refers to the growing realization that the attention mechanism itself may need to change. All the techniques above are engineering patches on a fundamentally O(N) memory mechanism (linear in sequence length, per layer). State-space models (Mamba, RWKV), linear attention variants, and hybrid architectures are the "barbarians" — they offer O(1) memory per token by replacing the KV cache with a fixed-size recurrent state.

The counterargument: the KV cache is not just overhead. It IS the model's explicit memory. Eviction techniques lose information. Recurrent states compress information lossy. The KV cache is a perfect, lossless memory of every token the model has seen. The frontier lab engineer's job is to make this memory affordable, not to throw it away.

The Engineer's Cheat Sheet

Chapter 9: Training at Scale [dist. training lesson]

You have a 70-billion parameter model. Each parameter is a BF16 number: 2 bytes. The model weights alone occupy 70B × 2 = 140 GB. A single H100 has 80 GB of HBM. The model does not even fit on one GPU — and we have not yet accounted for optimizer states, gradients, or activations.

This is the fundamental problem of large-scale training: the model, its optimizer state, its gradients, and its activations are collectively far too large for any single device. The solution is to split the work across many GPUs. But how you split matters enormously. A naive approach can leave 90% of your GPUs idle, waiting for communication. A well-engineered approach can achieve 40-55% of peak FLOPS utilization — and the difference between 20% and 50% MFU on a 2048-GPU cluster is millions of dollars in wasted compute.

This chapter covers the four pillars of distributed training: data parallelism, fully sharded data parallelism (FSDP), tensor parallelism, and pipeline parallelism. By the end, you will be able to derive the optimal parallelism configuration for any model-cluster pair.

The Memory Budget: What Goes Where

Before choosing a parallelism strategy, you must know what consumes memory. Let us derive the full breakdown for a 70B parameter model trained in BF16 with AdamW optimizer:

Notice the optimizer dominates: 840 GB out of 1,120 GB. This is why ZeRO and FSDP focus primarily on sharding optimizer states.

Strategy 1: Data Parallelism (DDP)

Data Parallelism is the simplest distributed strategy. Every GPU holds a complete copy of the model. Each GPU processes a different mini-batch. After the forward and backward passes, gradients are synchronized across all GPUs via an all-reduce operation, and every GPU applies the same optimizer step.

DDP works beautifully for models that fit on one GPU (up to ~10B parameters in BF16 on an 80GB H100). The communication cost is a single all-reduce of the gradient tensor per step.

All-Reduce: The Core Communication Primitive

The all-reduce operation takes a tensor that exists on every GPU, sums (or averages) them, and distributes the result back to every GPU. The efficient implementation is the ring all-reduce, which works in two phases:

Phase 1: Reduce-Scatter. Each GPU sends 1/N of its gradient tensor to the next GPU in the ring. After N-1 steps, each GPU holds the fully reduced sum for its 1/N shard.

Phase 2: All-Gather. Each GPU sends its reduced shard around the ring. After N-1 steps, every GPU holds the complete reduced result.

The total data transferred per GPU in a ring all-reduce is:

Strategy 2: Fully Sharded Data Parallelism (FSDP / ZeRO)

DDP requires every GPU to hold the full model, gradients, and optimizer states. For a 70B model, that is 1,120 GB per GPU — impossible. FSDP (PyTorch's implementation of DeepSpeed ZeRO) fixes this by sharding everything across GPUs. Each GPU stores only 1/N of the weights, 1/N of the gradients, and 1/N of the optimizer states.

But if each GPU only has 1/N of the weights, how does it compute the forward pass? It all-gathers the full weights just before each layer's forward pass, uses them, then discards them. During the backward pass, it all-gathers again, computes gradients, then reduce-scatters to keep only its 1/N shard of the gradients.

The communication cost is higher than DDP. Let us derive it:

Strategy 3: Tensor Parallelism (Megatron-Style)

FSDP shards weights but requires an all-gather before every layer. Tensor parallelism takes a different approach: it splits each individual layer's computation across GPUs so that no GPU ever needs the full weight matrix.

Consider a linear layer Y = XW, where X is the input [B, d_model] and W is the weight matrix [d_model, d_hidden]. Tensor parallelism splits W column-wise across T GPUs:

In a transformer, Megatron-LM chains column-parallel and row-parallel layers to minimize communication. Each transformer block requires exactly 2 all-reduces in the forward pass (one after the attention block, one after the FFN) and 2 all-reduces in the backward pass.

Strategy 4: Pipeline Parallelism

Pipeline parallelism splits the model vertically by layers. If you have 80 transformer layers and 8 pipeline stages, each GPU holds 10 consecutive layers. GPU 0 runs layers 0-9, GPU 1 runs layers 10-19, and so on.

The problem: naive pipelining has catastrophic bubble overhead. GPU 7 sits idle while GPUs 0-6 process the forward pass through their layers. Then GPU 0 sits idle while GPUs 1-7 finish the forward pass and start the backward pass.

The solution is microbatching. Instead of sending one large batch through the pipeline, split it into M microbatches. As soon as GPU 0 finishes the forward pass for microbatch 1, it starts microbatch 2 while GPU 1 works on microbatch 1. This fills the pipeline.

The Parallelism Decision Tree

Given N GPUs and a model of size M parameters, which strategy do you use? Here is the decision framework that frontier labs actually follow:

Worked Example: LLaMA 3 70B on 2048 H100s

Let us derive the optimal parallelism configuration for training LLaMA 3 70B on a cluster of 2048 NVIDIA H100 GPUs, each with 80 GB HBM, connected by NVLink within 8-GPU nodes and InfiniBand across nodes.

MFU: Model FLOPS Utilization

The ultimate metric for training efficiency is MFU — what fraction of the GPU's theoretical peak FLOPS you are actually using for useful computation (excluding communication, bubbles, recomputation, and overhead).

Gradient Accumulation and Its Interaction with Parallelism

Sometimes the global batch size you need is larger than what fits in GPU memory across all data-parallel replicas. Gradient accumulation solves this: run K forward-backward passes, accumulating gradients, then do one optimizer step.

Gradient accumulation is "free" in terms of communication — you only all-reduce once per K steps. But it increases latency per optimizer step by K×.

Mixed Precision Training: The Loss Scaling Dance

Training in BF16 saves memory and doubles throughput on tensor cores. But gradients can underflow to zero in BF16 (the smallest representable normal number is ~1.2×10^-38 for the exponent, but mantissa precision is only ~0.01). Small gradients from early layers can vanish.

Loss scaling solves this: multiply the loss by a large constant S before backpropagation. This scales all gradients by S, pushing them away from the underflow zone. After the backward pass, divide gradients by S before the optimizer step.

Communication Cost Summary

Interactive: Parallelism Strategy Visualizer

Activation Checkpointing: Trading Compute for Memory

During the forward pass, every layer's output is stored in memory so it can be used during the backward pass. For a 70B model with 80 layers, batch size 1, and sequence length 8192, the activation memory is enormous:

Activation checkpointing (also called gradient checkpointing or rematerialization) discards intermediate activations during the forward pass and recomputes them during the backward pass. The trade-off: 33% more compute (one extra forward pass per layer), but activation memory drops from O(L) to O(sqrt(L)) or O(1) depending on the strategy.

Sequence Parallelism and Context Parallelism

When training with very long sequences (32K, 128K, or even 1M tokens), the attention memory grows quadratically with sequence length. Even with FlashAttention, the KV cache for one layer is B × 2 × H_kv × S × d_head × 2 bytes. At S=128K with GQA (8 KV heads, d_head=128):

Context parallelism (also called sequence parallelism for the attention component) splits the sequence across GPUs. Each GPU processes a shard of the sequence and communicates partial attention results via all-to-all or ring exchanges. This allows training on sequences far longer than what fits on a single GPU.

Ring attention is the most elegant implementation: GPUs are arranged in a ring. Each GPU holds a chunk of Q and a chunk of K,V. The K,V chunks are passed around the ring, and each GPU accumulates the attention for its Q chunk against all K,V chunks as they arrive. After one full ring rotation, every GPU has its complete attention output.

Failure Recovery at Scale

When training on 2048+ GPUs for weeks, hardware failures are guaranteed. A single GPU failure can halt a training run that costs $50K per hour of idle time. Frontier labs engineer for this:

python
# Compute parallelism configuration for any model-cluster pair
def compute_parallel_config(
    num_params_B: float,   # billions of parameters
    num_gpus: int,         # total GPU count
    hbm_per_gpu_GB: float = 80,
    num_heads: int = 64,
    num_layers: int = 80,
):
    P = num_params_B * 1e9

    # Memory breakdown (BF16 training with AdamW)
    weights_GB = P * 2 / 1e9          # BF16 weights
    grads_GB = P * 2 / 1e9            # BF16 gradients
    opt_GB = P * 12 / 1e9             # FP32 master + m + v
    total_static_GB = weights_GB + grads_GB + opt_GB

    # Strategy 1: Can we fit on 1 GPU? → DDP
    if total_static_GB * 1.3 < hbm_per_gpu_GB:
        return {"strategy": "DDP", "DP": num_gpus,
                "mem_per_gpu": total_static_GB}

    # Strategy 2: Choose TP (must divide heads, max 8 for NVLink)
    tp = 1
    for t in [8, 4, 2]:
        if num_heads % t == 0 and weights_GB / t < hbm_per_gpu_GB * 0.5:
            tp = t; break

    # Strategy 3: Choose PP (enough stages to fit layers)
    pp = 1
    while total_static_GB / (tp * pp) > hbm_per_gpu_GB * 0.6:
        pp *= 2
        if pp > 16: break

    # Remaining GPUs are data parallel
    dp = num_gpus // (tp * pp)

    # Per-GPU memory with FSDP across DP dimension
    layers_per_stage = num_layers // pp
    frac = layers_per_stage / num_layers
    mem_weights = weights_GB * frac / tp          # TP shards weights
    mem_grads = grads_GB * frac / (dp * tp)       # FSDP shards grads
    mem_opt = opt_GB * frac / (dp * tp)           # FSDP shards opt
    mem_per_gpu = mem_weights + mem_grads + mem_opt

    return {
        "strategy": f"TP={tp} x PP={pp} x DP={dp}",
        "tp": tp, "pp": pp, "dp": dp,
        "mem_per_gpu_GB": round(mem_per_gpu, 2),
        "total_static_GB": round(total_static_GB, 1),
        "bubble_frac": (pp - 1) / (pp - 1 + 4 * pp) if pp > 1 else 0,
    }

# Example: LLaMA 3 70B on 2048 H100s
config = compute_parallel_config(70, 2048)
print(config)
# {'strategy': 'TP=8 x PP=4 x DP=64', 'mem_per_gpu_GB': 5.14, ...}

Real-World Training Configurations (2024-2025)

Let us examine the actual parallelism configurations used by frontier labs. These numbers come from published papers and technical reports:

DeepSpeed vs Megatron-LM vs PyTorch FSDP

Three major frameworks support distributed training. Knowing their trade-offs matters for system design interviews:

In practice, Meta uses Megatron-LM + FSDP for LLaMA training. Google uses its own internal framework (Pathways) on TPUs. DeepSpeed is popular for research labs and smaller companies. The choice often depends on what your cluster runs (NVIDIA GPUs → Megatron-LM, TPUs → JAX/Pax, budget-constrained → DeepSpeed).

The Learning Rate Schedule: An Overlooked Performance Factor

The learning rate schedule interacts with parallelism in non-obvious ways. When you change the global batch size (by changing DP or gradient accumulation), the optimal learning rate changes too.

Chapter 10: Serving Infrastructure [inference lesson] [handbook]

Training a model costs millions of dollars and takes weeks. Serving that model costs millions of dollars per month and runs indefinitely. At frontier labs, the serving infrastructure team often has a larger headcount than the training team. And the engineering challenges are fundamentally different: training optimizes for throughput (tokens per second across the cluster), serving optimizes for latency under load (time to first token for the user who just typed a query while 10,000 other users are also waiting).

This chapter covers the core systems that make LLM serving efficient: the two-phase serving model, continuous batching, speculative decoding, disaggregated serving, and KV cache management. By the end, you will be able to derive the cost of serving a 70B model and identify exactly where the bottlenecks are.

Prefill vs Decode: The Two-Phase Model

When a user sends a prompt to an LLM, the server does two fundamentally different things:

Phase 1: Prefill. The entire prompt (say, 500 tokens) is processed in a single forward pass. Every token can attend to all previous tokens in parallel. This is a large matrix multiply — compute-bound. The GPU tensor cores are busy. Memory bandwidth is not the bottleneck.

Phase 2: Decode. Output tokens are generated one at a time. Each forward pass processes exactly 1 new token (attending to all previous tokens via the KV cache). This is a tiny matrix multiply — memory-bandwidth-bound. The tensor cores are mostly idle, waiting for weights to be loaded from HBM.

Batching Theory: Deriving the Throughput-Latency Curve

Batching is the single most important optimization in LLM serving. Let us derive exactly why, starting from first principles.

Consider a 70B model on 4x H100 with TP=4. The model weights occupy 140/4 = 35 GB per GPU. During decode, each token requires loading all weights from HBM. The key insight: loading weights once for B tokens costs nearly the same as loading them for 1 token (the extra KV cache and activation memory is tiny in comparison).

Continuous Batching: Why Static Batching Wastes GPUs

In static batching, you collect B requests, process them together, and wait for ALL B requests to finish generating before accepting new ones. If request 1 generates 20 tokens and request 2 generates 500 tokens, request 1's GPU slot sits idle for 480 decode steps. For a batch of diverse-length requests, average GPU utilization can be as low as 30-40%.

Continuous batching (also called "in-flight batching" or "iteration-level batching") fixes this: as soon as any request finishes, its slot is immediately filled with a new request from the queue. The batch composition changes at every decode step.

Let us quantify the waste. Consider 4 requests with output lengths [30, 80, 150, 500]:

The implementation requires a scheduler that tracks the state of every request (prefilling, decoding, waiting, done) and makes admission decisions at every iteration:

python
class ContinuousBatchScheduler:
    def __init__(self, max_batch: int, max_tokens: int):
        self.max_batch = max_batch     # max concurrent requests
        self.max_tokens = max_tokens   # max total tokens in KV cache
        self.running = []              # requests currently generating
        self.waiting = []              # queue of pending requests

    def schedule_step(self):
        # Step 1: Remove finished requests (hit EOS or max_len)
        finished = [r for r in self.running if r.is_done()]
        self.running = [r for r in self.running if not r.is_done()]

        # Step 2: Free KV cache memory from finished requests
        for r in finished:
            self.kv_cache.free(r.request_id)

        # Step 3: Admit new requests from the queue
        while (self.waiting
               and len(self.running) < self.max_batch
               and self.kv_cache.available() > self.waiting[0].est_tokens):
            req = self.waiting.pop(0)
            req.state = "prefill"
            self.running.append(req)

        # Step 4: Separate prefill and decode requests
        prefill_reqs = [r for r in self.running if r.state == "prefill"]
        decode_reqs = [r for r in self.running if r.state == "decode"]

        # Step 5: Run prefill for new requests (compute-bound)
        if prefill_reqs:
            self.run_prefill(prefill_reqs)
            for r in prefill_reqs:
                r.state = "decode"

        # Step 6: Run one decode step for all decoding requests
        if decode_reqs:
            self.run_decode(decode_reqs)

        return finished

Chunked Prefill: Reducing Head-of-Line Blocking

A subtle problem with continuous batching: when a new request with a long prompt (say, 10,000 tokens) enters the batch, its prefill takes hundreds of milliseconds. During this time, the decode requests in the same batch are blocked — they cannot generate their next token until the prefill completes. This is head-of-line blocking.

Chunked prefill solves this by breaking long prefills into smaller chunks (e.g., 512 tokens at a time). Between chunks, decode requests get a chance to run. The result: prefill spreads over multiple iterations, and decode latency stays bounded.

Guided Decoding and Structured Output

Many applications need the LLM to output valid JSON, SQL, or code that follows a grammar. Guided decoding constrains the token-by-token generation to follow a formal grammar or regex pattern by masking out invalid tokens at each step.

python
# Guided JSON decoding: mask invalid tokens at each step
class JSONGuide:
    def __init__(self, schema: dict):
        self.schema = schema
        self.state = "start"  # FSM state: start, key, value, etc.
        self.stack = []        # nesting stack: [object, array, ...]

    def get_valid_tokens(self, generated_so_far: str) -> set:
        """Return the set of token IDs that are valid at this position."""
        if self.state == "start":
            return tokens_starting_with("{")
        elif self.state == "key":
            # Only allow keys from schema
            valid_keys = self.schema["properties"].keys()
            return tokens_for_any_of([f'"{k}"' for k in valid_keys])
        elif self.state == "value":
            expected_type = self.current_field_type()
            if expected_type == "integer":
                return digit_tokens | tokens_for("-")
            elif expected_type == "string":
                return tokens_starting_with('"')
            # ... more types
        return all_tokens  # fallback: no constraint

    def mask_logits(self, logits, generated_so_far):
        """Zero out logits for invalid tokens."""
        valid = self.get_valid_tokens(generated_so_far)
        mask = torch.full_like(logits, -1e9)
        mask[list(valid)] = 0
        return logits + mask

Speculative Decoding: Small Model Drafts, Large Model Verifies

Autoregressive decoding is slow because you generate one token at a time. Speculative decoding accelerates this with a clever insight: use a small, fast draft model to generate K candidate tokens, then have the large target model verify all K tokens in a single forward pass (which is as fast as verifying 1 token, since it is compute-bound like prefill).

Disaggregated Serving: Split Prefill and Decode

Prefill and decode have opposite hardware requirements. Prefill is compute-bound — it benefits from high FLOPS (many tensor cores). Decode is memory-bandwidth-bound — it benefits from high HBM bandwidth and large batch sizes. Running both on the same GPU creates contention: a long prefill blocks decode for all batched requests, causing latency spikes.

Disaggregated serving runs prefill and decode on separate GPU pools:

The critical challenge is transferring the KV cache from prefill to decode. For a 70B model with 80 layers, 8 KV heads, d_head=128, and a 4096-token prompt in BF16:

KV Cache Management in Production

The KV cache is the largest per-request memory consumer. For a serving system handling hundreds of concurrent requests, KV cache management becomes a critical systems problem.

PagedAttention (from vLLM) is the breakthrough that made efficient KV cache management possible. Instead of allocating contiguous memory for each request's KV cache (which leads to fragmentation when requests have variable lengths), PagedAttention divides KV cache memory into fixed-size pages (typically 16 tokens per page):

python
# PagedAttention KV Cache Manager (simplified)
class PagedKVCacheManager:
    def __init__(self, num_pages: int, page_size: int = 16,
                 num_layers: int = 80, num_heads: int = 8,
                 head_dim: int = 128):
        self.page_size = page_size
        # Physical pages: [num_pages, 2, num_layers, num_heads, page_size, head_dim]
        # 2 = key and value
        self.physical_pages = torch.zeros(
            num_pages, 2, num_layers, num_heads, page_size, head_dim,
            dtype=torch.bfloat16, device="cuda"
        )
        self.free_pages = list(range(num_pages))
        self.page_tables = {}  # request_id → [page_indices]

    def allocate(self, request_id: str, num_tokens: int):
        num_pages_needed = (num_tokens + self.page_size - 1) // self.page_size
        if len(self.free_pages) < num_pages_needed:
            return False  # out of memory — preempt a request
        pages = [self.free_pages.pop() for _ in range(num_pages_needed)]
        self.page_tables[request_id] = pages
        return True

    def append_token(self, request_id: str):
        # Check if current last page has space
        pages = self.page_tables[request_id]
        current_len = self.get_seq_len(request_id)
        if current_len % self.page_size == 0:
            # Need a new page
            if not self.free_pages:
                return False
            pages.append(self.free_pages.pop())
        return True

    def free(self, request_id: str):
        pages = self.page_tables.pop(request_id, [])
        self.free_pages.extend(pages)

Latency Metrics: What You Measure, What You Optimize

The Cost Model: $/1M Tokens

Cloud providers and API companies price LLM inference in dollars per million tokens. Let us derive this cost from first principles for serving LLaMA 3 70B:

Interactive: Continuous Batching Visualizer

Worked Example: Serving LLaMA 3 70B on 8x TPU v5e

Let us derive the full serving profile for LLaMA 3 70B on a TPU v5e pod slice of 8 chips, using tensor parallelism.

Request Preemption and Priority Scheduling

When KV cache memory runs out, the scheduler must decide: reject new requests (increasing queue wait time) or preempt an existing request (evict its KV cache, restart it later). Sophisticated schedulers use priority-based preemption:

python
# Priority preemption: evict lowest-priority request when memory is full
def preempt_if_needed(self, new_request):
    """Evict a running request to make room for a higher-priority one."""
    if self.kv_cache.available() >= new_request.est_tokens:
        return True  # enough memory, no preemption needed

    # Find lowest-priority running request
    candidates = sorted(
        self.running,
        key=lambda r: (r.priority, -r.tokens_generated)
    )
    # Evict requests with LOWER priority than the new one
    for victim in candidates:
        if victim.priority >= new_request.priority:
            break  # don't preempt equal or higher priority

        # Swap out: save KV cache to CPU memory (or discard + recompute later)
        self.kv_cache.swap_out(victim.request_id)
        self.running.remove(victim)
        victim.state = "preempted"
        victim.preempt_count += 1
        self.waiting.insert(0, victim)  # re-queue at front

        if self.kv_cache.available() >= new_request.est_tokens:
            return True

    return False  # cannot make room even after preemption

Quantized KV Cache: FP8 and INT4

The KV cache is a major memory consumer during serving. Quantizing it from BF16 to FP8 or even INT4 can double or quadruple the number of concurrent requests:

Multi-Turn Conversation Optimization

In a chatbot, users send multiple messages in a conversation. Each new message requires attending to the full conversation history. Without optimization, every user message triggers a complete prefill of the entire conversation — wasting compute on tokens already processed.

The solution is KV cache persistence: keep the KV cache from previous turns alive between requests. When the user sends a new message, only the new tokens need prefill. This requires:

Chapter 11: LLM-Guided Search

Every system we have studied so far uses LLMs as standalone generators: prompt in, text out. But the most exciting frontier in 2025 is using LLMs as components inside larger optimization loops. The LLM does not solve the problem — it proposes candidate solutions that an automated evaluator scores, and the best candidates are fed back to the LLM to produce better proposals. This is LLM-guided search.

This pattern has already produced genuine scientific discoveries: new algorithms for matrix multiplication, state-of-the-art solutions to open combinatorics problems, and improvements to real-world infrastructure systems. It is the closest thing we have to "AI doing research."

The General Pattern

Every LLM-guided search system follows the same four-step loop:

FunSearch: LLM as Mutation Operator [lesson] [applied] [comp.prog paper]

Google DeepMind's FunSearch (2023) was the first system to demonstrate that LLMs can make genuine mathematical discoveries. The name stands for "searching in the space of functions" — the LLM generates Python functions, and an evaluator scores them.

FunSearch tackled the cap set problem — a notorious open problem in combinatorics. A cap set is a subset of points in F₃ⁿ (the n-dimensional space over the field with 3 elements) such that no three points are collinear. The question: how large can a cap set be?

Mathematicians had found cap sets of size 512 in F₃⁸ through clever constructions. FunSearch found cap sets of size 512 as well — matching the best known — but using a completely different construction method that humans had not considered. For larger dimensions, it exceeded the previous best.

Worked Example: The Cap Set Problem

To understand what FunSearch actually discovered, we need to understand the problem it solved. A cap set in F₃ⁿ is a subset S of {0, 1, 2}ⁿ such that no three elements a, b, c in S satisfy a + b + c = 0 (mod 3). In other words, no three points form an arithmetic progression.

FunSearch does not search directly for cap sets. Instead, it evolves a construction function — a Python function that takes the dimension n and returns a set of points. The evaluator checks (a) that the returned set is actually a cap set (no collinear triples), and (b) how large it is. Larger = better score.

FunSearch's Island Model

FunSearch does not use a single population of solutions. It maintains multiple islands — independent populations that evolve separately. Periodically, the best solution from one island migrates to another. This prevents premature convergence (all candidates becoming too similar).

FunSearch Applied: Online Bin Packing

Beyond pure mathematics, FunSearch discovered new heuristics for online bin packing — a classic NP-hard optimization problem with real-world applications (packing containers, scheduling jobs, memory allocation).

The problem: items of various sizes arrive one at a time. You must place each item into a bin (capacity 1.0) immediately, without seeing future items. The goal: minimize the total number of bins used.

The best known online heuristic was First Fit Decreasing. FunSearch discovered a new priority function that outperformed all previously known heuristics on standard benchmarks. The discovered function was non-obvious — a piecewise function with specific breakpoints that no human had designed.

python
# FunSearch-discovered bin packing heuristic (simplified)
# The LLM evolved this priority function over 10,000+ generations
def priority(item_size: float, bin_remaining: float) -> float:
    """Score for placing item into a bin. Higher = preferred."""
    waste = bin_remaining - item_size
    if waste < 0:
        return -1e9  # does not fit
    if waste < 0.01:
        return 1e6   # nearly perfect fit — always prefer
    if waste < 0.1:
        return 500 - waste * 1000
    if waste < 0.33:
        return 100 - waste * 200
    # Key insight: prefer bins that leave specific "useful" gaps
    # (gaps that commonly-sized future items could fill)
    if 0.49 < waste < 0.51:
        return 80   # half-empty bins are useful
    if 0.32 < waste < 0.35:
        return 60   # one-third gaps are useful
    return 10 - waste * 5

AlphaEvolve: FunSearch's Successor [full paper →]

Google DeepMind's AlphaEvolve (2025) extends FunSearch in three critical ways:

AlphaEvolve's most famous result: discovering a new algorithm for 4×4 complex matrix multiplication that uses fewer scalar multiplications than Strassen's algorithm. Specifically, it found a decomposition that uses 48 multiplications for 4×4 matrices, improving on the 49 previously known. This is a genuine mathematical discovery — the first improvement in this specific problem in decades.

AlphaEvolve: A Deeper Look at the Architecture

AlphaEvolve's architecture is worth studying in detail because it represents the state of the art in LLM-guided search as of 2025. Let us trace the lifecycle of a single evolution step:

Beyond FunSearch: Other LLM-as-Optimizer Systems

FunSearch and AlphaEvolve are the most famous, but the LLM-guided search paradigm has spawned many systems:

The common thread: an LLM generates candidates in a domain where automated evaluation is possible, and selection pressure drives improvement over hundreds or thousands of generations.

The AutoResearch Pattern

Andrej Karpathy's concept of Autoresearch (2024-2025) extends LLM-guided search to the full research workflow. Instead of searching for a single function, the system:

Minimal FunSearch-Style Loop

Here is a complete, runnable implementation of the core FunSearch loop for evolving a sorting key function:

python
import random
import subprocess
import json
from openai import OpenAI

client = OpenAI()

# ── Problem: evolve a priority function for bin packing ──
SPEC = """You are evolving a Python function called `priority(item_size, bin_remaining)`
that returns a float score. Higher score = preferred bin for the item.
The function is used to solve online bin packing: items arrive one at a time,
each must be placed in a bin (capacity 1.0) immediately.
Goal: minimize total bins used across the test suite."""

EVAL_HARNESS = """
import random
random.seed(42)

# Generate test instances
def make_instance(n=100):
    return [random.uniform(0.1, 0.9) for _ in range(n)]

def evaluate(priority_fn, instances):
    total_bins = 0
    for items in instances:
        bins = []  # list of remaining capacities
        for item in items:
            scores = [(priority_fn(item, b), i) for i, b in enumerate(bins) if b >= item]
            if scores:
                _, best_i = max(scores)
                bins[best_i] -= item
            else:
                bins.append(1.0 - item)
        total_bins += len(bins)
    return -total_bins  # negative because lower bins = better, but we maximize score

instances = [make_instance(200) for _ in range(20)]
{candidate_code}
score = evaluate(priority, instances)
print(json.dumps({"score": score}))
"""

class Island:
    def __init__(self, capacity=50):
        self.programs = []  # (score, code) sorted by score descending
        self.capacity = capacity

    def insert(self, score, code):
        self.programs.append((score, code))
        self.programs.sort(key=lambda x: x[0], reverse=True)
        if len(self.programs) > self.capacity:
            self.programs.pop()

    def sample_best(self, k=2):
        # Tournament selection: pick from top 30%
        top = self.programs[:max(1, len(self.programs) // 3)]
        return random.sample(top, min(k, len(top)))

def evaluate_candidate(code: str, timeout: int = 30) -> float:
    """Run candidate in sandbox, return score."""
    full_code = EVAL_HARNESS.replace("{candidate_code}", code)
    try:
        result = subprocess.run(
            ["python3", "-c", full_code],
            capture_output=True, text=True, timeout=timeout
        )
        return json.loads(result.stdout)["score"]
    except:
        return -1e9  # crashed or timed out

def generate_candidate(parents: list) -> str:
    """Ask LLM to evolve a new candidate from parent programs."""
    parent_text = "\n\n".join(
        f"# Score: {s}\n{c}" for s, c in parents
    )
    response = client.chat.completions.create(
        model="gpt-4o",
        temperature=1.0,
        messages=[{
            "role": "user",
            "content": f"""{SPEC}

Here are the best solutions found so far:
{parent_text}

Write an improved `priority` function. Be creative — try a
different mathematical form. Output ONLY the Python function."""
        }]
    )
    return response.choices[0].message.content

# ── Main evolution loop ──
NUM_ISLANDS = 5
islands = [Island() for _ in range(NUM_ISLANDS)]

# Seed with a simple baseline
seed = """def priority(item_size, bin_remaining):
    if bin_remaining < item_size:
        return -1e9
    return -(bin_remaining - item_size)  # first-fit-decreasing"""
seed_score = evaluate_candidate(seed)
for island in islands:
    island.insert(seed_score, seed)

best_ever = seed_score
for gen in range(500):
    # Pick a random island
    island = random.choice(islands)
    parents = island.sample_best(2)

    # Generate and evaluate
    candidate = generate_candidate(parents)
    score = evaluate_candidate(candidate)
    island.insert(score, candidate)

    if score > best_ever:
        best_ever = score
        print(f"Gen {gen}: NEW BEST = {score}")

    # Migrate every 50 generations
    if gen % 50 == 0 and gen > 0:
        src, dst = random.sample(islands, 2)
        if src.programs:
            dst.insert(*src.programs[0])

When LLM Search Fails: The Antipatterns

Not every problem benefits from LLM-guided search. Understanding the failure modes is as important as understanding the successes:

The Prompt Engineering of LLM Search

The quality of the LLM's proposals depends critically on how you prompt it. FunSearch and AlphaEvolve both discovered that specific prompt structures dramatically improve the mutation quality:

Building Your Own LLM Search System

If you want to apply LLM-guided search to your own problem, here is the engineering checklist:

Interactive: Evolutionary Loop Visualizer

Chapter 12: The Scaling Book — TPU/GPU Rooflines [exercises →] [full book →]

You have written the training loop. You have configured the parallelism. The job is running. But the MFU reads 22%. Nearly 80% of the cluster's theoretical capability is wasted. Where is the performance going?

To answer this question, you need to understand how the hardware actually works — not at the CUDA-kernel level (Chapter 2 covered that), but at the architectural level: how data flows through the chip, where bottlenecks form, and how to predict performance before writing a single line of code. This is the roofline model, and it is the single most useful tool in a performance engineer's toolkit.

GPU Architecture: The H100

An NVIDIA H100 SXM is organized as a hierarchy of compute and memory:

TPU Architecture: The v5e

Google's TPU v5e is a fundamentally different design philosophy. Where GPUs evolved from graphics (many small cores, complex scheduling), TPUs were built from scratch for matrix multiplication.

Systolic Arrays: How TPU MXUs Work

A systolic array is a grid of processing elements (PEs) where data flows rhythmically through the grid, like a heartbeat (hence "systolic"). Each PE performs one multiply-accumulate per cycle and passes the result to its neighbor.

For a 128×128 systolic array computing C = A × B:

The Roofline Model

The roofline model predicts the maximum achievable performance of any operation based on a single metric: arithmetic intensity (operations per byte of memory traffic). The model has two regimes:

Let us compute the arithmetic intensity for the key operations in a transformer:

TPU vs GPU: Where Each Wins

NVLink/NVSwitch vs ICI: Interconnect Battle

For distributed training, the interconnect between devices is often the bottleneck, not the compute. The two architectures take fundamentally different approaches:

Worked Example: Same Matmul on H100 vs TPU v5e

Let us predict the performance of a [4096, 4096] × [4096, 4096] BF16 matmul on both chips, then verify against roofline predictions.

The Profiling Workflow: JAX + XLA

On TPUs, you do not write kernels. Instead, JAX traces your Python code, XLA compiles it to optimized HLO (High-Level Operations), and the TPU runtime schedules the execution. Profiling tells you where time is spent in this compiled execution.

python
import jax
import jax.numpy as jnp
from jax import profiler

# ── Step 1: Define the computation ──
def transformer_block(x, w_q, w_k, w_v, w_o, w_ff1, w_ff2, w_ln):
    # Layer norm (memory-bound)
    x_norm = jax.nn.standardize(x, axis=-1)

    # Attention projections (compute-bound for large batch)
    q = x_norm @ w_q   # [B, S, D] @ [D, D] → [B, S, D]
    k = x_norm @ w_k
    v = x_norm @ w_v

    # Attention scores (compute-bound for long sequences)
    d_k = q.shape[-1]
    scores = (q @ k.transpose(0, 1, 3, 2)) / jnp.sqrt(d_k)
    attn = jax.nn.softmax(scores, axis=-1)

    # Attention output
    out = (attn @ v) @ w_o
    x = x + out  # residual (memory-bound)

    # FFN (compute-bound)
    x_norm = jax.nn.standardize(x, axis=-1)
    h = jax.nn.gelu(x_norm @ w_ff1)  # [B,S,D] @ [D,4D] → [B,S,4D]
    x = x + h @ w_ff2                 # [B,S,4D] @ [4D,D] → [B,S,D]
    return x

# ── Step 2: JIT compile ──
transformer_jit = jax.jit(transformer_block)

# ── Step 3: Profile ──
B, S, D = 8, 2048, 4096
key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (B, S, D), dtype=jnp.bfloat16)
# ... initialize weight matrices ...

# Warm up (trigger compilation)
_ = transformer_jit(x, w_q, w_k, w_v, w_o, w_ff1, w_ff2, w_ln)

# Profile 10 iterations
with profiler.trace("/tmp/jax-trace"):
    for _ in range(10):
        x = transformer_jit(x, w_q, w_k, w_v, w_o, w_ff1, w_ff2, w_ln)
        x.block_until_ready()  # force synchronization

# View in TensorBoard: tensorboard --logdir=/tmp/jax-trace
# Look for:
#   - HLO operation breakdown (which ops take the most time)
#   - Memory usage timeline (peak vs sustained)
#   - ICI communication time (for multi-chip runs)
#   - MXU utilization % (target: >80%)

XLA Compilation: How JAX Lowers to Hardware

When you write `y = x @ w` in JAX, the path to execution is:

Sharded Matrices: How Matmul Distributes

When training on multiple TPU chips (or GPUs with tensor parallelism), the large weight matrices are sharded across devices. The matmul C = A × B becomes a distributed operation.

Interactive: TPU vs GPU Architecture Comparison

GPU Profiling with Nsight

On GPUs, the profiling workflow uses NVIDIA's Nsight Systems (system-level timeline) and Nsight Compute (kernel-level analysis):

bash
# Nsight Systems: capture a full training step timeline
nsys profile --trace=cuda,nvtx --output=train_step \
    python train.py --steps 10

# What to look for in the timeline:
# 1. GPU idle gaps → CPU-bound or waiting for data
# 2. Small kernels with gaps → launch overhead (use CUDA Graphs)
# 3. Long NCCL kernels → communication bottleneck
# 4. Memory copy (H2D/D2H) overlapping with compute → good!
# 5. Memory copy NOT overlapping → pipeline stall

# Nsight Compute: deep-dive on one kernel
ncu --set full --target-processes all \
    python -c "import torch; a=torch.randn(4096,4096,device='cuda',dtype=torch.bfloat16); b=a@a"

# Key metrics from ncu:
# - SM utilization (% of SMs active)
# - Tensor Core utilization (% of TC active)
# - Memory throughput (% of peak BW achieved)
# - Achieved occupancy (warps per SM)
# - Arithmetic intensity (actual, not theoretical)

Memory-Bound Optimization: The Fusion Imperative

Our roofline analysis showed that most transformer operations are memory-bound. Let us derive exactly how much time is wasted on non-compute operations for a single transformer block on an H100:

Practical Profiling: What to Do When MFU Is Low

Here is the systematic debugging checklist that production teams at frontier labs follow when MFU is below expectations:

Putting It All Together: From Code to FLOPS

Here is the mental model you should carry. Every computation you write passes through this stack:

At every layer, performance can be lost. The profiling tools let you pinpoint where. The roofline model tells you the theoretical maximum. The gap between theoretical and achieved is your optimization opportunity.

Chapter 13: Mixture-of-Experts Architectures

You have just been handed an uncomfortable reality. The dense transformer you have been training — every parameter active on every token — is hitting diminishing returns. You doubled the parameter count from 7B to 14B, doubled the compute budget, and gained... 0.3 points on your eval suite. The Chinchilla-optimal frontier says you need 4x the tokens to justify 4x the parameters, but your data pipeline cannot produce tokens that fast. There has to be a way to get more parameters without a proportional compute increase.

There is. It is called Mixture of Experts (MoE), and it is the architectural pattern behind every frontier lab's largest models: GPT-4, Gemini, DeepSeek-V3, Mixtral, and Grok. The idea is deceptively simple: instead of one big feedforward network (FFN) that processes every token, you have many smaller FFNs (the "experts"), and a lightweight router decides which experts handle which token. Most of the parameters are dormant for any given token. You pay the storage cost of a huge model but only the compute cost of a small one.

Where MoE Lives in the Transformer

A standard transformer block has two main components: a multi-head attention layer and a feedforward network (FFN). MoE replaces the FFN. The attention layer stays dense — every head sees every token. Only the FFN becomes sparse.

Why only the FFN? Two reasons. First, the FFN typically accounts for 2/3 of the parameters in a transformer block (it expands the hidden dimension by 4x, then contracts back). Replacing it with experts gives you the biggest parameter multiplier for the smallest architectural change. Second, attention is inherently a cross-token operation — all tokens need to interact. The FFN operates token-independently, so routing tokens to different experts does not break any dependency.

The Router: Gated Softmax Selection

The router is the brain of MoE. It takes a token's hidden state x and produces a routing decision. Here is the math, step by step.

Think of it as a committee vote where only the top-k committee members get to speak, and their influence is proportional to how confident the router is in their relevance.

Worked Example: Routing a Single Token

Let's trace through a concrete example. We have 8 experts, top-2 routing, and d_model = 4 (tiny for illustration).

Load Balancing: The Expert Collapse Problem

Here is the failure mode that kills naive MoE training. Without intervention, the router learns to send almost all tokens to 1-2 "favorite" experts. Those experts get better because they see more data. Because they are better, the router sends even more tokens to them. The other experts never get trained. Within a few thousand steps, you have a "Mixture of 2 Experts plus 6 dead weights." This is called expert collapse.

The standard fix is an auxiliary load-balancing loss that penalizes uneven expert usage. The loss has two terms:

Why multiply f_i by p_i? Because f_i depends on the argmax (not differentiable), but p_i is the smooth softmax probability (differentiable). The product lets gradients flow through p_i to encourage the router to spread probability mass across experts. The total loss becomes L_total = L_LM + L_aux.

Deriving the FLOPs Savings

Let's be precise about why MoE is efficient. A dense transformer FFN has two weight matrices: W₁ ∈ ℝ^{d × 4d} and W₂ ∈ ℝ^{4d × d}. The FLOPs for one token through this FFN are approximately 2 × d × 4d + 2 × 4d × d = 16d² (counting multiply-accumulate as 2 FLOPs).

Expert Parallelism: All-to-All Communication

MoE introduces a unique parallelism challenge. In a dense model, every GPU processes the same layers — you can split by data parallelism (different batches) or tensor parallelism (different parts of the same matrix). In MoE, different GPUs hold different experts, and tokens must physically move to the GPU hosting their assigned expert.

This is called expert parallelism, and the communication pattern is all-to-all: every GPU may need to send tokens to every other GPU, and every GPU may receive tokens from every other GPU. It is the MoE equivalent of traffic at an intersection where every car wants to go in a different direction.

The all-to-all operation has two phases. In the dispatch phase, each GPU sends its tokens to the GPUs hosting the selected experts. In the combine phase, each GPU sends the expert outputs back to the originating GPUs. Both phases are all-to-all collectives, and they are the dominant communication overhead in MoE training.

Worked Example: All-to-All Communication Cost

Let's calculate the actual communication overhead for a realistic MoE setup, because this is the number that determines whether MoE is practical at your cluster scale.

The Capacity Factor and Token Dropping

Load balancing is imperfect. Even with the auxiliary loss, some experts receive more tokens than others. If expert 3 is assigned 2x the average number of tokens, processing it takes 2x as long, and all other GPUs sit idle waiting — destroying parallelism.

The capacity factor C controls the maximum number of tokens each expert can process. If an expert is assigned more than C × (total_tokens / num_experts) tokens, the excess tokens are dropped — they pass through the layer with only the residual connection (no FFN computation). Typical values: C = 1.0 (strict) to C = 1.5 (generous).

DeepSeek-V3's MoE Design: Auxiliary-Free Balancing

DeepSeek-V3 (2024) introduced several MoE innovations that frontier labs are rapidly adopting. The most significant: no auxiliary balancing loss. Instead, they use a simple bias term.

The key insight behind DeepSeek's auxiliary-free approach: instead of adding a loss term that fights with the main training objective, they add a learnable bias b_i to each expert's routing logit. When an expert is underutilized, its bias increases slightly; when overutilized, it decreases. The bias is updated by a simple heuristic (not gradient descent), decoupling load balancing from the language modeling loss entirely.

Why 256 fine-grained experts instead of 8 coarse ones? Finer granularity means each expert can specialize more precisely. An expert that handles "Python variable assignment" is more useful than one that handles "all of programming." The cost is more complex routing, but with 8-of-256 selection, the activation ratio (3.1%) is even sparser than 2-of-8 (25%).

The shared expert is another innovation. One expert processes every single token — it captures common patterns that all tokens need (basic grammar, formatting, common words). The routed experts then only need to capture specialized knowledge. Think of it as a "base layer" that everyone gets, plus "specialist consultants" routed by topic.

jax.lax.ragged_dot: Variable-Length Expert Batches

Here is a practical problem that matters for your capstone implementation. When you route tokens to experts, each expert receives a different number of tokens. Expert 3 might get 50 tokens while expert 7 gets 120. You need to do batched matrix multiplication where each batch element has a different length. In dense models, every batch element has the same shape, so standard matmul works. In MoE, you need a ragged (variable-length) batched dot product.

JAX provides jax.lax.ragged_dot for exactly this. It takes:

python
import jax
import jax.numpy as jnp

# x: [total_tokens, d_model] — all tokens concatenated
# w: [num_experts, d_model, expert_dim] — stacked expert weights
# group_sizes: [num_experts] — how many tokens assigned to each expert
#   e.g., [50, 80, 120, 45, 90, 110, 65, 75] — sums to total_tokens

# ragged_dot computes expert_i's matmul on its assigned chunk of x
output = jax.lax.ragged_dot(
    x,              # [total_tokens, d_model]
    w,              # [num_experts, d_model, expert_dim]
    group_sizes,    # [num_experts] — token counts per expert
)
# output: [total_tokens, expert_dim]
# Equivalent to: for each expert i, slice x[start:start+group_sizes[i]]
# and compute matmul(slice, w[i]). Concatenate results.

The alternative without ragged_dot is to pad every expert's batch to the maximum length, waste compute on padding tokens, and then mask out the results. For highly imbalanced routing, this can waste 2-3x FLOPs. ragged_dot avoids the waste by handling variable-length segments natively on the accelerator.

The Switch Transformer: Where It All Began

Before Mixtral and DeepSeek-V3, there was the Switch Transformer (Fedus et al., 2022). It simplified the original MoE design (Shazeer et al., 2017) from top-2 routing to top-1 — each token goes to exactly one expert. This sounds wasteful, but it has a crucial advantage: simpler communication patterns and higher training stability.

The Switch Transformer also introduced the capacity factor and the auxiliary load-balancing loss formulation that became standard. Its key finding: MoE scaling is efficient. A Switch Transformer with 1.6 trillion parameters (but only ~100B active) trained 4x faster than a compute-equivalent dense T5 model, reaching the same quality in 1/4 the time.

Gating Mechanisms: Beyond Simple Softmax

The router's gating function is a design choice with significant consequences. Different papers use different gating mechanisms, each with tradeoffs.

The Expert Choice mechanism deserves special attention. Instead of each token choosing its experts, each expert chooses its top-k tokens. This guarantees perfect load balance (each expert gets exactly the same number of tokens) but creates a different problem: some tokens may be chosen by zero experts and others by many. In practice, Expert Choice often works better than auxiliary-loss-based balancing because it eliminates the hyperparameter sensitivity of the balancing coefficient.

Chinchilla Laws for MoE

Chinchilla scaling laws for dense models say: for a compute budget C, the optimal model has N_opt ∝ C^0.5 parameters trained on D_opt ∝ C^0.5 tokens. But MoE breaks this equation because active parameters ≠ total parameters.

The practical implication: you cannot just add infinite experts. At some point, the communication overhead and the diminishing scaling returns of inactive parameters make it better to increase expert size or add more active experts. DeepSeek-V3's choice of 256 experts with 8 active represents a point on this frontier that was determined empirically.

MoE Router Implementation in JAX

jax
import jax
import jax.numpy as jnp
import flax.linen as nn

class TopKRouter(nn.Module):
    """Top-k expert routing with load-balancing auxiliary loss."""
    num_experts: int = 8
    top_k: int = 2

    @nn.compact
    def __call__(self, x):
        # x: [batch, seq_len, d_model]
        B, S, D = x.shape

        # Router: single linear projection to expert logits
        logits = nn.Dense(self.num_experts, name="gate")(x)  # [B, S, N]

        # Top-k selection
        top_k_logits, top_k_indices = jax.lax.top_k(logits, self.top_k)

        # Routing weights: softmax over selected experts only
        gates = jax.nn.softmax(top_k_logits, axis=-1)  # [B, S, k]

        # ── Auxiliary load-balancing loss ──
        # f_i = fraction of tokens routed to expert i
        router_probs = jax.nn.softmax(logits, axis=-1)  # [B, S, N]
        expert_mask = jax.nn.one_hot(top_k_indices, self.num_experts)  # [B,S,k,N]
        expert_mask = expert_mask.sum(axis=-2)  # [B, S, N] — 1 if expert selected

        f_i = expert_mask.mean(axis=(0, 1))     # [N] — fraction of tokens per expert
        p_i = router_probs.mean(axis=(0, 1))  # [N] — mean router prob per expert

        aux_loss = self.num_experts * (f_i * p_i).sum()

        return gates, top_k_indices, aux_loss

class MoELayer(nn.Module):
    """Mixture of Experts FFN layer."""
    num_experts: int = 8
    top_k: int = 2
    expert_dim: int = 2048  # hidden dim of each expert FFN

    @nn.compact
    def __call__(self, x):
        B, S, D = x.shape

        # Route tokens to experts
        gates, indices, aux_loss = TopKRouter(
            self.num_experts, self.top_k
        )(x)  # gates: [B,S,k], indices: [B,S,k]

        # Compute ALL expert outputs (simple but wasteful)
        # Production code uses scatter-gather for true sparsity
        expert_outputs = []
        for i in range(self.num_experts):
            # Each expert is a 2-layer FFN with GELU
            h = nn.Dense(self.expert_dim, name=f"expert_{i}_up")(x)
            h = jax.nn.gelu(h)
            h = nn.Dense(D, name=f"expert_{i}_down")(h)
            expert_outputs.append(h)

        # Stack: [num_experts, B, S, D]
        all_experts = jnp.stack(expert_outputs, axis=0)

        # Gather selected expert outputs using indices
        # indices: [B, S, k] — which experts for each token
        selected = jnp.take_along_axis(
            all_experts,
            indices.transpose(2, 0, 1)[..., None].broadcast_to(
                self.top_k, B, S, D),
            axis=0
        ).transpose(1, 2, 0, 3)  # [B, S, k, D]

        # Weighted sum by gate values
        output = (selected * gates[..., None]).sum(axis=-2)  # [B, S, D]

        return output, aux_loss

MoE and Multi-Modality: Routing Across Modalities

A frontier topic that appears in research discussion interviews: how do you route tokens from different modalities (text, images, audio) through the same MoE? This is relevant because frontier models like Gemini and GPT-4o are multi-modal.

The challenge: text tokens and image patch tokens have very different statistical properties. Text tokens have a discrete, high-entropy distribution. Image patch tokens are continuous, locally correlated, and lower-entropy. If you use a single router for both, the router may route all image tokens to one set of experts and all text tokens to another — creating accidental modality-specific experts rather than concept-specific experts.

This is an active research area with no consensus answer. In an interview, the best response is to present the tradeoffs (shared routing gives more flexibility but risks modality collapse; reserved experts are safer but limit cross-modal transfer) and propose an experiment to decide (train both configurations, measure loss on text-only, image-only, and multi-modal benchmarks).

MoE and Reinforcement Learning from Human Feedback

An important practical consideration: MoE models behave differently under RLHF than dense models. The router's decisions can shift dramatically during the reward model training phase, because RLHF changes the distribution of "good" outputs.

MoE Routing Visualization

Expert Specialization: What Do Experts Learn?

A natural question: do MoE experts specialize in meaningful ways, or is the routing arbitrary? Empirical studies on trained MoE models reveal striking patterns.

In language models, experts often specialize by domain. One expert handles code, another handles scientific text, another handles conversational language. This happens without any explicit supervision — the router learns to route tokens to the expert that produces the best loss for that domain. Think of it as emergent department formation in a company: people naturally gravitate toward the work they are best at.

In the addition task from our capstone, we expect even finer specialization. Some experts should learn to handle carry operations (the hard part), while others handle simple single-digit addition (the easy part). Some might specialize by digit position — handling the ones place, tens place, etc. Analyzing this specialization after training is one of the most interesting parts of the capstone.

MoE Failure Modes: A Field Guide

When MoE training goes wrong, the symptoms are distinctive. Here is a diagnostic guide that will save you hours of debugging, both in the capstone and in production.

MoE at Inference Time: The Serving Challenge

Training MoE is hard. Serving MoE is harder. During training, you have a large batch of tokens that statistically covers all experts — every expert gets enough work to justify its GPU. During inference, a single request might only activate 2 experts, leaving the other 6 GPUs idle.

The simulation below shows tokens flowing through an MoE layer. Each token is routed to its top-2 experts. Watch how load balancing distributes tokens, and see what happens when you disable it — experts collapse.

Mixtral vs DeepSeek-V3 vs GShard: Comparing MoE Designs

Three landmark MoE designs represent different points on the design space. Understanding their differences teaches you the tradeoffs that matter.

The trend is clear: newer designs use more experts, finer granularity, and more sophisticated balancing. GShard proved MoE works at scale. Mixtral proved it works in the open-weights community. DeepSeek-V3 pushed the frontier with novel routing and training techniques.

MoE Training Budget: How Much Compute Do You Save?

Let's make the savings concrete with a worked example comparing training costs.

Implementing Permute-Compute-Unpermute

The efficient MoE implementation avoids computing all experts by physically sorting tokens by expert assignment. This is the permute-compute-unpermute pattern used in production systems.

jax
def efficient_moe_forward(x, expert_weights_up, expert_weights_down, gates, indices):
    """Efficient MoE: sort tokens by expert, compute, unsort."""
    B, S, D = x.shape
    x_flat = x.reshape(-1, D)  # [B*S, D]
    T = B * S

    # Step 1: PERMUTE — sort tokens by expert assignment
    # For top-2: each token appears twice (once per selected expert)
    expert_ids = indices.reshape(-1)  # [T*k] — flat list of expert assignments
    token_ids = jnp.repeat(jnp.arange(T), 2)  # [T*k] — which token each entry came from

    # Sort by expert id to group tokens for each expert together
    sort_order = jnp.argsort(expert_ids)
    sorted_expert_ids = expert_ids[sort_order]
    sorted_token_ids = token_ids[sort_order]

    # Gather sorted tokens
    sorted_x = x_flat[sorted_token_ids]  # [T*k, D]

    # Count tokens per expert (group_sizes for ragged_dot)
    group_sizes = jnp.bincount(sorted_expert_ids, length=NUM_EXPERTS)

    # Step 2: COMPUTE — run each expert on its tokens
    # Using ragged_dot or sequential per-expert matmul
    h = jax.lax.ragged_dot(sorted_x, expert_weights_up, group_sizes)
    h = jax.nn.gelu(h)
    y = jax.lax.ragged_dot(h, expert_weights_down, group_sizes)

    # Step 3: UNPERMUTE — scatter results back to original token order
    # Weight by gate values and sum the k=2 contributions per token
    gate_flat = gates.reshape(-1)  # [T*k]
    sorted_gates = gate_flat[sort_order]
    y_weighted = y * sorted_gates[:, None]

    # Scatter-add back to original positions
    output = jnp.zeros_like(x_flat)
    output = output.at[sorted_token_ids].add(y_weighted)

    return output.reshape(B, S, D)

Chapter 14: Pallas Kernels

Your MoE model is training on a TPU v4 pod. You profile the training step and discover that 35% of the step time is spent inside the MoE layer — specifically, inside jax.lax.ragged_dot. The operation handles variable-length expert batches correctly, but its generic implementation is not exploiting the specific structure of your problem. Your experts all have the same hidden dimension. Your batch sizes per expert, while variable, fall into a predictable range. You know that if you could write a custom kernel that tiles the computation exactly right for your access pattern, you could cut that 35% down to 20%.

This is where Pallas comes in. Pallas is JAX's framework for writing custom accelerator kernels — for both TPUs and GPUs — directly in Python. No CUDA C++. No XLA HLO surgery. You write Python functions decorated with a few Pallas-specific annotations, and the framework compiles them to efficient machine code for your target hardware.

What is Pallas?

Pallas is a kernel authoring API that sits between JAX's high-level NumPy-like interface and the low-level hardware instruction set. Think of it as JAX's answer to Triton (for NVIDIA GPUs) — except Pallas targets both TPUs and GPUs with a unified programming model.

The key insight of Pallas is declarative data movement. Instead of manually loading tiles from global memory into shared memory (as you would in CUDA), you declare what data each block needs via a BlockSpec, and the compiler figures out how to move it. This is a fundamental difference from CUDA, where you write the loads yourself.

The BlockSpec Abstraction

A BlockSpec describes the shape and indexing pattern of one input or output. It tells Pallas: "For grid cell (i, j), give me this slice of the array." The compiler then generates the appropriate memory load/store instructions.

Why is this better than writing loads manually? Three reasons. First, the compiler can prefetch: it knows what data the next grid cell will need before you ask. Second, it can double-buffer: while the kernel computes on one tile, the hardware asynchronously loads the next tile. Third, it can optimize memory placement: on TPU, it decides whether to put your tile in VMEM (fast, small) or HBM (slow, large). You do not have to think about any of this.

Grid and Block Dimensions

A Pallas kernel executes on a grid — a multi-dimensional array of blocks. Each block processes one tile of the computation. The grid dimensions determine the parallelism and tiling strategy.

The k-dimension loop is the reduction axis. On each iteration, we load a tile of A and a tile of B, multiply them, and accumulate into the output tile. The accumulator stays in fast on-chip memory (VMEM on TPU, registers on GPU) across all k iterations.

Memory Spaces: VMEM, SMEM, HBM

Understanding the memory hierarchy is essential for writing fast kernels.

The key ratio is compute-to-memory. A TPU v4 core does ~275 TFLOPS but can only load ~1.2 TB/s from HBM. That means for every byte you load, you need to do ~230 FLOPs to keep the compute units busy. If your kernel does fewer FLOPs per byte loaded, it is memory-bound and the compute units are starving. If it does more, it is compute-bound and the memory system can keep up.

A First Pallas Kernel: Vector Add

Let's start with the simplest possible Pallas kernel — adding two vectors. This is memory-bound and will not be faster than jnp.add, but it teaches the API structure.

jax/pallas
import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl

def vector_add_kernel(x_ref, y_ref, o_ref):
    """Pallas kernel body. Operates on one block at a time.
    x_ref, y_ref, o_ref are Ref objects — views into the block.
    """
    o_ref[...] = x_ref[...] + y_ref[...]

def vector_add(x, y):
    """Launch the Pallas kernel over a 1D grid."""
    n = x.shape[0]
    BLOCK = 512  # process 512 elements per block
    grid = (n // BLOCK,)  # 1D grid

    return pl.pallas_call(
        vector_add_kernel,
        grid=grid,
        in_specs=[
            pl.BlockSpec(block_shape=(BLOCK,), index_map=lambda i: (i,)),
            pl.BlockSpec(block_shape=(BLOCK,), index_map=lambda i: (i,)),
        ],
        out_specs=pl.BlockSpec(block_shape=(BLOCK,), index_map=lambda i: (i,)),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
    )(x, y)

# Test
x = jnp.ones(4096)
y = jnp.ones(4096) * 2
result = vector_add(x, y)  # [3.0, 3.0, 3.0, ...]
print(result[:5])  # [3. 3. 3. 3. 3.]