the xla path · 0/14
start the path
the gym
xla floor

serves xla-path chapters 04 · 05 · 08 · 10 · streaks check themselves off; the same skill counts on every floor that drills it · all floors

The compiler's decisions, drilled.

The xla path teaches you to read what the compiler did and why. These stations hold the evidence: decisions the framework made silently, fusion plans captured from a real TPU compile, and a hardware profile with an op nobody in the program wrote. Everything here is generated or captured, never retyped.

drills

Active recall, one tap

GYM·03 spot the decision
MLP blockstreak 0

click the first line where the framework makes a hidden broadcast explicit

program: mlp
find the line the framework wrote for you: upcasts, hidden broadcasts, remat marks · targets computed by scanning the real dumps · assigned in xla chapter 04
GYM·12 which pipeline filed this pass?
passadd-random-host-offloading
streak 0two real dumps

which pipeline filed this one?

cpu: jax 0.4.38, CPU, the course machine · tpu: jax 0.11, Colab TPU v6e, LAB·X2 and LAB·X4
17 names read out of two real dumps: this course's CPU compile and a Colab v6e · the shared names are the interesting ones · serves chapters 04 and 05
GYM·13 count the fusions
one elementwise opy = jnp.exp(x)
streak 0compiled and counted

how many fusion computations does the optimized module hold?

counted on cpu, jax 0.4.38 · a different backend draws these boundaries differently, which is the point of reading your own dump rather than trusting a rule
eight programs compiled and counted on jax 0.4.38 CPU · the reveal shows what each fusion actually holds · serves chapter 05
before /
after

Fusion, side by side

The StableHLO XLA received on the left, the optimized HLO it decided on the right, captured on a real TPU. Fusion ops light up in copper, and the lines still carrying each program's big intermediate are underlined: in the attention program that is the spill, sitting in the compiler's own plan. Chapter 05 teaches the policy that drew these boundaries.

GYM·07 the fusion x-ray
21 lines in → 31 lines decided · 3 fusion ops · 3 carry bf16[1024,1024]

find the fusion ops that carry bf16[1024,1024]: the spill, in the compiler's own plan

what XLA received (StableHLO)
module @jit_naive_attention attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
 func.func public @main(%arg0: tensor<1024x128xbf16>, %arg1: tensor<1024x128xbf16>, %arg2: tensor<1024x128xbf16>) -> (tensor<1024x128xbf16> {jax.result_info = "result"}) {
 %0 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<1024x128xbf16>) -> tensor<128x1024xbf16>
 %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1024x128xbf16>, tensor<128x1024xbf16>) -> tensor<1024x1024xbf16>
 %cst = stablehlo.constant dense<0xFF80> : tensor<bf16>
 %2 = stablehlo.reduce(%1 init: %cst) applies stablehlo.maximum across dimensions = [1] : (tensor<1024x1024xbf16>, tensor<bf16>) -> tensor<1024xbf16>
 %3 = stablehlo.broadcast_in_dim %2, dims = [0] : (tensor<1024xbf16>) -> tensor<1024x1xbf16>
 %4 = stablehlo.broadcast_in_dim %3, dims = [0, 1] : (tensor<1024x1xbf16>) -> tensor<1024x1024xbf16>
 %5 = stablehlo.subtract %1, %4 : tensor<1024x1024xbf16>
 %6 = stablehlo.exponential %5 : tensor<1024x1024xbf16>
 %7 = stablehlo.convert %6 : (tensor<1024x1024xbf16>) -> tensor<1024x1024xf32>
 %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
 %8 = stablehlo.reduce(%7 init: %cst_0) applies stablehlo.add across dimensions = [1] : (tensor<1024x1024xf32>, tensor<f32>) -> tensor<1024xf32>
 %9 = stablehlo.broadcast_in_dim %8, dims = [0] : (tensor<1024xf32>) -> tensor<1024x1xf32>
 %10 = stablehlo.convert %9 : (tensor<1024x1xf32>) -> tensor<1024x1xbf16>
 %11 = stablehlo.broadcast_in_dim %10, dims = [0, 1] : (tensor<1024x1xbf16>) -> tensor<1024x1024xbf16>
 %12 = stablehlo.divide %6, %11 : tensor<1024x1024xbf16>
 %13 = stablehlo.dot_general %12, %arg2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1024x1024xbf16>, tensor<1024x128xbf16>) -> tensor<1024x128xbf16>
 return %13 : tensor<1024x128xbf16>
 }
}
what XLA decided (optimized HLO)
HloModule jit_naive_attention, is_scheduled=true, entry_computation_layout={(bf16[1024,128]{1,0:T(8,128)(2,1)}, bf16[1024,128]{1,0:T(8,128)(2,1)}, bf16[1024,128]{1,0:T(8,128)(2,1)})->bf16[1024,128]{1,0:T(8,128)(2,1)}}, allow_spmd_sharding_propagation_to_parameters={true,true,true}, allow_spmd_sharding_propagation_to_output={true}, frontend_attributes={arg_layout_modes="default;default;default",arg_memory_spaces="0;0;0",out_layout_modes="default",out_memory_spaces="0"}
 
%bitcast_fusion (bitcast_input: bf16[1024,128]) -> bf16[1024,128] {
 %bitcast_input = bf16[1024,128]{1,0:T(8,128)(2,1)S(1)} parameter(0)
 ROOT %bitcast = bf16[1024,128]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input)
}
 
%bitcast_fusion.1 (bitcast_input.1: bf16[1024,128]) -> bf16[1024,128] {
 %bitcast_input.1 = bf16[1024,128]{1,0:T(8,128)(2,1)} parameter(0)
 ROOT %bitcast.1 = bf16[1024,128]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.1)
}
 
%fused_computation (param_0: bf16[1024,128], param_1: bf16[1024,128]) -> bf16[1024,1024] {
 %param_0 = bf16[1024,128]{1,0:T(8,128)(2,1)S(1)} parameter(0)
 %fusion.1 = bf16[1024,128]{1,0:T(8,128)(2,1)} fusion(%param_0), kind=kLoop, calls=%bitcast_fusion
 %param_1 = bf16[1024,128]{1,0:T(8,128)(2,1)} parameter(1)
 %fusion.2 = bf16[1024,128]{1,0:T(8,128)(2,1)} fusion(%param_1), kind=kLoop, calls=%bitcast_fusion.1
 ROOT %convolution.1 = bf16[1024,1024]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.1, %fusion.2), dim_labels=bf_oi->bf, metadata={op_name="jit(naive_attention)/dot_general" stack_frame_id=20}
}
 
ENTRY %main.3 (q.1: bf16[1024,128], k.1: bf16[1024,128], v.1: bf16[1024,128]) -> bf16[1024,128] {
 %q.1 = bf16[1024,128]{1,0:T(8,128)(2,1)} parameter(0), metadata={op_name="q"}
 %copy-start = (bf16[1024,128]{1,0:T(8,128)(2,1)S(1)}, bf16[1024,128]{1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%q.1), cross_program_prefetch_index=0, backend_config={"dma_priority":1,"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
 %v.1 = bf16[1024,128]{1,0:T(8,128)(2,1)} parameter(2), metadata={op_name="v"}
 %copy-start.1 = (bf16[1024,128]{1,0:T(8,128)(2,1)S(1)}, bf16[1024,128]{1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%v.1), backend_config={"dma_priority":1,"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
 %k.1 = bf16[1024,128]{1,0:T(8,128)(2,1)} parameter(1), metadata={op_name="k"}
 %copy-done = bf16[1024,128]{1,0:T(8,128)(2,1)S(1)} copy-done(%copy-start), backend_config={"dma_priority":1,"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
 %fusion = bf16[1024,1024]{1,0:T(8,128)(2,1)S(1)} fusion(%copy-done, %k.1), kind=kOutput, calls=%fused_computation, metadata={op_name="jit(naive_attention)/dot_general" stack_frame_id=20}, backend_config={"aliasing_operands":{"lists":[]},"convolution_algorithm_config":{"emitter":"EmitAllBatchInSublanes"},"flag_configs":[],"retry_config":{"retry_count":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"262144"}],"window_config":{"buffering_level":"2","cost_model_type":"COST_MODEL_TYPE_CLASSIC","estimated_bundle_count":"0","estimated_cycles":"3528","estimated_scoped_vmem_bytes":"0","estimated_vmem_bytes":"0","input_window_bounds":["128","1"],"is_mask":false,"iteration_bounds":["1","1","1"],"kernel_window_bounds":["128","1"],"ml_estimated_microseconds":0,"output_window_bounds":["128","8"],"pad_input_on_minor_dim":"0","pad_output_on_minor_dim":"0"}}
 %copy-done.1 = bf16[1024,128]{1,0:T(8,128)(2,1)S(1)} copy-done(%copy-start.1), backend_config={"dma_priority":1,"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
 ROOT %online-softmax = bf16[1024,128]{1,0:T(8,128)(2,1)} custom-call(%fusion, %copy-done.1), custom_call_target="tpu_custom_call", operand_layout_constraints={bf16[1024,1024]{1,0}, bf16[1024,128]{1,0}}, frontend_attributes={tiling="1024,1024"}, metadata={op_name="online-softmax"}, backend_config={"aliasing_operands":{"lists":[]},"custom_call_config":{"allow_input_fusion":[],"body":"bW9kdWxlIHsKICBmdW5jLmZ1bmMgQG1haW4oJWFyZzA6IGkzMiwgJWFyZzE6IGkzMiwgJWFyZzI6IG1lbXJlZjwxMDI0eDEwMjR4YmYxNiwgI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4+LCAlYXJnMzogbWVtcmVmPDEwMjR4MTI4eGJmMTYsICN0cHUubWVtb3J5X3NwYWNlPHZtZW0+PiwgJWFyZzQ6IG1lbXJlZjwxMDI0eDEyOHhiZjE2LCAjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPj4sICVhcmc1OiBtZW1yZWY8MTAyNHgxeGJmMTYsICN0cHUubWVtb3J5X3NwYWNlPHZtZW0+PiwgJWFyZzY6IG1lbXJlZjwxMDI0eDF4YmYxNiwgI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4+KSBhdHRyaWJ1dGVzIHtkaW1lbnNpb25fc2VtYW50aWNzID0gWyN0cHUuZGltZW5zaW9uX3NlbWFudGljczxwYXJhbGxlbD4sICN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+XSwgaXRlcmF0aW9uX2JvdW5kcyA9IGFycmF5PGk2NDogMSwgMT4sIHNjcmF0Y2hfb3BlcmFuZHMgPSAyIDogaTY0LCB3aW5kb3dfcGFyYW1zID0gW3t0cmFuc2Zvcm1faW5kaWNlcyA9IEBxa19mbiwgd2luZG93X2JvdW5kcyA9IGFycmF5PGk2NDogMTAyNCwgMTAyND59LCB7dHJhbnNmb3JtX2luZGljZXMgPSBAdl9mbiwgd2luZG93X2JvdW5kcyA9IGFycmF5PGk2NDogMTAyNCwgMTI4Pn0sIHt0cmFuc2Zvcm1faW5kaWNlcyA9IEBvaV9mbiwgd2luZG93X2JvdW5kcyA9IGFycmF5PGk2NDogMTAyNCwgMTI4Pn1dfSB7CiAgICAlYzAgPSBhcml0aC5jb25zdGFudCAwIDogaW5kZXgKICAgICVjMF9pMzIgPSBhcml0aC5jb25zdGFudCAwIDogaTMyCiAgICAlMCA9IGFyaXRoLmNtcGkgZXEsICVjMF9pMzIsICVhcmcxIDogaTMyCiAgICBzY2YuaWYgJTAgewogICAgICAlY3N0XzMgPSBhcml0aC5jb25zdGFudCBkZW5zZTwwLjAwMDAwMGUrMDA+IDogdmVjdG9yPDEwMjR4MTI4eGJmMTY+CiAgICAgIHZlY3Rvci5zdG9yZSAlY3N0XzMsICVhcmc0WyVjMCwgJWMwXSA6IG1lbXJlZjwxMDI0eDEyOHhiZjE2LCAjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPj4sIHZlY3RvcjwxMDI0eDEyOHhiZjE2PgogICAgICAlY3N0XzQgPSBhcml0aC5jb25zdGFudCBkZW5zZTwweEZGODA+IDogdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgICB2ZWN0b3Iuc3RvcmUgJWNzdF80LCAlYXJnNVslYzAsICVjMF0gOiBtZW1yZWY8MTAyNHgxeGJmMTYsICN0cHUubWVtb3J5X3NwYWNlPHZtZW0+PiwgdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgICAlY3N0XzUgPSBhcml0aC5jb25zdGFudCBkZW5zZTwwLjAwMDAwMGUrMDA+IDogdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgICB2ZWN0b3Iuc3RvcmUgJWNzdF81LCAlYXJnNlslYzAsICVjMF0gOiBtZW1yZWY8MTAyNHgxeGJmMTYsICN0cHUubWVtb3J5X3NwYWNlPHZtZW0+PiwgdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgfQogICAgJTEgPSB2ZWN0b3IubG9hZCAlYXJnMlslYzAsICVjMF0gOiBtZW1yZWY8MTAyNHgxMDI0eGJmMTYsICN0cHUubWVtb3J5X3NwYWNlPHZtZW0+PiwgdmVjdG9yPDEwMjR4MTAyNHhiZjE2PgogICAgJWNzdCA9IGFyaXRoLmNvbnN0YW50IGRlbnNlPDB4RkY4MD4gOiB2ZWN0b3I8MTAyNHhiZjE2PgogICAgJTIgPSB2ZWN0b3IubXVsdGlfcmVkdWN0aW9uIDxtYXhpbXVtZj4sICUxLCAlY3N0IFsxXSA6IHZlY3RvcjwxMDI0eDEwMjR4YmYxNj4gdG8gdmVjdG9yPDEwMjR4YmYxNj4KICAgICUzID0gdmVjdG9yLnNoYXBlX2Nhc3QgJTIgOiB2ZWN0b3I8MTAyNHhiZjE2PiB0byB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICAlNCA9IHZlY3Rvci5sb2FkICVhcmc1WyVjMCwgJWMwXSA6IG1lbXJlZjwxMDI0eDF4YmYxNiwgI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4+LCB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICAlNSA9IGFyaXRoLm1heGltdW1mICU0LCAlMyA6IHZlY3RvcjwxMDI0eDF4YmYxNj4KICAgICVjc3RfMCA9IGFyaXRoLmNvbnN0YW50IGRlbnNlPDAuMDAwMDAwZSswMD4gOiB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICAlNiA9IGFyaXRoLmNtcGYgb2VxLCAlNCwgJTUgOiB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICAlNyA9IGFyaXRoLnN1YmYgJTQsICU1IDogdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgJTggPSBhcml0aC5zZWxlY3QgJTYsICVjc3RfMCwgJTcgOiB2ZWN0b3I8MTAyNHgxeGkxPiwgdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgJTkgPSB2ZWN0b3IuYnJvYWRjYXN0ICU1IDogdmVjdG9yPDEwMjR4MXhiZjE2PiB0byB2ZWN0b3I8MTAyNHgxMDI0eGJmMTY+CiAgICAlMTAgPSBhcml0aC5zdWJmICUxLCAlOSA6IHZlY3RvcjwxMDI0eDEwMjR4YmYxNj4KICAgICUxMSA9IG1hdGguZXhwICUxMCA6IHZlY3RvcjwxMDI0eDEwMjR4YmYxNj4KICAgICVjc3RfMSA9IGFyaXRoLmNvbnN0YW50IGRlbnNlPDAuMDAwMDAwZSswMD4gOiB2ZWN0b3I8MTAyNHhiZjE2PgogICAgJTEyID0gdmVjdG9yLm11bHRpX3JlZHVjdGlvbiA8YWRkPiwgJTExLCAlY3N0XzEgWzFdIDogdmVjdG9yPDEwMjR4MTAyNHhiZjE2PiB0byB2ZWN0b3I8MTAyNHhiZjE2PgogICAgJTEzID0gdmVjdG9yLnNoYXBlX2Nhc3QgJTEyIDogdmVjdG9yPDEwMjR4YmYxNj4gdG8gdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgJTE0ID0gdmVjdG9yLmxvYWQgJWFyZzZbJWMwLCAlYzBdIDogbWVtcmVmPDEwMjR4MXhiZjE2LCAjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPj4sIHZlY3RvcjwxMDI0eDF4YmYxNj4KICAgICUxNSA9IG1hdGguZXhwICU4IDogdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgJTE2ID0gYXJpdGgubXVsZiAlMTUsICUxNCA6IHZlY3RvcjwxMDI0eDF4YmYxNj4KICAgICUxNyA9IGFyaXRoLmFkZGYgJTE2LCAlMTMgOiB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICAlMTggPSB2ZWN0b3IubG9hZCAlYXJnNFslYzAsICVjMF0gOiBtZW1yZWY8MTAyNHgxMjh4YmYxNiwgI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4+LCB2ZWN0b3I8MTAyNHgxMjh4YmYxNj4KICAgICUxOSA9IG1hdGguZXhwICU4IDogdmVjdG9yPDEwMjR4MXhiZjE2PgogICAgJTIwID0gYXJpdGgubXVsZiAlMTksICUxNCA6IHZlY3RvcjwxMDI0eDF4YmYxNj4KICAgICUyMSA9IHZlY3Rvci5icm9hZGNhc3QgJTIwIDogdmVjdG9yPDEwMjR4MXhiZjE2PiB0byB2ZWN0b3I8MTAyNHgxMjh4YmYxNj4KICAgICUyMiA9IGFyaXRoLm11bGYgJTIxLCAlMTggOiB2ZWN0b3I8MTAyNHgxMjh4YmYxNj4KICAgICUyMyA9IHZlY3Rvci5sb2FkICVhcmczWyVjMCwgJWMwXSA6IG1lbXJlZjwxMDI0eDEyOHhiZjE2LCAjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPj4sIHZlY3RvcjwxMDI0eDEyOHhiZjE2PgogICAgJTI0ID0gdmVjdG9yLnNoYXBlX2Nhc3QgJTIzIDogdmVjdG9yPDEwMjR4MTI4eGJmMTY+IHRvIHZlY3RvcjwxMDI0eDEyOHhiZjE2PgogICAgJTI1ID0gdmVjdG9yLnNoYXBlX2Nhc3QgJTExIDogdmVjdG9yPDEwMjR4MTAyNHhiZjE2PiB0byB2ZWN0b3I8MTAyNHgxMDI0eGJmMTY+CiAgICAlMjYgPSB2ZWN0b3Iuc2hhcGVfY2FzdCAlMjIgOiB2ZWN0b3I8MTAyNHgxMjh4YmYxNj4gdG8gdmVjdG9yPDEwMjR4MTI4eGJmMTY+CiAgICAlMjcgPSBhcml0aC5leHRmICUyNiA6IHZlY3RvcjwxMDI0eDEyOHhiZjE2PiB0byB2ZWN0b3I8MTAyNHgxMjh4ZjMyPgogICAgJTI4ID0gdHB1Lm1hdG11bCAlMjUsICUyNCwgJTI3IHtkaW1lbnNpb25fbnVtYmVycyA9ICN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzBdLCBbMF0sIFsxXSwgWzAsIDAsIDEsIDFdLCBbXSwgW10+LCBwcmVjaXNpb24gPSAjdHB1LmNvbnRyYWN0X3ByZWNpc2lvbjxiZjE2PiwgdHJhbnNwb3NlX2xoc19oaW50ID0gZmFsc2V9IDogdmVjdG9yPDEwMjR4MTAyNHhiZjE2PiwgdmVjdG9yPDEwMjR4MTI4eGJmMTY+LCB2ZWN0b3I8MTAyNHgxMjh4ZjMyPiAtPiB2ZWN0b3I8MTAyNHgxMjh4ZjMyPgogICAgJTI5ID0gYXJpdGgudHJ1bmNmICUyOCA6IHZlY3RvcjwxMDI0eDEyOHhmMzI+IHRvIHZlY3RvcjwxMDI0eDEyOHhiZjE2PgogICAgJTMwID0gdmVjdG9yLnNoYXBlX2Nhc3QgJTI5IDogdmVjdG9yPDEwMjR4MTI4eGJmMTY+IHRvIHZlY3RvcjwxMDI0eDEyOHhiZjE2PgogICAgJWNzdF8yID0gYXJpdGguY29uc3RhbnQgZGVuc2U8MS4wMDAwMDBlKzAwPiA6IHZlY3RvcjwxMDI0eDF4YmYxNj4KICAgICUzMSA9IGFyaXRoLmRpdmYgJWNzdF8yLCAlMTcgOiB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICAlMzIgPSB2ZWN0b3IuYnJvYWRjYXN0ICUzMSA6IHZlY3RvcjwxMDI0eDF4YmYxNj4gdG8gdmVjdG9yPDEwMjR4MTI4eGJmMTY+CiAgICAlMzMgPSBhcml0aC5tdWxmICUzMCwgJTMyIDogdmVjdG9yPDEwMjR4MTI4eGJmMTY+CiAgICAlMzQgPSB2ZWN0b3Iuc2hhcGVfY2FzdCAlMzMgOiB2ZWN0b3I8MTAyNHgxMjh4YmYxNj4gdG8gdmVjdG9yPDEwMjR4MTI4eGJmMTY+CiAgICB2ZWN0b3Iuc3RvcmUgJTM0LCAlYXJnNFslYzAsICVjMF0gOiBtZW1yZWY8MTAyNHgxMjh4YmYxNiwgI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4+LCB2ZWN0b3I8MTAyNHgxMjh4YmYxNj4KICAgIHZlY3Rvci5zdG9yZSAlNSwgJWFyZzVbJWMwLCAlYzBdIDogbWVtcmVmPDEwMjR4MXhiZjE2LCAjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPj4sIHZlY3RvcjwxMDI0eDF4YmYxNj4KICAgIHZlY3Rvci5zdG9yZSAlMTcsICVhcmc2WyVjMCwgJWMwXSA6IG1lbXJlZjwxMDI0eDF4YmYxNiwgI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4+LCB2ZWN0b3I8MTAyNHgxeGJmMTY+CiAgICByZXR1cm4KICB9CiAgZnVuYy5mdW5jIEBxa19mbiglYXJnMDogaTMyLCAlYXJnMTogaTMyKSAtPiAoaTMyLCBpMzIpIHsKICAgIHJldHVybiAlYXJnMCwgJWFyZzEgOiBpMzIsIGkzMgogIH0KICBmdW5jLmZ1bmMgQHZfZm4oJWFyZzA6IGkzMiwgJWFyZzE6IGkzMikgLT4gKGkzMiwgaTMyKSB7CiAgICAlYzBfaTMyID0gYXJpdGguY29uc3RhbnQgMCA6IGkzMgogICAgcmV0dXJuICVhcmcxLCAlYzBfaTMyIDogaTMyLCBpMzIKICB9CiAgZnVuYy5mdW5jIEBvaV9mbiglYXJnMDogaTMyLCAlYXJnMTogaTMyKSAtPiAoaTMyLCBpMzIpIHsKICAgICVjMF9pMzIgPSBhcml0aC5jb25zdGFudCAwIDogaTMyCiAgICByZXR1cm4gJWFyZzAsICVjMF9pMzIgOiBpMzIsIGkzMgogIH0KfQ==","input_memory_space_colors":[],"needs_hlo_passes":false,"needs_layout_passes":true,"output_memory_colors":[],"output_memory_space_colors":[]},"flag_configs":[],"retry_config":{"retry_count":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"3305472"}]}
}

copper lines are fusion ops · underlined lines carry the program’s big intermediate · dim lines are regions, parameters, and bookkeeping

hover any line on either side · captured on TPU v6 lite, jax 0.11.0, by labs/capture-hlo.ipynb
three programs captured by labs/capture-hlo.ipynb on v6e-1 · ops verbatim, frame tables trimmed, runtime paths scrubbed
the machine's
own account

A real profile, read op by op

Twenty iterations of naive attention, traced on a v6e-1 and aggregated from the device plane: hardware time, not a cost model. Two things to find. The envelope equals the two compute ops alone, so the 89.6 µs copy hid entirely behind compute. And the second-largest op is a custom call nobody wrote: XLA:TPU pattern-matched the softmax and dispatched its own online-softmax kernel, exactly the escape hatch chapter 09 teaches and the measurement habit chapter 10 closes on.

GYM·08 the timeline x-ray
naive attention, seq 8192, d 128, bf16 · 20 iterations · per-op hardware time from the device plane

The reading: the envelope is 217.6 µs and so is fusion + online-softmax alone, which means the 89.6 µs copy hid completely behind compute. And the second-largest op is a kernel nobody in this program wrote: the compiler dispatched its own.

hover an op: it explains itself · click for its full HLO text · captured on TPU v6 lite, jax 0.11.0
captured by labs/profile-bytes.ipynb on v6e-1, jax 0.11 · durations from the device plane, verbatim