the path · 0/15
start the path

the path · chapter 11 of 15 · part ii, the track

stage 01
title Pallas fundamentals
schedule weeks 2-4
gate PASSED

Write, benchmark, and profile correct single-chip Pallas kernels; read production kernels fluently.

mastery work · this chapter0/8
  2. go →auto
  3. go →auto
  4. go →auto
  5. go →
  6. go →
  7. go →auto
  8. auto
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
the chapter

Refs, blocks, and the grid

A Pallas kernel is a function over refs: windows into arrays that the runtime has already staged into VMEM. You never write loads or stores against HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu →; you write what one grid step does to its window, and two other objects decide everything else. The BlockSpec says how an array is carved into blocks and which block a given grid position sees (its index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas →). The grid says how many steps run and in what order. Keeping "what the step computes" separate from "which data the step sees" in your head from day one is the discipline everything later builds on; the site colors the first copper and the second steel everywhere.

the chapter

The grid is a pipeline

The part nobody tells you early enough: on TPU, the grid is a software pipeline, not a batch of parallel threads. Step k computes on its blocks while the runtime streams step k+1's blocks in behind it. You get double buffering without writing a DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas →, but only if your blocks are sized so transfer time hides under compute time. EX·02 plays the pipelined schedule; EX·07 plays the same matmul with the pipeline off, so you can watch the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → idle between loads. The difference between those two animations is most of single-chip performance engineering.

the chapter

The carried accumulator

The other load-bearing pattern is the carried accumulator. A reduction axis in the grid means the same output block is revisited once per step along that axis, initialized on the first visit and accumulated into after. Matmul carries C over the K axis here; flash attention will carry (m, l, acc) over KV blocks in stage 3. Same shape, bigger algebra.

the chapter

Reading the layer below

When a kernel misbehaves, drop one layer and read what Pallas actually built. pallas_call(..., debug=True) prints two artifacts at lowering time: the kernel jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → (your body, traced against refs) and the MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → module, the MLIR the TPU backend receives. Everything you wrote is still visible in it: BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → index mapsThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → become @transform functions, pl.when becomes an scf.if, jnp.dot becomes tpu.matmul, and every vector type wears the shape the (8, 128) lattice demands. Chapter 07 walks the layer in depth, the GYM·05 instrument holds three of this stage's kernels open at it with hover sync, and LAB·1.4 has you print your own.

The reason to build the habit now: the two hardest exhibits in the museum, the VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → overflow and the lattice violation, are this layer speaking. An engineer who has read a MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → module recognizes "window allocation" as a BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → made physical and "divisible by 8 and 128" as the register tile. One who has not sees noise.

the chapter

Numbers carry provenance

And the habit: every measured number states its chip, dtype, shapes, and method (warmup, block_until_ready, median of N). Numbers without provenance are noise, and the bench page refuses them.

go deeper, in order

Lessons

  1. 01BlockSpec and the index mapA carve is not a slicing convenience. It is the contract that decides which bytes are resident on which step, and what garbage sits past the edge. ·
  2. 02The grid and the pipelineThe grid promises that every point runs. Every other thing you rely on comes from the backend underneath it. ·
  3. 03The scalar worldA kernel that decides has to keep the number it decides on somewhere, and a vector register is the wrong place. ·
hands on the machine

Instruments

EX·02 tiled matmul, double-buffered
HBMAA0A1A2A3BB0B1B2B3VMEMa0emptyb0emptya1emptyb1emptyaccC (zeros)MXU · idle

01/06grid step 0 · DMA A0, B0 into the first slot pair; the MXU has no data yet and waits
frames computed from the BlockSpec schedule (K-blocked, depth-2 pipeline) · no invented values; compute is shown as expressions
EX·07 the same matmul with the pipeline off
HBMAA0A1A2A3BB0B1B2B3VMEMaemptybemptyaccC (zeros)MXU · idle

01/09load 1 · DMA A0, B0 into VMEM; the MXU sits idle the whole transfer
the depth-1 schedule: compare frame count and MXU idle time against EX·02 above · no invented values
EX·09 the BlockSpec sandbox

legal: BlockSpec(256, 256) tiles the (2048, 2048) bf16 array

M = 2048N = 20480,0

8 × 8 = 64 blocks launched

grid
8 × 8
sublane rule
bm % 16 == 0
lane rule
bn % 128 == 0
bytes/block
131 KB
VMEM working set (3 blocks in flight)393 KB / 134 MB

VMEM budget 134,217,728 B for TPU v5e, source: jax-ml.github.io/scaling-book/tpus/ (retrieved 2026-08-06).
all legality arithmetic computed live: lattice rules, grid, and the v5e VMEM budget · diagnostic wording is ours, the arithmetic is the chip's
EX·12 the schedule and its lines, locked together
tiled matmul, double-buffered

grid step 0 · DMA A0, B0 into the first slot pair; the MXU has no data yet and waits

a0
empty
b0
empty
a1
empty
b1
empty
acc
C (zeros)
transfers
A0 → a0, B0 → b0
compute
idle
1def matmul_kernel(a_ref, b_ref, o_ref):
2    k = pl.program_id(2)
3    @pl.when(k == 0)
4    def _init():
5        o_ref[...] = jnp.zeros_like(o_ref)
6    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32).astype(o_ref.dtype)
7
8def matmul(a, b, bm=128, bn=128, bk=128):
9    m, k = a.shape
10    _, n = b.shape
11    return pl.pallas_call(
12        matmul_kernel,
13        grid=(m // bm, n // bn, k // bk),
14        in_specs=[
15            pl.BlockSpec((bm, bk), lambda i, j, kk: (i, kk)),
16            pl.BlockSpec((bk, bn), lambda i, j, kk: (kk, j)),
17        ],
18        out_specs=pl.BlockSpec((bm, bn), lambda i, j, kk: (i, j)),
19        out_shape=jax.ShapeDtypeStruct((m, n), a.dtype),
20        interpret=INTERP,
21    )(a, b)

the schedule step on the left is these lines on the right

01/06
each trace frame maps to the exact lines of the matmul walk that produce it · step either side, both move
line by line, at your pace

Guided walks

WALK·matmul tiled matmul: the reduction kernel and its pallas_call
tiled matmul: the reduction kernel and its pallas_call
1def matmul_kernel(a_ref, b_ref, o_ref):
2    k = pl.program_id(2)
3    @pl.when(k == 0)
4    def _init():
5        o_ref[...] = jnp.zeros_like(o_ref)
6    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32).astype(o_ref.dtype)
7
8def matmul(a, b, bm=128, bn=128, bk=128):
9    m, k = a.shape
10    _, n = b.shape
11    return pl.pallas_call(
12        matmul_kernel,
13        grid=(m // bm, n // bn, k // bk),
14        in_specs=[
15            pl.BlockSpec((bm, bk), lambda i, j, kk: (i, kk)),
16            pl.BlockSpec((bk, bn), lambda i, j, kk: (kk, j)),
17        ],
18        out_specs=pl.BlockSpec((bm, bn), lambda i, j, kk: (i, j)),
19        out_shape=jax.ShapeDtypeStruct((m, n), a.dtype),
20        interpret=INTERP,
21    )(a, b)

01/07The kernel takes three references, not two arrays. a_ref and b_ref hold the input blocks for this grid step; o_ref holds the block of the output this step owns. Pallas has already sliced the full A and B matrices down to a (bm, bk) and (bk, bn) tile before the call, and k = pl.program_id(2) reads which K-block this grid step landed on. That line is schedule: it says where in the reduction the kernel is, not what to compute.
code verbatim from the lab notebook · step with the buttons or j/k
WALK·softmax row softmax: one kernel, one grid axis
row softmax: one kernel, one grid axis
1def softmax_kernel(x_ref, o_ref):
2    x = x_ref[...].astype(jnp.float32)
3    m = jnp.max(x, axis=-1, keepdims=True)
4    e = jnp.exp(x - m)
5    o_ref[...] = (e / jnp.sum(e, axis=-1, keepdims=True)).astype(o_ref.dtype)
6
7def softmax(x, rows_per_block=8):
8    n, d = x.shape
9    return pl.pallas_call(
10        softmax_kernel,
11        grid=(n // rows_per_block,),
12        in_specs=[pl.BlockSpec((rows_per_block, d), lambda i: (i, 0))],
13        out_specs=pl.BlockSpec((rows_per_block, d), lambda i: (i, 0)),
14        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
15        interpret=INTERP,
16    )(x)

01/07softmax_kernel takes just two refs, one input block and one output block, no accumulator to carry across grid steps the way the matmul kernel needed. Casting x to float32 up front is the same move as in the flash kernels: the exponential a couple of lines down needs full precision even if the data arrives and leaves in a lower-precision dtype.
code verbatim from the lab notebook · step with the buttons or j/k
deliverables

What gets built

  • Six kernels in order: add, transpose, tiled matmul, fused softmax, LayerNorm, matmul+GELU
  • A benchmarking habit: warmup, block_until_ready, median-of-N, chip stated with every number
  • Annotated readings of production kernels, every line marked algorithm or schedule
sources

Readings

the plan

The work, in order

week 2 · mechanics

  • LAB·1.1: add and transpose; internalize refs, index maps, and interpret mode
  • LAB·1.2 first half: the tiled matmul with a carried accumulator, correctness only
  • Answer in writing: why is K the innermost grid dim, and what breaks if it is outermost?

week 3 · fusion and reductions

  • LAB·1.3: fused softmax and LayerNorm; one HBM round-trip each
  • Blocked rows, unblocked features: write down why the reduction axis must fit the block
  • Read one Tokamax kernel line by line, margin-marking algorithm vs schedule

week 4 · pipelining and measurement

  • LAB·1.2 second half: the block-size sweep against XLA on a TPU runtime
  • LAB·1.4: profile the matmul, find the pipeline in the trace, starve it on purpose
  • Predict the VMEM budget line from stage 0 constants and confirm where compilation actually fails
run it yourself

Labs

single source of truth

From the notebook

LAB·1.1First kernels: add, transpose · extracted from the notebook
def add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]

def add(x, y):
    return pl.pallas_call(
        add_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        interpret=INTERP,
    )(x, y)

x = jax.random.normal(jax.random.key(0), (256, 256), jnp.float32)
y = jax.random.normal(jax.random.key(1), (256, 256), jnp.float32)
check("add", add(x, y), x + y)
def transpose_kernel(x_ref, o_ref):
    o_ref[...] = x_ref[...].T

def transpose(x, bm=128, bn=128):
    m, n = x.shape
    return pl.pallas_call(
        transpose_kernel,
        grid=(m // bm, n // bn),
        in_specs=[pl.BlockSpec((bm, bn), lambda i, j: (i, j))],
        out_specs=pl.BlockSpec((bn, bm), lambda i, j: (j, i)),
        out_shape=jax.ShapeDtypeStruct((n, m), x.dtype),
        interpret=INTERP,
    )(x)

check("transpose", transpose(x), x.T)
source of truth: labs/stage-1/lab-1.1-first-kernels.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates
LAB·1.2Tiled matmul with a carried accumulator · extracted from the notebook
def matmul_kernel(a_ref, b_ref, o_ref):
    k = pl.program_id(2)
    @pl.when(k == 0)
    def _init():
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32).astype(o_ref.dtype)
def matmul(a, b, bm=128, bn=128, bk=128):
    m, k = a.shape
    _, n = b.shape
    return pl.pallas_call(
        matmul_kernel,
        grid=(m // bm, n // bn, k // bk),
        in_specs=[
            pl.BlockSpec((bm, bk), lambda i, j, kk: (i, kk)),
            pl.BlockSpec((bk, bn), lambda i, j, kk: (kk, j)),
        ],
        out_specs=pl.BlockSpec((bm, bn), lambda i, j, kk: (i, j)),
        out_shape=jax.ShapeDtypeStruct((m, n), a.dtype),
        interpret=INTERP,
    )(a, b)

a = jax.random.normal(jax.random.key(0), (512, 512), jnp.float32)
b = jax.random.normal(jax.random.key(1), (512, 512), jnp.float32)
check("matmul 512", matmul(a, b), a @ b, tol=1e-3)
source of truth: labs/stage-1/lab-1.2-tiled-matmul.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates
LAB·1.3Fused softmax and LayerNorm · extracted from the notebook
def softmax_kernel(x_ref, o_ref):
    x = x_ref[...].astype(jnp.float32)
    m = jnp.max(x, axis=-1, keepdims=True)
    e = jnp.exp(x - m)
    o_ref[...] = (e / jnp.sum(e, axis=-1, keepdims=True)).astype(o_ref.dtype)

def softmax(x, rows_per_block=8):
    n, d = x.shape
    return pl.pallas_call(
        softmax_kernel,
        grid=(n // rows_per_block,),
        in_specs=[pl.BlockSpec((rows_per_block, d), lambda i: (i, 0))],
        out_specs=pl.BlockSpec((rows_per_block, d), lambda i: (i, 0)),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        interpret=INTERP,
    )(x)

x = jax.random.normal(jax.random.key(0), (256, 512), jnp.float32)
check("softmax", softmax(x), jax.nn.softmax(x, axis=-1), tol=1e-5)
def layernorm_kernel(x_ref, g_ref, b_ref, o_ref):
    x = x_ref[...].astype(jnp.float32)
    mu = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.mean((x - mu) ** 2, axis=-1, keepdims=True)
    y = (x - mu) * jax.lax.rsqrt(var + 1e-6)
    o_ref[...] = (y * g_ref[...] + b_ref[...]).astype(o_ref.dtype)

def layernorm(x, g, b, rows_per_block=8):
    n, d = x.shape
    return pl.pallas_call(
        layernorm_kernel,
        grid=(n // rows_per_block,),
        in_specs=[
            pl.BlockSpec((rows_per_block, d), lambda i: (i, 0)),
            pl.BlockSpec((1, d), lambda i: (0, 0)),
            pl.BlockSpec((1, d), lambda i: (0, 0)),
        ],
        out_specs=pl.BlockSpec((rows_per_block, d), lambda i: (i, 0)),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        interpret=INTERP,
    )(x, g.reshape(1, -1), b.reshape(1, -1))

g = jnp.ones(512); b = jnp.zeros(512)
ref = (x - x.mean(-1, keepdims=True)) * jax.lax.rsqrt(x.var(-1, keepdims=True) + 1e-6)
check("layernorm", layernorm(x, g, b), ref, tol=1e-4)
source of truth: labs/stage-1/lab-1.3-fused-softmax-layernorm.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates
LAB·1.4Pipelining and the profile · extracted from the notebook
source of truth: labs/stage-1/lab-1.4-pipelining-profile.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates
the self-test, before the gate

Can you

0/8
checked state lives in your browser only · the gate below is the public half
pass or fail, in public

The gate

GATE 01 PASSED
criterionmeasured
Tiled matmul within 15% of XLA for 4096³ bf16 measured on v6e-1: the first schedule ran 1.65x behind XLA; retuned with dimension_semantics and a (2048, 1024, 512) block, 384.6 µs vs 353.0 µs, 1.09x. Inside the bar, and the tuning IS the lesson.
Fused softmax beats the unfused XLA chain at rows of 32k+ measured on v6e-1: rows=1024 blocks give 222.7 µs vs 379.5 µs unfused, and beat XLA's own fusion (267.6 µs). The rows=64 schedule that lost by 27% was the same kernel; only the schedule changed.
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/