Show anyone a new approach to heterogeneous programming and you get the same question inside thirty seconds. Sometimes it is “how is this not SYCL,” sometimes “how is this not Kokkos,” and if the p...

In September, I am presenting a session at CppCon that challenges conventional frontend design: “Escaping the AST: A Data-Oriented, Lock-Free Parallel Compiler Architecture.” The core thesis requi...

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...