This is what leaves JAX and crosses the seam, the form every chapter at the waist works on.
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%1 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%2 = stablehlo.dot_general %arg0, %arg3, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%3 = stablehlo.transpose %1, dims = [1, 0] : (tensor<16x32xf32>) -> tensor<32x16xf32>
%4 = stablehlo.dot_general %0, %3, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x16xf32>) -> tensor<16x16xf32>
%cst = stablehlo.constant dense<3.200000e+01> : tensor<f32>
%5 = stablehlo.sqrt %cst : tensor<f32>
%6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<f32>) -> tensor<16x16xf32>
%7 = stablehlo.divide %4, %6 : tensor<16x16xf32>What your code becomes
Your Python is not what runs. JAX traces it into a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → (a flat list of primitive ops), lowers that to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → (the portable tensor IR every framework converges on), and hands it to XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/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.
Fusion, and its exact limit
XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s big move is fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →: 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 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 → traffic that exists purely because the algorithm is multi-pass.
An algorithm failure
That sentence deserves its own paragraph, because this whole track pivots on it: the spill is not a fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → 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 machine's account
This stage now closes its own loop on hardware. LAB·2.3 captures a trace of the very program whose spill you just found, reads the device plane in code, and converts hardware time to bytes through measured bandwidth: the gate below passed at 1.2% agreement on exactly that chain. The same trace held a discovery worth knowing before stage 3: the current compiler pattern-matches this program and dispatches its own online-softmax kernel, which chapter 08's guide and the GYM·08 timeline teach you to see for yourself.
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.
Instruments
hover or tab through any line: the same computation lights up at every layer
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
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 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)
Labs
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())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.")Can you
The gate
| criterion | measured |
|---|---|
| Spill-size estimate from the HLO dump matches the profiler within 20% | profiled on v6e-1, closed at 1.2% on a 20% bar: the S-carrying ops cost 217.6 µs of hardware time, and at the measured copy bandwidth (1287.5 GB/s, 80.5% of nameplate) that is 280.2 MB against the 276.8 MB estimate. Two lessons rode along: the cost model said 155.2 MB because it cannot see inside custom calls, and the timeline revealed XLA:TPU swapping in its own online-softmax kernel for this exact pattern. |