XLA
Also known as Accelerated Linear Algebra, StableHLO, XLA compiler
In plain English
XLA is the compiler that turns a model graph into TPU machine code, deciding how to tile, pad, fuse, and schedule every operation.
Technical definition
XLA is the domain-specific compiler for linear algebra that consumes a StableHLO graph and produces executable code for TPUs and other accelerators, handling fusion, layout, tiling, and scheduling.
Engineering details
Every TPU serving path ends in XLA. In the JAX path, jax.jit captures a step and hands it to XLA. In TorchTPU, TorchDynamo and AOTAutograd produce an FX graph that is lowered to StableHLO, the standardized intermediate representation, and XLA compiles it. XLA is the compiler in both cases; Inductor and Triton are not involved. The compiler pads any dimension smaller than the MXU tile, picks default layouts, and fuses elementwise work, which is efficient for regular shapes and costly for shapes that fight the tile geometry. Pallas exists for the cases where XLA’s defaults are not good enough.
Why it matters
Co-design between the compiler and the hardware is a large part of Google’s cost-per-token advantage: the chip, fabric, and compiler are designed together so computation and communication can be optimized as one system. The same tight coupling means model shapes and compiled shape buckets have outsized performance effects.
How to read it in InferenceX
Several Official Preview optimizations are about steering XLA rather than replacing it, including packing expert ID and token index into one sort key so XLA performs a simpler sort, which cut sort latency from 106.6 to 21.7 microseconds and enabled an FP8 all-gather.