From Python to the MXU

Everything in Chapters 1–8 assumed the compiler would turn a sharded program into well-scheduled MXU work, DMAs and collectives. This chapter opens that box. It follows a small sharded MLP through every stage, using real compiler output, then covers the two front ends that feed the same compiler: JAX and TorchTPU.

The stages

The Scaling Book's summary of the TPU stack [SB9]:

1. Python and JAX trace a function into a jaxpr, JAX's own intermediate form.
2. jax.jit lowers the jaxpr to StableHLO, a versioned, platform-independent MLIR dialect. Sharding annotations travel with it.
3. XLA converts StableHLO to HLO and runs its passes:
   • sharding propagation and SPMD partitioning;
   • fusion;
   • layout assignment;
   • memory-space assignment between HBM and VMEM;
   • scheduling. The optimized HLO is what a profile shows.
4. The TPU backend lowers HLO to LLO, a low-level IR that programs the TensorCore directly: copies between memories, pushes into the systolic arrays, DMA scheduling.
5. LLO becomes machine code, loaded into the TensorCore's instruction memory.

Pallas kernels enter as custom calls at stage 3 and are compiled by Mosaic (Chapter 10). The TPU backend of XLA, stages 4–5, ships in the closed-source libtpu [X-arch].

Each stage can be printed: jax.make_jaxpr(f)(*args), jax.jit(f).lower(*args).as_text() for StableHLO, and .compile().as_text() for optimized HLO [J-aot]. XLA_FLAGS=--xla_dump_to=DIR writes every HLO pass to disk [X-dumps].

One MLP, every stage

The listings below come from running scripts/ch09_dump.py with JAX 0.8.2 on CPU with 4 simulated devices, on a 2×2 mesh named ('data', 'model'). The function is gelu(x @ w1) @ w2 with x sharded P('data', None), w1 sharded P(None, 'model') and w2 sharded P('model', None): Megatron tensor parallelism over model, data parallelism over data. TPU output differs in layout tiling, fusion choices and memory-space annotations. The structure is the same.

{ lambda ; a:bf16[8,512] b:bf16[512,2048] c:bf16[2048,512]. let
    d:bf16[8,2048] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] a b
    e:bf16[8,2048] = integer_pow[y=3] d
    f:bf16[8,2048] = mul 0.044677734375:bf16[] e
    g:bf16[8,2048] = add d f
    h:bf16[8,2048] = mul 0.796875:bf16[] g
    i:bf16[8,2048] = tanh h
    j:bf16[8,2048] = add 1.0:bf16[] i
    k:bf16[8,2048] = mul 0.5:bf16[] j
    l:bf16[8,2048] = mul d k
    m:bf16[8,512] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] l c
  in (m,) }

  sdy.mesh @mesh = <["data"=2, "model"=2]>
  func.func public @main(%arg0: tensor<8x512xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {}]>}, %arg1: tensor<512x2048xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"model"}]>}, %arg2: tensor<2048x512xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{"model"}, {}]>}) -> (tensor<8x512xbf16> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{"data"}, {}]>}) {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<8x512xbf16>, tensor<512x2048xbf16>) -> tensor<8x2048xbf16>
    %1 = stablehlo.multiply %0, %0 : tensor<8x2048xbf16>
    %2 = stablehlo.multiply %1, %0 : tensor<8x2048xbf16>
    %4 = stablehlo.multiply %3, %2 : tensor<8x2048xbf16>
    %5 = stablehlo.add %0, %4 : tensor<8x2048xbf16>
    %7 = stablehlo.multiply %6, %5 : tensor<8x2048xbf16>
    %8 = stablehlo.tanh %7 : tensor<8x2048xbf16>
    %10 = stablehlo.add %9, %8 : tensor<8x2048xbf16>
    %12 = stablehlo.multiply %11, %10 : tensor<8x2048xbf16>
    %13 = stablehlo.multiply %0, %12 : tensor<8x2048xbf16>
    %14 = stablehlo.dot_general %13, %arg2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<8x2048xbf16>, tensor<2048x512xbf16>) -> tensor<8x512xbf16>
    return %14 : tensor<8x512xbf16>

ENTRY %main.0_spmd (param: bf16[4,512], param.1: bf16[512,1024], param.2: bf16[1024,512]) -> bf16[4,512] {
  %param = bf16[4,512]{1,0} parameter(0), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, metadata={op_name="x"}
  %param.1 = bf16[512,1024]{1,0} parameter(1), sharding={devices=[1,2,2]<=[2,2]T(1,0) last_tile_dim_replicate}, metadata={op_name="w1"}
  %param.2 = bf16[1024,512]{1,0} parameter(2), sharding={devices=[2,1,2]<=[2,2]T(1,0) last_tile_dim_replicate}, metadata={op_name="w2"}
  %dot = f32[4,1024]{1,0} dot(%wrapped_convert, %wrapped_convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=2}
  %dot.1 = f32[4,512]{1,0} dot(%convert_convert_fusion, %wrapped_convert.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=3}
  %all-reduce.1 = f32[4,512]{1,0} all-reduce(%dot.1), channel_id=1, replica_groups=[2,2]<=[4], use_global_device_ids=true, to_apply=%add.clone_promoted, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=3}
  ROOT %wrapped_convert.3 = bf16[4,512]{1,0} fusion(%all-reduce.1), kind=kLoop, calls=%wrapped_convert_computation.3, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=3}

Read the third listing as Chapter 5 would:

• The first matmul is case 1: x is sharded on its batch rows, w1 on its output columns, and there is no communication.
• The second is case 3: both operands are sharded on the contracting dimension over model, so each device holds a partial sum and XLA adds an all-reduce.
• replica_groups=[2,2]<=[4] describes which devices reduce together: the two model peers in each data row.
• The convert ops to f32 are a CPU artifact; TPU MXUs take bf16 and accumulate in f32 directly.

What you control, and what you do not

You control:

• Shardings of inputs and outputs, and intermediate constraints via jax.lax.with_sharding_constraint. The Scaling Book calls adding constraints to steer propagation "compiler tickling", and estimates it is "about 60% of JAX parallel programming in the automatic partitioning world". The usual symptom of bad propagation is an unexpected AllGather taking most of a profile [SB10].

• Layouts. Passing AUTO input layouts to jax.jit lets XLA pick, avoiding relayout copies [SB9].

• Rematerialization policies via jax.checkpoint.

• Buffer donation via donate_argnums, which lets outputs reuse input buffers. This is how KV caches update without copies.

• XLA flags, passed on TPU through the LIBTPU_INIT_ARGS environment variable, the convention MaxText and Google's recipes use unverified. Examples:

  • the latency-hiding scheduler (xla_tpu_enable_latency_hiding_scheduler);
  • async collective fusion;
  • the windowed-einsum threshold for collective matmul;
  • scoped VMEM limits (xla_tpu_scoped_vmem_limit_kib);
  • memory-space-assignment tuning.

  The OpenXLA flags guide lists them [X-flags].

• Pallas kernels, for anything the compiler does not generate well (Chapter 10).

You do not control: TPU code generation, fusion decisions beyond what flags expose, VLIW bundle scheduling and VMEM allocation inside fusions. The backend is closed, so these are visible only in their effects, through profiles and HLO dumps.

Sharding: GSPMD, Shardy and three modes

GSPMD introduced the model JAX still uses [P-GSPMD, §3–4]. You write a single-device program and annotate a few tensors with shardings. The compiler propagates shardings to every other tensor, then rewrites the program into one per-device program with collectives inserted. The paper reported 50–62% compute utilization up to 2,048 TPUv3 cores.

Shardy replaced GSPMD's propagation as JAX's default in version 0.7.0 [J-log]. It is MLIR-based, uses named mesh axes instead of device-assignment tensors, and distinguishes open dimensions (which propagation may shard further) from closed ones. It also adds propagation priorities. Its annotations are the sdy.sharding attributes in the StableHLO listing above [X-shardy-jax].

JAX now offers three ways to say how a program is sharded [SB10] [J-shard]:

Explicit mode moves sharding errors from profile time to trace time. Multiplying arrays whose contracting dimensions are both sharded is ambiguous: the result could be AllReduced or ReduceScattered. JAX refuses to guess:

ShardingTypeError: Contracting dimensions are sharded and it is ambiguous how the output should be sharded. Please specify the output sharding via the `out_sharding` parameter. Got lhs_contracting_spec=('Y',) and rhs_contracting_spec=('Y',)

Manual mode is the escape hatch. Inside jax.shard_map, each device sees its own block and every collective is explicit. The compiled HLO then contains exactly the collectives you wrote:

ENTRY %main.0_spmd (param.2: bf16[4,256], param.3: bf16[256,512]) -> bf16[4,512] {
  %param.2 = bf16[4,256]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="a_blk"}
  %param.3 = bf16[256,512]{1,0} parameter(1), sharding={devices=[2,1,2]<=[2,2]T(1,0) last_tile_dim_replicate}, metadata={op_name="b_blk"}
  %dot_general.3 = f32[4,512]{1,0} dot(%wrapped_convert, %wrapped_convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(matmul_psum)/shard_map/dot_general" stack_frame_id=2}
  %psum_invariant.8 = f32[4,512]{1,0} all-reduce(%dot_general.3), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, to_apply=%region_0.0_promoted, metadata={op_name="jit(matmul_psum)/shard_map/psum_invariant" stack_frame_id=3}
  ROOT %wrapped_convert.2 = bf16[4,512]{1,0} fusion(%psum_invariant.8), kind=kLoop, calls=%wrapped_convert_computation.2, metadata={op_name="jit(matmul_psum)/shard_map/psum_invariant" stack_frame_id=3}

jax.shard_map became the public API in JAX 0.8.0, replacing jax.experimental.shard_map [J-log]. The Scaling Book's collective matmul (Chapter 5) is written in it.

TorchTPU

TorchTPU is Google's native PyTorch backend for TPUs, built with Meta's PyTorch team. It was described in April 2026 and open-sourced on 2026-10-20 [T-blog1] [T-oss]. Google's docs state that "TorchTPU replaces PyTorch/XLA" [D-pytorch], and PyTorch/XLA's repository is archived on 2026-11-15 [T-xla]. For a practical guide to moving existing code, see PyTorch/XLA is deprecated: what to use instead.

Where it plugs in

TorchTPU registers as an out-of-tree PyTorch backend on the PrivateUse1 dispatch key, renamed tpu. Importing it enables torch.device("tpu"). Tensors are ordinary PyTorch tensors, with "no subclasses, no wrappers" [T-how]. That is the main difference from torchax, whose tensors are a subclass wrapping JAX arrays, and from PyTorch/XLA, whose xla device held lazy tensors.

How an op executes

When an ATen op dispatches to tpu, TorchTPU does not run it. It creates a deferred op holding a StableHLO builder; aten::mul becomes a stablehlo.multiply under construction. The returned tensor's storage points at it. PyTorch semantics that StableHLO lacks are emulated: views and in-place mutation are handled by decomposing and inverting views, since StableHLO is single-assignment [T-how].

When deferred ops are compiled and run depends on the mode [T-exec]:

Two details matter for performance reasoning:

• No Inductor. TorchTPU states that it "does not use TorchInductor". The FX graph goes to XLA, so the fusion and scheduling of this chapter apply unchanged [T-exec].
• Fused eager's speedup is reported inconsistently. The launch blog says 50% to more than 100% over strict eager; the docs say up to 50% [T-blog1] [T-exec].

Execution goes through PjRT and libtpu on a background thread, so errors can surface after the op that caused them. Compiled programs are cached at three levels: in the process, in shared memory on the machine, and in a remote persistent store. Keys cover the graph, shapes, compiler options and library versions. A first compile can take minutes; the caches make it a one-time cost [T-exec].

Distributed: the SPMD handshake

TorchTPU provides a tpu_dist process group. DDP, FSDP2, DTensor and SimpleFSDP run on it, and torch.tpu.distributed.topology_aware_mesh() builds a DeviceMesh aligned with ICI rings. Collectives lower to StableHLO collective ops compiled by XLA, "no separate library like NCCL" [T-dist].

This creates a problem JAX never had. XLA's collectives assume every device runs the same program (SPMD). PyTorch's distributed model is one process per device, and ordinary PyTorch code often branches on rank: rank 0 logs, or only some ranks do an extra step. That is MPMD. If two ranks compile different programs around a collective, the collective can deadlock or mismatch. TorchTPU's default is to compile collectives in isolation, which is safe but prevents fusing them with neighbouring compute. The SPMD handshake lets ranks confirm they are running identical programs before fusing across collectives. It is controlled by TORCH_TPU_SPMD_HANDSHAKE_STAGE (OFF, COMPILE_STAGE or DISPATCH_STAGE), and the @spmd_safe decorator marks code known to be rank-independent [T-dist]. JAX avoids the problem by construction: one traced program describes all devices.

Kernels and status

Pallas kernels are registered with torch.tpu.pallas.jax_op and called as ordinary ops. The April blog named the API custom_jax_kernel, so check the version you install. Tokamax kernels are usable directly, and Helion support is experimental [T-kern]. Google's PyTorch blog reports two results [T-pt]:

• Qwen3-MoE-235B trained on 8,192 TPU v7 chips at 27.1% MFU.
• A 30B model prefilling at 59% MFU on one chip.

A "52% MFU on 8,192 chips" figure circulating in press coverage does not appear in the primary sources unverified.

Four ways PyTorch code reaches a TPU

Sources: [T-xla], [R-torchax], [T-how], [T-exec].