Pallas for people who know JAX but not kernels yet
jax.numpy ergonomics.
Pallas currently lowers to Mosaic on TPUs and Mosaic GPU on newer NVIDIA GPUs. There is also a Triton GPU backend, but it is maintained on a best-effort basis and not recommended for new use. Mosaic GPU is currently aimed at Hopper and newer GPUs.
What is a kernel?
In operating systems, “kernel” means the core program managing the machine. In numerical programming, a kernel is different, it is a small function intended to run close to the hardware, often many times in parallel.
A good mental model is:
A kernel is not “the whole computation”. A kernel is the work done by one worker on one piece of the computation.
For example, if you add two vectors of length 1024, ordinary JAX lets you write:
x + y
Pallas asks a more explicit question:
Which slice of
xandydoes this particular program instance read, and where does it write the answer?
Pallas is not just “JAX but faster”. It is JAX with memory and tiling model exposed.
Hello world in Pallas: adding two vectors
import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl
def add_vectors_kernel(x_ref, y_ref, o_ref):
x = x_ref[...]
y = y_ref[...]
o_ref[...] = x + y
This looks like JAX, but two important things are different.
First, the kernel does not receive normal jax.Arrays. It receives Refs. A Ref is a reference to a mutable buffer. You read from it using indexing syntax, and that read gives you a jax.Array. You write back into another Ref.
Second, the kernel does not return anything. The output is written into o_ref.
x = jnp.arange(10)
y = jnp.arange(10)
out = pl.pallas_call(
add_vectors_kernel,
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
)(x, y)
pallas_call turns the kernel into something callable from JAX. It needs out_shape because the kernel itself does not return a value, so JAX needs to know the shape and dtype of the output buffer ahead of time. The JAX docs describe pallas_call as lifting a Pallas kernel into a JAX operation, with out_shape determining the shape and dtype of the output Ref.
Ref vs JAX array
This is the most important beginner distinction.
x_ref # a mutable reference to memory
x_ref[...] # read from memory; gives a jax.Array
o_ref[...] = value # write a jax.Array back to memory
A normal JAX function is usually pure:
def f(x):
return x + 1
A Pallas kernel is more like:
def kernel(x_ref, o_ref):
x = x_ref[...]
o_ref[...] = x + 1
So, inside a kernel:
Refmeans “where the data lives”.jax.Arraymeans “the value I loaded and can compute with”.- assigning to an output
Refmeans “store the result”.
Now the real idea: grids
The first example uses the whole vector at once. That is useful for learning the API, but it hides the main idea of kernels.
A kernel usually handles one block. The grid says how many blocks to run.
The JAX docs give the simplest mental model: grid=(n,) is like running the kernel in a loop n times; grid=(n, m) is like nested loops; each invocation is called a “program”, and inside the kernel you can ask which program you are using pl.program_id(axis).
A beginner-friendly way to say it:
The kernel is the recipe. The grid is how many cooks you launch.
program_idtells each cook which plate they are responsible for.
Example:
BLOCK_SIZE = 4
def add_vectors_block_kernel(x_ref, y_ref, o_ref):
pid = pl.program_id(0)
start = pid * BLOCK_SIZE
offsets = start + jnp.arange(BLOCK_SIZE)
x = x_ref[offsets]
y = y_ref[offsets]
o_ref[offsets] = x + y
Then call it with a grid:
x = jnp.arange(16)
y = jnp.arange(16)
out = pl.pallas_call(
add_vectors_block_kernel,
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
grid=(x.size // BLOCK_SIZE,),
)(x, y)
This is the first moment where Pallas starts feeling like kernel programming. You are no longer saying “add these arrays”. You are saying:
Launch 4 programs. Program 0 handles elements 0–3. Program 1 handles elements 4–7. Program 2 handles elements 8–11. Program 3 handles elements 12–15.
 |
|---|
| Figure 1: Mental model of blocks and program ids |
What about BlockSpec and GridSpec?
At first, grids are enough. But soon you will want to say: “for each program, pass only the tile it needs.”
That is where BlockSpec comes in. A BlockSpec tells Pallas how to map a program id to a slice/block of an input or output. The docs summarize it as the way to provide the mapping between a grid iteration and the block of each input/output operated on by that invocation.
I like to think of it this way:
grid: how many program instances exist?program_id: which instance am I?BlockSpec: which block of each array should this instance see?GridSpec: the packaged version of grid + input specs + output specs.
Running on TPUs
For TPU use, you typically install the TPU-enabled JAX package and run on a TPU runtime:
pip install -U "jax[tpu]" -qq
Then sanity check:
import jax
jax.devices()
You should see TPU devices.
On TPU, Pallas lowers through Mosaic. The TPU programming model exposes memory spaces such as HBM, VMEM, and SMEM. The TPU pipelining docs describe HBM as device memory, VMEM as vector SRAM/cache, and SMEM as scalar SRAM/cache. They also note that Pallas TPU exposes these memory hierarchy levels to users.
Beginner translation:
On TPU, performance often comes from moving the right blocks of data from large, slower memory into smaller, faster memory, doing work there, and writing back.
You do not need to master this on day one. But you should know that this is why Pallas exists: it gives you control over these choices when normal JAX/XLA is not enough.
What is Mosaic GPU vs Mosaic for TPU?
This confused me at first.
Pallas is the frontend: the thing you write. Mosaic and Mosaic GPU are backend compiler paths. The JAX quickstart says Pallas lowers to Mosaic GPU on GPUs and Mosaic on TPUs.
So:
Pallas kernel
↓
Mosaic GPU backend → NVIDIA GPU code
Mosaic TPU backend → TPU code
Mosaic GPU is not “Triton renamed”. It is a lower-level GPU backend. The Mosaic GPU reference says it is similar in spirit to Triton’s programming model, but more low-level, giving more control at the cost of more work.
Debugging tips
1. Use interpret=True
out = pl.pallas_call(
add_vectors_kernel,
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
interpret=True,
)(x, y)
interpret=True runs the Pallas call through a JAX interpretation path instead of compiling to accelerator code. The API docs say this is useful for debugging and is the only way to run Pallas kernels on CPU.
2. Use debug=True
out = pl.pallas_call(
add_vectors_kernel,
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
debug=True,
)(x, y)
The debug flag prints intermediate forms of the kernel while Pallas processes it.
3. Print from inside a kernel
Pallas has pl.debug_print for printing values from inside kernels, although formatting rules differ by backend.
def kernel(x_ref, o_ref):
x = x_ref[...]
pl.debug_print("x = {}", x)
o_ref[...] = x + 1
4. Compare compiled mode against interpreted mode
If a TPU kernel compiles but gives surprising results, compare it with interpret=True. The TPU kernel guide recommends this and asks users to file a bug if interpreted and compiled results diverge.
The mental model I wish I had earlier
When writing Pallas, ask these questions in order:
- What is one program instance responsible for?
- How many such program instances do I need?
- What block of each input does each program read?
- What block of each output does each program write?
- Are my blocks shaped in a way the hardware likes?
- Can I debug it first with
interpret=True? - Only then: can I make it fast?
That is the difference between writing JAX and writing kernels. In JAX, you start with the whole array. In Pallas, you start with one worker, one tile, one memory movement, and one write-back.
That is the beautiful part.
