In September I am giving a talk at CppCon called “Escaping the AST: A Data-Oriented, Lock-Free Parallel Compiler Architecture”1. The thesis is in the title. Throw out the pointer-chasing AST, put t...

I have eight devices. I want a matmul whose contracting dimension is split across some of them. So I write the matmul normally, say where the two operands live, and compile: mesh = Mesh(np.array(j...

This is the last post on the JAX/XLA stack. We went from Python to a jaxpr, from a jaxpr to optimized single-device HLO, and from single-device HLO to a mesh of devices. One door left, and it is th...

Coming from PyTorch, the first thing JAX does is offend you. You write a simple function with an if statement, wrap it in jax.jit, and call it. Instead of running, it throws a TracerBoolConversion...

The TPU has less public hardware documentation than a GPU. There is no vendor ISA manual, no die-level memory-model spec; you get peak FLOPs, HBM capacity, and a block diagram. But if you write JAX...

Here is a small function and the entire program XLA compiled it into: def simp(x): y = (x + 0.0) * 1.0 return jnp.transpose(jnp.transpose(y)) ENTRY %main (x: f32[3,3]) -> f32[3,3] { ...

If you train or serve models, XLA is probably compiling them, whether or not you have ever named it. It is the compiler behind JAX, behind TensorFlow, and — through PyTorch/XLA — behind much of PyT...

The first time I put torch.compile in front of an LLM inference server, it spent the better part of a minute compiling and handed back a 6% speedup. I nearly wrote the feature off. The next morning...

Triton is a compiler with a Python frontend. The @triton.jit decorator does not decorate a function. It parses the function’s AST, runs it through an MLIR pipeline, and emits a GPU binary. The Pyth...

When you compile a CUDA kernel, the final artifact is a cubin — a CUDA binary. It is a standard ELF64 file with NVIDIA-specific sections that encode everything the CUDA driver needs to load and lau...