XLA and Just-In-Time Compilation

A naive PyTorch training step launches hundreds of separate CUDA kernels per iteration. Each launch carries overhead, and between kernels the GPU stalls while intermediate tensors are written to and read back from HBM. XLA's central bet is that if you see the whole computation graph before executing any of it, you can collapse those hundreds of round-trips into a handful of fused operations that never touch slow memory between steps. That bet turned out to be correct enough to power Google's entire TPU stack.

What "compiling a graph" actually means

Every ML framework, at some level, represents a forward pass as a directed acyclic graph of operations (matmul, softmax, layer norm, ...). An eager runtime executes each node immediately as Python reaches it. A JIT compiler instead traces or captures that graph, hands it to an optimiser, and emits hardware-specific code for the whole thing.

XLA (Accelerated Linear Algebra) follows a three-stage pipeline:

1. Framework lowering. JAX, TensorFlow, or PyTorch/XLA converts Python ops into StableHLO, a portable intermediate representation (IR) with a stable opset. StableHLO is then lowered further into XLA's internal HLO (High-Level Operations) dialect.
2. Target-independent optimisation. XLA runs algebraic simplification, common-subexpression elimination (CSE), and - crucially - operation fusion. Two element-wise ops that would each read and write a full tensor become one fused op with a single memory round-trip.
3. Backend code generation. The GPU backend emits PTX/SASS via LLVM; the TPU backend produces TPU assembly. Each backend can apply further target-specific fusion decisions before final emission.

The result is a compiled artefact bound to the exact tensor shapes seen during tracing. New shapes require re-compilation, which is why shape dynamism is one of XLA's most persistent pain points.

HLO fusion: the mechanism behind the speedup

Fusion is not a vague "it batches things together" idea. It has a precise meaning: two HLO operations are fused when their combined kernel reads inputs once, computes both ops in registers, and writes outputs once.

Consider GELU applied after a linear projection:

# unfused - two kernel launches, two HBM round-trips
x = matmul(W, h)          # write x to HBM
y = gelu(x)               # read x from HBM, write y to HBM

# fused - one kernel launch, one HBM round-trip
y = fused_matmul_gelu(W, h)

For a 4096-wide activation on an H100, each HBM round-trip costs roughly 1-2 microseconds at peak bandwidth (3.35 TB/s). Multiply that across hundreds of element-wise ops per transformer layer, thousands of layers, and millions of training steps, and the savings compound.

XLA's fusion heuristic groups producers and consumers greedily, subject to register pressure limits. The compiler also performs layout assignment - choosing row-major vs. column-major for each tensor - to avoid transpose overhead when feeding into cuBLAS.

JAX's jit and the tracing contract

JAX exposes XLA compilation directly through jax.jit. The first call to a jitted function traces it: JAX runs the Python function with abstract tracers (values that carry shape and dtype but no data), records the HLO graph, and ships it to XLA for compilation.

import jax
import jax.numpy as jnp

@jax.jit
def step(x, W):
    return jnp.tanh(x @ W)

# First call: traces + compiles for shape (128, 512) @ (512, 256)
y = step(jnp.ones((128, 512)), jnp.ones((512, 256)))

# Second call: uses cached compiled artefact - no Python overhead
y = step(jnp.ones((128, 512)), jnp.ones((512, 256)))

The cache key is (function_identity, input_shapes, input_dtypes, static_arguments). Change the batch size and you pay compilation cost again. This tracing contract also means you cannot use Python control flow that branches on runtime tensor values inside jit - the tracer sees abstract shapes, not numbers.

torch.compile and TorchDynamo's different approach

PyTorch's answer, torch.compile (introduced in PyTorch 2.0), solves the same problem but with a stronger emphasis on handling arbitrary Python code. TorchDynamo inspects Python bytecode at runtime, extracts the largest possible sub-graph that is safe to compile, and falls back to eager execution for anything it cannot handle (closures over Python state, graph breaks, etc.).

The extracted graph goes to TorchInductor by default, which generates Triton (for GPU) or C++ (for CPU) code, performing its own fusion and tiling decisions. Internally, TorchInductor also has an XLA-compatible lowering path.

import torch

@torch.compile
def step(x, W):
    return torch.tanh(x @ W)

The key difference from JAX/XLA: TorchDynamo allows graph breaks, so you keep working code even when the compiler gives up on a subgraph. JAX's jit is stricter - a non-traceable operation is a hard error unless explicitly marked static. Neither approach is universally superior; the choice depends on how dynamic your model's Python logic is.

When it falls down

Shape dynamism. XLA recompiles for each unique input shape. Models with variable-length sequences (common in NLP) can generate dozens of compiled artefacts and spend significant wall time in XlaCompile. The standard mitigation is padding inputs to a small set of bucketed lengths, but this wastes FLOPs on padding tokens.

Compilation latency. A large transformer can take 30-120 seconds to compile on first run. This is acceptable for a multi-day training job; it is disqualifying for low-latency inference where cold-start matters. Some teams pre-compile and serialise the artefact, but this couples the binary to exact shapes and hardware generation.

Python-side control flow. Any if statement whose condition depends on a tensor value inside a jitted function forces either a recompile (if the value is marked static) or a hard trace error. Complex sampling loops (beam search with dynamic stopping, speculative decoding) require either restructuring as scan/while-loop primitives or falling outside jit entirely.

Debugging difficulty. Once a computation is inside a compiled region, Python debuggers and intermediate print calls stop working as expected. JAX provides jax.debug.print as an escape hatch, but the mental model of "my program is a graph, not a sequence of statements" takes time to internalise.

Operator coverage gaps. Custom CUDA ops written in C++/CUDA without XLA-compatible lowering rules are opaque to the compiler. The graph breaks at those boundaries, preventing fusion across them. Registering a proper HLO custom-call or writing the op in Triton is necessary to close those gaps.

Further reading

• XLA Architecture (openxla.org)
• JAX JIT Compilation (docs.jax.dev)
• OpenXLA Project Overview (openxla.org)