Pallas
Also known as Pallas kernels, JAX Pallas, TPU custom kernels
In plain English
Pallas is the way engineers write hand-tuned kernels for TPUs when the compiler alone does not schedule an operation well enough.
Technical definition
Pallas is the JAX extension for writing custom TPU and GPU kernels that control tiling, memory movement between HBM and VMEM, and the scheduling of MXU and vector unit work explicitly.
Engineering details
XLA handles most of a model, but performance-critical operations such as ragged paged attention, grouped MoE matmuls, and Gated DeltaNet recurrences need explicit control over block sizes, double buffering, and lane layout. Pallas provides that control in Python. The TPU inference optimizations in the Official Preview are mostly Pallas kernels: GDN v3 fuses Conv1D and GDN into one kernel, the grouped matmul triple-buffers expert weights, and the attention kernel adopts a sequence-on-lane layout. Helion’s TPU backend also emits Pallas. TorchTPU can call Pallas kernels from native PyTorch, so kernels written for TorchAX transfer with wrapper and layout adjustments.
Why it matters
Pallas plays the role that CUDA C++, CUTLASS, and Triton play on NVIDIA. Native PyTorch support does not remove the need for it; tensor shapes and layouts still have to be tuned to feed the MXU. How fast the ecosystem writes Pallas kernels for new models sets the pace of TPU externalization.
How to read it in InferenceX
Reported Pallas kernel wins on Ironwood include GDN v3 speedups of 1.41x decode, 1.60x prefill, and 2.14x mixed batches at the kernel level, an 11.3% throughput gain at concurrency 512 from asynchronous state transfers, and a 49% decode throughput gain from splitting attention fetch and compute block sizes.