Pallas

JAX’s kernel language and the libraries built on it: BlockSpecs and a grid in Python, lowered to Mosaic on a TPU and to Triton or Mosaic GPU on a GPU.

Explainers

• Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot — One keyword argument runs a TPU kernel on a laptop, and almost every confusing Pallas afternoon comes from not knowing what that run checked. This reads the machinery at a pinned commit of jax-ml/jax. There are three interpreters behind one argument, not one: a generic scan over the grid that models the arithmetic and removes memory entirely, a TPU mode that simulates HBM, VMEM, DMAs and semaphores and carries a happens-before race detector, and a real chip, the only rung with a compiler that can refuse you. What each rung catches and what it structurally cannot, drawn as a walk. Which import names which chip, with the memory-space tables side by side and the SMEM collision that means scalar memory on one and shared memory on the other. What interpret=True literally executes: a while loop over the grid, dynamic slices for blocks, outputs and scratch pre-filled with NaN. The padding the interpreter adds before the loop starts, computed live from the real function. Why the grid promises a count and not an order, and the one parameter that randomizes it. Then the case worth committing to: an index map one block past the end, which raises nothing, returns finite numbers, and quietly duplicates half the output. The debugging surface as a picker, one question in and one switch out. And the six-cell notebook shape, including the reset cell that exists because the interpreter keeps its simulated memory after an exception on purpose.
• pallas_call, Line by Line: One Call, Three Calling Conventions — Every Pallas kernel on every backend enters JAX through one 1381-line Python file. This walks it whole, at a pinned commit of jax-ml/jax. The file declares a single primitive, then spends the rest of itself teaching JAX what that primitive means under each transform: an abstract evaluation that decides the output types and refuses a non-manual mesh, a lowering rule that swaps its own body for a scan when you ask for interpret mode, a batching rule that answers vmap by growing the grid rather than looping, and a public function whose real work is turning a pile of keyword arguments into one GridMapping. The signature exhibit runs the real slot arithmetic out of that file: dial the grid, the scalar prefetch operands, the inputs, the outputs and the scratch, and watch the three orderings the primitive, the kernel and the index map each demand.
• The Mosaic TPU Pipeline, Line by Line — A Pallas kernel on a TPU computes from VMEM, a scratchpad of on-chip memory small enough that the array you passed in does not fit. So every block has to be copied in from HBM before the core can touch it, and copied back out afterwards, and if the core waits for either of those the matrix unit sits idle for most of the schedule. One file decides when each copy is issued, when the core is allowed to wait for it, and how many copies of a block VMEM holds at once. This walks that file at a pinned commit of jax-ml/jax: the buffered ref and its four independent cursors, the semaphore per slot, the four boolean predicates that are the entire schedule, the prologue that primes the buffers stage by stage, the loop body whose line ordering is the only reason anything overlaps, and the epilogue that waits for the last send nobody else waited for. The signature exhibit runs the real predicates over a grid you dial and prints the resulting schedule, slot by slot.
• Pallas on a TPU: From a Jaxpr to the Mosaic Dialect — A Pallas kernel is a Python function, but by the time a TPU sees it the Python is long gone: what the backend receives is a jaxpr, and what it emits is an MLIR module in the Mosaic dialect wrapped inside one HLO custom call. This reads that translation at a pinned commit of jax-ml/jax. The type mapping first, where a ref becomes a memref carrying a memory-space attribute and a value becomes a vector, and the four-step chain that turns a BlockSpec with no memory space into the literal string vmem. Then the window parameters that decide whether a block is copied in for you or left in HBM for your own DMA, and the two allocators behind run_scoped, one calling alloca and one calling a semaphore allocator. Then the table: 135 primitives carry a TPU rule at this commit, computed from the real registrations and filterable by core type and by the MLIR dialect each rule builds into. Then the two doors to the MXU, one of them a single tpu.matmul and the other five primitives that name an accumulator and an MXU index. It closes on refusal: the dtypes the lowering will not take, the block shapes it rejects before lowering starts, and the one message that means no rule exists at all.
• BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces — A Pallas kernel never sees your array. It sees one block of it, chosen by a Python function you wrote and a loop the compiler ran, and this page reads the arithmetic that connects the two at a pinned commit of jax-ml/jax. The grid is a loop nest and the kernel is its body, so the kernel runs prod(grid) times. A BlockSpec is a block shape plus an index map, and the index map returns block indices rather than element indices: the start of the slice is the block size times the number your map returned, which is one match statement in core.py and nothing more. From there the page walks the five kinds of block dimension and what each one asks your map to return, the three defaults and why omitting a BlockSpec entirely does not mean the whole array but an unblocked ref in whichever memory the backend chose, the two memory-space vocabularies and where DEFAULT lands on a TPU, and the conversion that turns your closure into a traced jaxpr along with the three rules that conversion enforces. Then the honest parts: a block shape that does not divide the array runs anyway, on a full-sized block, reading padding whose values the documentation tells you to assume are garbage, discarded on output and not on input; the grid runs with the last axis fastest on the interpreter and in an order you may not assume on a chip; and a TPU wants the last two dimensions of your block to be multiples of 8 and 128. The signature exhibit ports the real block-index arithmetic, the ceiling division, the start indices and the padding rule, so you can dial an array shape against a block shape and read the slice the kernel would receive.
• Pallas on a GPU: Two Backends, and What Each One Lets You Write — One pallas_call on an NVIDIA card reaches one of two compilers, and they do not accept the same kernel body. Each backend is, quite literally, a dictionary from JAX primitive to lowering function, and a primitive with no entry raises before a single instruction is emitted. This reads both dictionaries at a pinned commit of jax-ml/jax: 101 primitives on the Triton side, 117 on the Mosaic GPU side, 59 in common, and a matmul that lowers on one and has no rule at all on the other. Then it reads a production kernel library at a pinned commit of google/tokamax to see what a team does with that: a Triton attention kernel in one file, a Mosaic GPU dispatcher that owns no kernel, and two kernel bodies forked by GPU generation because the accumulator moved out of the registers on Blackwell. The signature exhibit checks one softmax body against the real rule tables, line by line.
• Flash Attention on a TPU, Line by Line — The flash attention algorithm fits on a napkin. The Pallas kernel that runs it on a TPU is 1715 lines: four kernel bodies, three separate launches, eleven block sizes, and a four-line diagonal test consulted from ten places. This walks the whole file at a pinned commit of jax-ml/jax.
• Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop — Flash attention treats a causal mask as a value: compute the whole block of logits, then overwrite half of them with negative infinity. Splash attention treats the same mask as a schedule. Before the kernel is ever traced, a Python pass cuts the mask into blocks the size of one grid step, labels each block empty, partial or full, and turns that label grid into three small integer arrays that live in TPU scalar memory. The kernel then reads its own loop structure out of them: an empty block is not computed and its data is not fetched, a full block skips the masking entirely, and for a banded mask the grid itself is rebuilt narrower. The signature exhibit and the dial both run the real classification out of the pinned file, so the block counts on this page are the counts the compiler would produce.
• The tokamax Op and Its Autotuner, Line by Line — A fast kernel is never one kernel. It is a family of implementations and a family of tilings, and something has to choose between them on every call. Tokamax, the fused-kernel library Google builds on Pallas, answers that in one 701-line Python file. This walks it whole at a pinned commit: the frozen Op dataclass and the two hooks a backend overrides, the twelve lines that turn a call into a dispatch key made only of argument shapes and dtypes, the five-step ladder that decides whether you get an explicit config, a cached measurement, an unmeasured guess or an error, the measurement loop that fills the cache in, and the payload the file writes into the compiled program so a captured production trace can be tuned on a machine that never runs your model. The signature exhibit is that real key, opened against the library’s own shipped cache: pick a ragged-dot workload and a chip, and read every cached implementation with its winning tiling and its measured median.
• Attention for sm90, Line by Line — A production flash attention kernel for the H100, read whole at a pinned commit of google/tokamax. The algorithm is the one everybody knows; what a real kernel adds on top of it is six hundred lines of scheduling. One thread block launches three warpgroups and each branches immediately into a different role: two do the arithmetic on 232 registers a thread, and the third drops to 40 and does nothing but issue transfers, running a stage ahead of the maths. Automatic synchronisation is switched off at launch, so every handshake in the file is one of sixteen numbered hardware barriers placed by hand, and running out of them is a real failure mode. The page walks the constants, the shared-memory plan, both configuration heuristics, both warpgroup roles and the launch, then steps through one query tile: which slot each warpgroup is holding, how far ahead the loader has run, and where the causal mask starts biting. A second exhibit runs the file own shared-memory estimate so you can watch a tiling stop fitting in 227 KiB.