Pallas, Mosaic, and where Triton fits

XLA fuses at the granularity of HLO ops using a cost model you do not control, and some kernels, flash attention is the canonical one, cannot be expressed as a good fusion at that granularity. When you hit that wall you stop describing what to compute and start describing how. In JAX that door is Pallas, and it lowers through two very different backends: Triton on GPU and Mosaic on TPU.

You write to references, not values

A Pallas kernel does not look like normal JAX. Normal JAX is functional: values in, values out. A Pallas kernel is imperative and operates on references: you are handed input and output Refs and you load from and store to them explicitly.

def kernel(x_ref, o_ref):
    o_ref[...] = jnp.maximum(x_ref[...] * 2.0 + 1.0, 0.0)

n, blk = 1024, 256
def f(x):
    return pl.pallas_call(kernel,
        out_shape=jax.ShapeDtypeStruct((n,), jnp.float32),
        grid=(n // blk,),
        in_specs=[pl.BlockSpec((blk,), lambda i: (i,))],
        out_specs=pl.BlockSpec((blk,), lambda i: (i,)))(x)

The shape of the thing is the whole point: a pallas_call node in the jaxpr with a grid mapping, block mappings, and the kernel as a subcomputation. The grid is how many programs run, and the BlockSpecs say which slice of the input each program gets.

Two backends, one language

Pallas lowers the same kernel source through two very different backends:

• Triton on GPU: the kernel becomes Triton, which becomes PTX, which becomes a cubin. This is the same path a hand-written Triton kernel takes, but with JAX's type system and shape inference in front.
• Mosaic on TPU: the kernel becomes Mosaic, which becomes code for libtpu. The TPU backend is entirely different because the hardware is entirely different.

This is the escape hatch from XLA's fusion decisions. When the compiler's cost model picks the wrong fusion granularity, you write the kernel yourself, and Pallas gives you one language for both GPU and TPU.

When you hit the fusion wall, you stop describing what to compute and start describing how.

Back to the blog