When XLA Isn't Enough: Pallas, Mosaic, and Triton

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 the one you take when the compiler’s automatic decisions are not good enough: writing the kernel yourself. In JAX that door is Pallas, and it lowers through two very different backends — Triton on GPU a...