Pallas, From the Inside
Pallas is the kernel language you reach for when XLA fusion is not enough: a Python function over blocks of a grid, lowered to Mosaic on a TPU and to Triton or Mosaic GPU on a GPU. The track starts with a kernel you run in interpret mode on a laptop, reads the language and both lowerings, derives attention, and ends in a production library that autotunes it.
Phase 1 — Run it first
Interpret mode runs a Pallas kernel on any machine; the first article is how to write one that also runs on a chip.
- LAB·1.1 · First kernels ↗ — Add, scale, and reduce as Pallas kernels in interpret mode, then read the jaxpr each one produced.
- Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot — Three interpreters sit behind one argument, and each models a different amount of the machine: read this before you trust a kernel that only ever ran on a laptop.
- Pallas, the layer ↗ — The kernels-site layer that introduces the language and its place in the descent.
Phase 2 — The language
pallas_call, BlockSpec, the grid, and the memory spaces: the four ideas every kernel is made of.
- pallas_call, Line by Line: One Call, Three Calling Conventions — Read the entry point whole before anything else: one primitive, and the transform rules that teach JAX what it means.
- BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces — The two decisions every kernel starts with, read from the source: index maps return block indices, and the block size turns one into a slice.
- LAB·1.2 · Tiled matmul ↗ — A matmul as a grid of tiles, then the same kernel with a different block shape, and what the roofline says about each.
- jax/_src/pallas/core.py ↗ — BlockSpec, GridSpec, memory spaces, and the abstract refs a kernel body sees.
- jax/_src/pallas/primitives.py ↗ — The primitives a kernel body may use: loads, stores, atomics, program ids.
Phase 3 — The TPU lowering
What a kernel becomes on a TPU: a pipelined schedule of DMAs into VMEM and a Mosaic module for the core.
- The Mosaic TPU Pipeline, Line by Line — pipeline.py line by line: the buffered ref, the four predicates that are the whole schedule, and a computed timeline of the DMA and compute slots.
- Pallas on a TPU: From a Jaxpr to the Mosaic Dialect — The rule table that turns a kernel jaxpr into Mosaic ops, read at the pin: 135 primitives, the memory-space chain, and the two doors to the MXU.
- LAB·1.4 · Pipelining profile ↗ — Profile a pipelined kernel on a TPU and read the overlap of DMA and compute in the trace.
- jax/_src/pallas/mosaic/core.py ↗ — TPU-specific specs: memory spaces, semaphores, the TPU grid.
- jax/_src/pallas/mosaic/pallas_call_registration.py ↗ — Where the TPU lowering is registered against pallas_call.
- jax/_src/pallas/mosaic/lowering.py ↗ — The lowering itself, thousands of lines; the pages above tell you which rules to open.
- Mosaic, the layer ↗ — The kernels-site layer on the TPU dialect and its tiling vocabulary.
- The TPU, the layer ↗ — The chip the lowering targets: cores, VMEM, the MXU, and the fabric.
Phase 4 — The GPU lowerings
Two lowerings for one language: Triton, and Mosaic GPU with its sm90 and sm100 paths.
- Pallas on a GPU: Two Backends, and What Each One Lets You Write — Each backend is a dictionary from primitive to lowering rule, and a missing key is an error, not a slow path: read both dictionaries before you write a GPU kernel.
- LAB·2.1 · The lowering ladder ↗ — One kernel through every IR on the way down, read at each rung.
- LAB·2.2 · Finding the spill ↗ — A kernel that spills registers, found from the compiler output rather than guessed.
- jax/_src/pallas/triton/lowering.py ↗ — The Triton lowering: jaxpr rules to Triton IR.
- jax/_src/pallas/mosaic_gpu/pipeline.py ↗ — The Mosaic GPU pipeline: warpgroups, barriers, the software pipeline.
Phase 5 — Attention, derived
The kernel everyone writes, as it is actually written: flash attention on a TPU and the splash kernel that makes masks cheap.
- Anatomy of a FlashAttention Kernel — Why attention is memory-bound and how tiling with online softmax removes the N by N matrix.
- Flash Attention on a TPU, Line by Line — The derivation as real code: four kernel bodies, three launches, and a four-line diagonal test that decides which blocks ever run.
- Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop — The mask stops being a value and becomes the loop: a compile-time pass labels every block empty, partial or full, and the kernel reads its own schedule out of scalar memory.
- LAB·3.2 · Flash attention, blind ↗ — Write flash attention from the derivation without looking at the reference kernel, then diff.
- LAB·3.3 · The backward pass ↗ — The backward kernel, and why it recomputes what the forward pass threw away.
- jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py ↗ — The splash kernel itself, with the mask machinery beside it.
Phase 6 — A kernel library
tokamax: the same kernels wrapped in an Op with a dispatch key, an autotuner, and a benchmark loop.
- The tokamax Op and Its Autotuner, Line by Line — op.py line by line: the Op dataclass, the dispatch key of argument shapes, the config ladder, and the cache the tuner fills.
- Attention for sm90, Line by Line — The whole kernel at a pinned commit: three warpgroups, sixteen hand-placed barriers, and a shared-memory budget you can dial until the tiling stops fitting.
- LAB·4.2 · Ring attention ↗ — Attention across devices as a ring, with the collective as part of the kernel.
- tokamax/_src/ops/attention/base.py ↗ — The attention Op base every backend kernel plugs into.
- tokamax/_src/benchmarking.py ↗ — The measurement loop the autotuner and the benchmarks share.