Why a kernel language lives inside JAX

XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → compiles most machine-learning programs well, and then some program hits its limit. The documented way out has been CustomCall: hand the compiler an opaque function you wrote yourself. The design document is blunt about what that costs. Writing one means writing C++, and on GPU it means learning CUDA, which the authors call arguably too low-level for many machine learning GPU kernels, matrix multiplication among them. Their reading is that even expert users have trouble implementing an efficient matmul or multi-headed attention that way.

There is a second motive listed right next to the first, and it is about time rather than difficulty. Advances in systems research take a while to land inside XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →, and people want to run ahead of the compiler. A hand-written kernel is how you run ahead; later the compiler absorbs whatever the kernel proved out. So the escape hatch is not only for the cases XLA will never cover. It is also for the cases XLA has not covered yet.

The kernel path already had you writing Pallas by week two, at /s/pallas, without ever asking why the language exists in this shape. That question is this lesson.

Triton changed the terms of the argument by showing that a kernel could be written as array programming instead of thread programming. You describe what happens to a tile, the compiler handles the parallelism underneath, and the result is good enough that Triton became the primary code generation route for torch.compile through Torch Inductor. The design document treats this as settled evidence, not as a competitor to be argued with.

The obvious next move would be to point Triton at a TPU, and the document says exactly why that does not work. Triton exposes a TPU-like model already, programs written for tiles of arrays in L1 cache, but it is specialized enough to GPU that it cannot be compiled directly for TPU. The example given is atomics: Triton offers atomic operations meant to handle parallel writes, and parallel writes are not what a TPU does. So the tile-level model needed a front end one level above Triton, abstract enough that a single kernel could reach two very different machines.

Having ruled out the C++ hatch and the Triton port, the document asks its own question and answers it in one sentence. The open question, as written: is JAX a good fit for a kernel language at all?

Triton demonstrates that an array programming language can be practical for writing GPU kernels and JAX is just that.

That line, from the Pallas design document, is the whole argument compressed. JAX is a mature tracing front end for numerical computing, its users already write NumPy-style array code, and its transformations are the reason people use it. Pallas is then described as three extensions and nothing more: Ref types so you can talk about memory, a handful of new primitives like program_id, and pallas_call to run the body over a grid.

Reusing tracing buys something an AST-parsing front end cannot offer. Your kernel is ordinary Python at trace time, so closures, higher-order functions, and any templating you can express in Python all work, and the document says outright that this makes Pallas far more amenable to templating than Triton. A kernel factory is a function that returns a function. Nothing in the toolchain needs to know it happened.

def make_eltwise_add(eltwise):
    def kernel(x_ref, y_ref, o_ref):
        o_ref[...] = eltwise(x_ref[...] + y_ref[...])
    return kernel

doubled = make_eltwise_add(lambda v: v * 2)
exponentiated = make_eltwise_add(jnp.exp)   # two kernels, one source

Pallas accepts a subset of JAX primitives, and the design document names two exclusions with their reasons attached. On conv_general: convolution usually is not offered as a primitive in the underlying hardware. On gather/scatter: the underlying compiler may not support noncontiguous memory reads and writes. Both reasons point at the machine, not at the implementation calendar.

Read as a gap, that list looks like work someone has not gotten to. Read as a position, it says something firmer: a kernel language whose substrate is tiles and contiguous moves should not offer you an operation that hides an arbitrary access pattern. The gather you actually need does not disappear, it changes form. It becomes an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → fed by prefetched scalars, which is a schedule you can look at and cost, and the scalar-world lesson at /s/pallas/scalar-world is that mechanism in full.

The TPU reference applies the same position at finer grain. Integer reductions are unsupported. Elementwise operations carry a published cost ranking where jnp.sin and jnp.cos sit in the expensive tier and jnp.exp in the middle. Loop primitives get fully unrolled during compilation, so a large trip count is a compile-time problem. None of that reads like a feature backlog. It reads like a language declining to pretend the hardware is uniform.

A language that refuses things can promise things. The TPU page states the promise plainly: while the features are experimental, a kernel accepted by the compiler must return the expected results, and if your outputs look wrong the instruction is to compare against a run with interpret=True and file a bug report. Correctness is not negotiable in exchange for the experimental label; expressiveness is.

The cost is that everything the compiler used to decide is now yours to state. Which bytes are resident, in what order, in which memory space, with which axis of the grid reused. The next five pages are those decisions one at a time, and each of them is a decision only because the language declined to guess.