Pallas: custom kernels in JAX

Most of the time, you don’t write kernels: you write jax.numpy programs, and the compiler decides how to turn them into device code. Usually it does well. But sometimes you know something the compiler doesn’t: a fusion it won’t find, a memory-access pattern it won’t discover (as in FlashAttention), or a sparsity structure it can’t exploit. Then the way to get the remaining performance is to write the kernel yourself.

Pallas is JAX’s kernel language: an extension of JAX that lets you write custom kernels for GPU and TPU, with fine-grained control over the generated code while keeping JAX tracing and the jax.numpy API. Kernels are written as functions over Refs in fast on-chip memory, launched over a grid with pl.kernel, and they compose with the rest of JAX: you can jit, vmap, and differentiate around them.

Pallas has its own extensive documentation site, which we won’t duplicate here:

pallas.jax.dev

Good entry points:

• Pallas quickstart — kernels, Refs, grids, and BlockSpecs, on both GPU and TPU.

• Pipelining and grids and BlockSpecs — the core concepts for expressing how data is carved up and streamed through on-chip memory.

• The TPU backend guide and the Mosaic GPU backend guide — per-platform details, lowering paths, and platform-specific features.

• The jax.experimental.pallas module API reference.

Pallas is experimental and changes frequently; see the Pallas changelog for recent developments.