Pallas: a JAX kernel language#
Pallas is an extension to JAX for writing custom kernels for GPU and TPU. Use it to achieve peak performance when XLA’s automated optimizations fall short—specifically when you need to:
Write custom fusions that XLA cannot generate automatically.
Take explicit control over memory movement and pipelining.
Directly access low-level hardware features.
Both backends build on shared ideas (Refs, BlockSpecs, and pipelining)
but provide distinct, hardware-specific APIs.
Note
Pallas is under active development. See the Pallas Changelog for recent changes.
Quickstart for GPU
Custom kernels for NVIDIA Hopper (H100) and Blackwell (B200) GPUs.
Quickstart for TPU
Custom kernels for TensorCore and SparseCore on TPUs.