Which backend does pallas_call use on a GPU?

At this commit, Mosaic GPU by default. The shared GPU lowering picks Mosaic GPU when you pass a Mosaic GPU CompilerParams, or when you pass no compiler params at all and the jax_pallas_use_mosaic_gpu config flag is true, and that flag defaults to true. It picks Triton when you pass a Triton CompilerParams, or when compiler params are absent and the flag is false. Both paths warn at this commit: using pl.pallas_call for a Mosaic GPU kernel is deprecated in favour of plgpu.kernel, and the Triton backend is deprecated altogether. There is one hard override in the other direction: Mosaic GPU raises on AMD ROCm devices and the error tells you to pass Triton compiler params instead, so on an AMD card the deprecated backend is the only backend.

Why does jnp.dot fail in my Mosaic GPU Pallas kernel?

Because there is no lowering rule for it. Both Pallas GPU backends lower a kernel by walking the jaxpr and looking each primitive up in a dictionary, and a missing key raises NotImplementedError naming the primitive. The Triton table has an entry for lax.dot_general_p, roughly two hundred lines of it. The Mosaic GPU lowering has no occurrence of dot_general anywhere, under any of its four semantics tables. The matmul is not missing from the backend, it is spelled differently: you call plgpu.wgmma against an accumulator reference on Hopper, or plgpu.tcgen05_mma against a tensor-memory accumulator on Blackwell, and then wait on the asynchronous instruction before reading the accumulator. Two smaller ones surprise people for the same reason: jnp.sqrt has no Mosaic GPU rule although lax.rsqrt does, and jnp.cumsum, jnp.argmax and lax.transpose are all Triton-only.

What is warp specialisation in Pallas, and why does Triton not have it?

It is giving different warpgroups in the same thread block different jobs. plgpu.kernel takes a num_threads argument whose docstring says the number does not correspond to CUDA threads but to warpgroups, so num_threads=3 launches three warpgroups, each with its own index, and the kernel body branches on that index. A production attention kernel for Hopper uses two compute warpgroups and one memory warpgroup that does nothing but issue asynchronous copies, and it raises the compute warpgroups to 232 registers while dropping the memory warpgroup to 40, because the register file is a fixed pool shared across the block. Triton has no equivalent. Its num_warps parameter defaults to 4, meaning 128 threads that all run the same code over different lanes of the same tile. A warp there is a slice of the data, not a role.

Why do kernel libraries ship separate sm90 and sm100 kernels?

Because the accumulator lives somewhere else. Hopper issues its matmul with wgmma and accumulates into registers, which is why the compute warpgroups in a Hopper attention kernel need a large register budget. Blackwell issues tcgen05_mma and accumulates into tensor memory, a separate on-chip memory, which frees the register file for the softmax warpgroup and changes what every warpgroup around it can afford. Layered on that, Blackwell adds cluster-level collective MMA where two thread blocks cooperate on one instruction, a four-warp mesh with one warp pinned per tensor stream, and a refusal to run unstable softmax at all. The two configs end up as different classes with different field names rather than different values of one class, which is why autotuning caches are kept in separate per-device directories, and why the dispatcher branches on the type of the config rather than asking the device.

Does a kernel that works in interpret mode work on the GPU?

Not necessarily, and the reason is exactly this page. Interpret mode evaluates the kernel jaxpr on the host, so it never consults either backend rule table. A body full of jnp.dot and jnp.cumsum will produce correct numbers in interpret mode and then raise during lowering on Mosaic GPU. The compiler is the one component interpret mode does not model, and the backend rule table is the compiler. The practical habit is to read which of two failure classes an error belongs to: a message naming a primitive means the body is wrong for the backend, and a NotImplementedError raised before any primitive is named means the call is wrong for the backend, such as a dynamic grid bound or a scalar prefetch operand.

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.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Pallas on a GPU: Two Backends, and What Each One Lets You Write — built to be understood, not skimmed.

Questions

Which backend does pallas_call use on a GPU?: At this commit, Mosaic GPU by default. The shared GPU lowering picks Mosaic GPU when you pass a Mosaic GPU CompilerParams, or when you pass no compiler params at all and the jax_pallas_use_mosaic_gpu config flag is true, and that flag defaults to true. It picks Triton when you pass a Triton CompilerParams, or when compiler params are absent and the flag is false. Both paths warn at this commit: using pl.pallas_call for a Mosaic GPU kernel is deprecated in favour of plgpu.kernel, and the Triton backend is deprecated altogether. There is one hard override in the other direction: Mosaic GPU raises on AMD ROCm devices and the error tells you to pass Triton compiler params instead, so on an AMD card the deprecated backend is the only backend.
Why does jnp.dot fail in my Mosaic GPU Pallas kernel?: Because there is no lowering rule for it. Both Pallas GPU backends lower a kernel by walking the jaxpr and looking each primitive up in a dictionary, and a missing key raises NotImplementedError naming the primitive. The Triton table has an entry for lax.dot_general_p, roughly two hundred lines of it. The Mosaic GPU lowering has no occurrence of dot_general anywhere, under any of its four semantics tables. The matmul is not missing from the backend, it is spelled differently: you call plgpu.wgmma against an accumulator reference on Hopper, or plgpu.tcgen05_mma against a tensor-memory accumulator on Blackwell, and then wait on the asynchronous instruction before reading the accumulator. Two smaller ones surprise people for the same reason: jnp.sqrt has no Mosaic GPU rule although lax.rsqrt does, and jnp.cumsum, jnp.argmax and lax.transpose are all Triton-only.
What is warp specialisation in Pallas, and why does Triton not have it?: It is giving different warpgroups in the same thread block different jobs. plgpu.kernel takes a num_threads argument whose docstring says the number does not correspond to CUDA threads but to warpgroups, so num_threads=3 launches three warpgroups, each with its own index, and the kernel body branches on that index. A production attention kernel for Hopper uses two compute warpgroups and one memory warpgroup that does nothing but issue asynchronous copies, and it raises the compute warpgroups to 232 registers while dropping the memory warpgroup to 40, because the register file is a fixed pool shared across the block. Triton has no equivalent. Its num_warps parameter defaults to 4, meaning 128 threads that all run the same code over different lanes of the same tile. A warp there is a slice of the data, not a role.
Why do kernel libraries ship separate sm90 and sm100 kernels?: Because the accumulator lives somewhere else. Hopper issues its matmul with wgmma and accumulates into registers, which is why the compute warpgroups in a Hopper attention kernel need a large register budget. Blackwell issues tcgen05_mma and accumulates into tensor memory, a separate on-chip memory, which frees the register file for the softmax warpgroup and changes what every warpgroup around it can afford. Layered on that, Blackwell adds cluster-level collective MMA where two thread blocks cooperate on one instruction, a four-warp mesh with one warp pinned per tensor stream, and a refusal to run unstable softmax at all. The two configs end up as different classes with different field names rather than different values of one class, which is why autotuning caches are kept in separate per-device directories, and why the dispatcher branches on the type of the config rather than asking the device.
Does a kernel that works in interpret mode work on the GPU?: Not necessarily, and the reason is exactly this page. Interpret mode evaluates the kernel jaxpr on the host, so it never consults either backend rule table. A body full of jnp.dot and jnp.cumsum will produce correct numbers in interpret mode and then raise during lowering on Mosaic GPU. The compiler is the one component interpret mode does not model, and the backend rule table is the compiler. The practical habit is to read which of two failure classes an error belongs to: a message naming a primitive means the body is wrong for the backend, and a NotImplementedError raised before any primitive is named means the call is wrong for the backend, such as a dynamic grid bound or a scalar prefetch operand.

Pallas on a GPU: Two Backends, and What Each One Lets You Write

Questions

Related explainers