Stage 2: The IR stack — Rudrite Kernels

the chapter

What your code becomes

Your Python is not what runs. JAX traces it into a jaxpr (a flat list of primitive ops), lowers that to StableHLO (the portable tensor IR every framework converges on), and hands it to XLA, which makes the performance decisions and emits code for the backend. Each layer is readable, and reading them is a skill worth drilling until it is boring, because the IR states everything Python hides: broadcasts made explicit, dtypes chosen for accumulators, the exact contraction dimensions of every matmul. EX·05 below holds one attention program open at three layers at once; hover any line and its counterparts light up.

the chapter

Fusion, and its exact limit

XLA's big move is fusion: merging adjacent ops so intermediates stay in fast memory instead of round-tripping through HBM. Fusion is powerful and it has a precise limit: XLA fuses along dataflow edges, but it cannot change the algorithm. Naive softmax needs the row max before any exponential, so the S = QK^T score matrix is computed, spilled, and re-read no matter how well XLA fuses around it. At seq 8192 in bf16 that spill is 134 MB written and read back: 268 MB of HBM traffic that exists purely because the algorithm is multi-pass.

the chapter

An algorithm failure

That sentence deserves its own paragraph, because this whole track pivots on it: the spill is not a fusion failure, it is an algorithm failure. No pass in the compiler can fix it; a theorem can, and stage 3 proves that theorem. This stage's job is to see the spill in the compiler's own output so you never again confuse "the compiler fused it" with "the memory traffic is gone."

The spill is not a fusion failure. It is an algorithm failure.

the chapter

The fluency drill

The fluency drill: dot_general dimension_numbers, decoded on sight. Which axes contract, which batch. It feels pedantic for a week and then every IR you ever read becomes legible.

hands on the machine

Instruments

EX·05 the IR x-ray · one program, three layers

JAX source

jaxpr

StableHLO

hover or tab through any line: the same computation lights up at every layer

generated with jax 0.4.38 on cpu · jaxpr and StableHLO are backend-independent · shapes (128, 64) bf16

deliverables

What gets built

A table mapping one program across jaxpr, StableHLO, and optimized HLO
An annotated HLO dump of naive attention: every fusion marked, the HBM spill found and sized
A trace of a real production operator, cataloged against the toy version

sources

Readings

StableHLO spec · skim the op list in one sitting; it is smaller than you fear
MaxText · pull one real operator and trace it; production jaxprs teach what toys cannot

the plan

The work, in order

week 5

LAB·2.1: trace one function through jaxpr, StableHLO, and optimized HLO; build the mapping table
Drill dot_general dimension_numbers until decoding is instant
LAB·2.2: predict the naive-attention spill in bytes, then find it in the TPU dump and confirm within 20% in the profiler
Trace a MaxText operator and catalog what the production jaxpr contains that your toy does not (remat markers, sharding constraints, dtype casts)

run it yourself

Labs

LAB·2.1 ~3 h · any machine (interpret mode)

The lowering ladder, traced

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

LAB·2.2 ~3 h · Colab TPU

Finding the spill in naive attention

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

single source of truth

From the notebook

LAB·2.1The lowering ladder, traced · extracted from the notebook

def f(x, w, b):
    h = x @ w + b
    return jax.nn.relu(h).mean(axis=-1)

args = [jax.ShapeDtypeStruct(s, jnp.bfloat16) for s in [(32, 64), (64, 128), (128,)]]

print(jax.make_jaxpr(f)(*args))

print(jax.jit(f).lower(*args).as_text())

source of truth: labs/stage-2/lab-2.1-lowering-ladder.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates

LAB·2.2Finding the spill in naive attention · extracted from the notebook

S, D = 8192, 128

def naive_attention(q, k, v):
    s = q @ k.T
    m = jnp.max(s, axis=-1, keepdims=True)
    p = jnp.exp(s - m)
    l = jnp.sum(p, axis=-1, keepdims=True)
    return (p / l) @ v

args = [jax.ShapeDtypeStruct((S, D), jnp.bfloat16)] * 3

# the spill, predicted from first principles: the SxS score matrix in bf16,
# written once after the first matmul and read back for the softmax chain
score_bytes = S * S * 2
print(f"score matrix: {score_bytes / 1e6:.0f} MB; at least written + read once = {2 * score_bytes / 1e6:.0f} MB of HBM traffic")
print(f"at v5e HBM bandwidth (8.2e11 B/s): {2 * score_bytes / 8.2e11 * 1e3:.1f} ms of pure spill time")

if ON_TPU:
    hlo = jax.jit(naive_attention).lower(*args).compile().as_text()
    lines = [ln for ln in hlo.splitlines() if "8192,8192" in ln and ("fusion" in ln or "custom-call" in ln or "= bf16" in ln)]
    print(f"{len(lines)} lines mention the full score matrix; the ones that are fusion outputs are your spills:")
    for ln in lines[:12]:
        print(ln.strip()[:160])
else:
    print("Optimized-HLO hunt needs the TPU runtime; the prediction above works anywhere.")

source of truth: labs/stage-2/lab-2.2-finding-the-spill.ipynb · cells tagged site:show render here; edit the notebook, rebuild, the page updates

the self-test, before the gate

Can you

0/7

·Decode the dimension_numbers of any dot_general on sight: which axes contract, which batch.
·Given a JAX function, produce its jaxpr, StableHLO, and optimized HLO yourself, and map one equation across all three.
·Find the HBM spill in naive attention's HLO dump and estimate its size in bytes from the sequence length and dtype alone.
·Explain in one sentence why fusion cannot fix the spill you just found, and what would.
·Given a StableHLO op you have never seen, tell whether it is a real computation or compiler bookkeeping (a broadcast, a convert, a reshape).
·Pull one real operator from a production model library and name what its jaxpr carries that your toy version doesn't (remat markers, sharding constraints, dtype casts).
·Tell a lax.scan from an unrolled lax.associative_scan just by reading the jaxpr, with no scan primitive present in the unrolled case.

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

pass or fail, in public

The gate

GATE 02 QUEUED

criterion	measured
Spill-size estimate from the HLO dump matches the profiler within 20%	computed: 268 MB round-trip at seq 8192 bf16; profiler confirm pending

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: what your JAX becomes, one program shown at every layer.