XLA (Accelerated Linear Algebra) is a domain-specific compiler for machine learning that optimizes computation graphs by fusing multiple operations into single, highly efficient kernel executions. Originally developed by Google for TensorFlow and TPUs, XLA is now an open-source project under the OpenXLA umbrella, supporting PyTorch, JAX, and TensorFlow across GPUs, CPUs, and custom ML accelerators.
The fundamental insight behind XLA is that ML frameworks like PyTorch and TensorFlow execute operations one at a time, launching a separate GPU kernel for each addition, multiplication, or activation function. Each kernel launch has overhead — reading input tensors from memory, computing, and writing results back. XLA eliminates this waste by analyzing the entire computation graph, fusing compatible operations into single kernels, and generating machine code optimized for the specific target hardware.
For AI engineers, XLA represents the shift from “eager execution” (run each op immediately) to “compiled execution” (analyze everything first, then run optimized code). This is the same approach that makes JAX fast by default and enables Google’s TPU pods to train frontier models efficiently. Understanding XLA helps you reason about why some model architectures are faster than others and how to structure code for maximum hardware utilization.
XLA operates on a high-level intermediate representation called HLO (High Level Optimizer IR). When a framework like JAX or TensorFlow traces a computation, it produces an HLO graph — a platform-independent description of the mathematical operations and their data dependencies. Each node in the HLO graph is a primitive operation like dot (matrix multiply), add, reduce, or broadcast.
HLO is XLA’s internal graph representation. It captures the mathematical structure of a computation without framework-specific details. StableHLO is a portable version of this IR designed as a stable interface between ML frameworks and ML compilers, supporting dynamism, quantization, and sparsity.
The HLO graph goes through a series of optimization passes before being lowered to machine code. This pipeline is where XLA’s performance gains originate:
Fusion is XLA’s single most important optimization. Consider a simple computation: y = relu(x @ W + b). Without fusion, this executes as three separate kernel launches:
# Without XLA: 3 kernel launches, 3 memory round-trips
temp1 = matmul(x, W) # Kernel 1: read x,W from memory → compute → write temp1
temp2 = add(temp1, b) # Kernel 2: read temp1,b from memory → compute → write temp2
y = relu(temp2) # Kernel 3: read temp2 from memory → compute → write yEach kernel reads from and writes to global GPU memory (HBM), which is the primary bottleneck. With fusion, XLA combines the addition and ReLU into a single kernel that keeps intermediate values in registers:
# With XLA fusion: matmul kernel + 1 fused kernel
temp1 = matmul(x, W) # Kernel 1: matrix multiply (not fused — uses cuBLAS)
y = fused_add_relu(temp1, b) # Kernel 2: add + relu in one pass, no intermediate writesThis eliminates one full round-trip to memory. In practice, XLA fuses chains of element-wise operations, reductions, broadcasts, and reshapes into single kernels. For transformer attention blocks, this can mean fusing LayerNorm, projection, and activation into a single kernel — reducing dozens of memory round-trips to a handful.
On modern GPUs, compute throughput (TFLOPS) far exceeds memory bandwidth (TB/s). An A100 GPU delivers 312 TFLOPS of FP16 compute but only 2 TB/s of HBM bandwidth. This means most operations are “memory-bound” — they spend more time moving data than computing. Fusion keeps intermediate results in registers or L1/L2 cache (orders of magnitude faster than HBM), which is why it produces such dramatic speedups.
In JAX, XLA compilation is the default — every function decorated with @jax.jit is traced and compiled via XLA:
import jax
import jax.numpy as jnp
@jax.jit
def forward(x, W, b):
return jax.nn.relu(jnp.dot(x, W) + b)
# First call: traces and compiles via XLA (slow)
# Subsequent calls: runs cached compiled code (fast)
y = forward(x, W, b)In PyTorch, XLA is used via the torch_xla package for TPU support, or through torch.compile which can use XLA as a backend:
import torch
import torch_xla.core.xla_model as xm
# Move model to XLA device (TPU)
device = xm.xla_device()
model = model.to(device)
input = input.to(device)
# Operations are traced into an XLA graph
output = model(input)
# Explicitly trigger compilation and execution
xm.mark_step()In TensorFlow, XLA can be enabled globally or per-function:
import tensorflow as tf
@tf.function(jit_compile=True)
def forward(x, W, b):
return tf.nn.relu(tf.matmul(x, W) + b)XLA compilation has a cost: the first execution of a new computation shape triggers a compilation that can take seconds to minutes. This is amortized over subsequent calls with the same shapes. However, if input shapes change frequently (common in NLP with variable sequence lengths), recompilation can become a bottleneck. Strategies to mitigate this include padding inputs to a set of fixed “bucket” sizes and using XLA’s persistent compilation cache.
XLA is the only compiler path for Google TPUs. All major models trained on TPU pods — including Google’s PaLM, Gemini, and open-source models trained on TPU Research Cloud — are compiled through XLA. The compiler’s understanding of TPU’s matrix units (MXUs) and high-bandwidth memory hierarchy is what makes TPU training competitive with NVIDIA GPUs despite very different hardware architectures.
JAX’s entire execution model is built on XLA. Libraries like Flax, Haiku, and Pax produce JAX computations that XLA compiles and optimizes. This is why JAX code can achieve near-hardware-peak performance with relatively simple Python code — XLA handles the low-level optimization that would otherwise require hand-tuned CUDA kernels.
XLA integrates with GSPMD (General and Scalable Parallelization for ML Computation Graphs), Google’s automatic parallelization system. Engineers annotate a small number of tensors with sharding specifications, and XLA+GSPMD automatically distributes the computation across thousands of devices, inserting the necessary communication primitives (all-reduce, all-gather). This powers training runs across entire TPU pods.
For serving, XLA’s ahead-of-time (AOT) compilation produces optimized binaries tuned for specific hardware. Combined with quantization and the compiler’s ability to fuse operations, this delivers inference throughput close to hand-optimized implementations. Google’s production serving infrastructure for Gemini and other models relies on XLA-compiled TPU binaries.
The OpenXLA project — backed by Google, Apple, AMD, Intel, NVIDIA, Meta, and others — aims to make XLA the universal ML compiler. Through StableHLO, models compiled once can target GPUs, TPUs, CPUs, and custom accelerators. This reduces the need for hardware-specific optimization work and gives ML engineers a single compilation path regardless of deployment target.
torch_xla) and TensorFlow (via jit_compile=True).