The last post argued that prefill and decode want different computers, and that the industry has started buying them separately. This one is about the handoff, which sounds like the easy part but i...

Three announcements inside a year, all making the same architectural bet. AWS is pairing Trainium with Cerebras: Trainium runs prefill, a Cerebras CS-3 runs decode, and the result ships as a premi...

As of this month, Claude watermarks its text output. Anthropic’s documentation is direct about the mechanism: for text, “it weaves an imperceptible watermark directly into the text itself. You won’...

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