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

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

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

mastery work · this chapter0/5
  2. go →
  3. go →
  4. auto
  5. 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 HBM; 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 map). 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 DMA, 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 MXU 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

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.

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-07-26).
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 QUEUED
criterionmeasured
Tiled matmul within 15% of XLA for 4096³ bf16 est. ~1.1x XLA at 4096³ (pending run)
Fused softmax beats the unfused XLA chain at rows of 32k+ est. beats unfused chain (pending 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: Pallas from zero, with the algorithm/schedule split color-coded.