Stage 0: The machine — Rudrite Kernels

mastery work · this chapter0/3

auto
auto

manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only

the chapter

Two speeds, one ridge

A chip has two speeds: how fast it computes and how fast it feeds. A TPU v5e can do 1.97e14 bf16 FLOPs per second, but its HBM can only deliver 8.2e11 bytes per second. Divide the two and you get the ridge: about 240 FLOPs per byte. Any op that does fewer FLOPs per byte it moves is memory-bound, and no amount of compute cleverness will speed it up; any op above the ridge is compute-bound, and no amount of bandwidth will. This one division explains most TPU performance mysteries you will ever meet.

the chapter

The units and the scratchpad

The compute lives in two units. The MXU is a systolic array (128x128 on most generations, 256x256 on v6e) built for matmuls; the VPU is a vector unit that handles everything elementwise. Between them and HBM sits VMEM, roughly 128 MiB of software-managed scratchpad on v5e. Data must be staged in VMEM before compute can touch it, and almost all kernel engineering is choreography of that staging. EX·06 below shows the hierarchy with its real widths.

the chapter

The lattice rule

One lattice rule to memorize now, because every kernel in this track obeys it: TPU tiles want their last two dimensions in multiples of (8, 128), and packing doubles the sublane count for narrower dtypes (bf16 packs (16, 128)). Blocks that violate the lattice compile badly or not at all.

the chapter

Predict, then measure

The habit this stage installs: predict before you measure, every time. Work the arithmetic in EX·01, commit to a number, then run the lab on a real chip and explain every miss. An engineer who predicts within 2x and can name the reason for each miss understands the machine; one who only measures is collecting trivia.

Predict before you measure, every time.

hands on the machine

Instruments

EX·06 the memory hierarchy, to its real widths

constants from jax-ml.github.io/scaling-book (retrieved 2026-07-26) · nothing invented

EX·01 the roofline playground

chipopN = 2048

intensity: 682.7 FLOPs/byte
verdict: compute-bound
attainable: 197 TFLOP/s
floor latency: 87.2 µs

chip constants from jax-ml.github.io/scaling-book (retrieved 2026-07-26) · all values computed, nothing invented

deliverables

What gets built

Roofline estimates for five ops, by hand, reconciled against XProf measurements
A step-time floor estimate for a 7B forward pass from chip constants alone
The working vocabulary: MXU, VPU, VMEM, HBM, the (8, 128) tiling lattice, ICI

sources

Readings

Scaling book · All about TPUs · the backbone text; take chip constants from its tables, not from memory
Scaling book · Rooflines · the reasoning framework this whole stage drills
Cloud TPU system architecture · skim for v5e/v6e specifics

the plan

The work, in order

week 1

Read both scaling-book chapters in full before touching a notebook
By hand: classify five ops (large matmul, skinny matmul, elementwise add, long-axis softmax, embedding lookup) as compute- or memory-bound per chip generation
Estimate the step-time floor for a 7B forward pass at batch 8, seq 4096 on one v5e from constants alone
Run LAB·0.1 on a Colab TPU and reconcile three estimates against XProf; explain every miss in a sentence

run it yourself

Labs

LAB·0.1 ~2 h · Colab TPU

Rooflines by hand

open in Colab → runs from this repo, save a copy to keep it

single source of truth

From the notebook

LAB·0.1Rooflines by hand · extracted from the notebook

# scaling-book chip constants (bf16), retrieved 2026-07-26
CHIPS = {
    "v5e": {"flops": 1.97e14, "hbm_bw": 8.2e11},
    "v6e": {"flops": 9.20e14, "hbm_bw": 1.6e12},
}
for name, c in CHIPS.items():
    print(f"{name}: ridge = {c['flops'] / c['hbm_bw']:.0f} FLOPs/byte")

def predict(name, flops, bytes_moved, chip):
    c = CHIPS[chip]
    intensity = flops / bytes_moved
    ridge = c["flops"] / c["hbm_bw"]
    bound = "compute" if intensity >= ridge else "memory"
    floor_s = max(flops / c["flops"], bytes_moved / c["hbm_bw"])
    return {"op": name, "intensity": intensity, "bound": bound, "floor_us": floor_s * 1e6}

N = 4096
ops = [
    predict("matmul NxNxN",      2 * N**3,     2 * 3 * N**2, "v5e"),
    predict("matmul skinny 8xNxN", 2 * 8 * N**2, 2 * (8*N + N*N + 8*N), "v5e"),
    predict("elementwise add",   N**2,         2 * 3 * N**2, "v5e"),
    predict("softmax rows",      5 * N**2,     2 * 2 * N**2, "v5e"),
    predict("reduce_sum",        N**2,         2 * (N**2 + N), "v5e"),
]
for o in ops:
    print(f"{o['op']:>18}: {o['intensity']:8.1f} F/B  {o['bound']:>7}-bound  floor {o['floor_us']:9.1f} us")

def measure(fn, *args, reps=20):
    fn(*args).block_until_ready()  # compile + warm
    times = []
    for _ in range(reps):
        t0 = time.perf_counter()
        fn(*args).block_until_ready()
        times.append(time.perf_counter() - t0)
    return float(np.median(times)) * 1e6  # us

results = []
if ON_TPU:
    key = jax.random.key(0)
    a = jax.random.normal(key, (N, N), jnp.bfloat16)
    b = jax.random.normal(key, (N, N), jnp.bfloat16)
    s = jax.random.normal(key, (8, N), jnp.bfloat16)

    cases = {
        "matmul NxNxN": (jax.jit(lambda x, y: x @ y), a, b),
        "matmul skinny 8xNxN": (jax.jit(lambda x, y: x @ y), s, b),
        "elementwise add": (jax.jit(lambda x, y: x + y), a, b),
        "softmax rows": (jax.jit(lambda x: jax.nn.softmax(x, axis=-1)), a),
        "reduce_sum": (jax.jit(lambda x: jnp.sum(x, axis=-1)), a),
    }
    for name, (fn, *args) in cases.items():
        us = measure(fn, *args)
        pred = next(o for o in ops if o["op"] == name)
        results.append({"op": name, "measured_us": round(us, 1),
                        "predicted_us": round(pred["floor_us"], 1),
                        "ratio": round(us / pred["floor_us"], 2)})
        print(f"{name:>18}: measured {us:9.1f} us  predicted {pred['floor_us']:9.1f} us  ratio {us / pred['floor_us']:5.2f}x")
else:
    print("Skipped: no TPU in this runtime.")

source of truth: labs/stage-0/lab-0.1-rooflines.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates

the self-test, before the gate

Can you

0/7

·Compute the ridge point (FLOPs per byte) for a TPU generation from its published FLOP/s and HBM bandwidth, without looking either number up twice.
·Given an op and its shapes, classify it compute-bound or memory-bound in under a minute, and say which chip generation you assumed.
·Name what lives in VMEM versus HBM versus SMEM, and say why VMEM residency is the entire point of a Pallas kernel.
·State the (8, 128) tiling lattice rule from memory, and say what changes for bf16 packing.
·Estimate a step-time floor for a forward pass from chip constants alone, given batch size, sequence length, and parameter count, no profiler needed.
·Given a roofline prediction and a measured XProf number that disagree by more than 2x, name a concrete reason for the miss instead of shrugging at "overhead."
·Tell the difference between ICI and DCN, and say which one a collective inside one pod uses.

checked state lives in your browser only · the gate below is the public half

pass or fail, in public

The gate

GATE 00 OPEN

criterion	measured
Latency predictions for five ops within 2x of measured, every miss explained	predictions computed (see bench); measured half pending a TPU run

gates pass only on real hardware, published pass or fail · est./computed rows are predictions holding the slot until a run replaces them · records in bench/

Public artifact for this stage: how a TPU actually spends its time.