Pallas

Pallas is a Python DSL for writing TPU kernels, analogous to Triton for GPUs. It is part of the JAX ecosystem and provides a higher-level abstraction than raw XLA. The key concept is the Ref (mutable array reference) which lives in a specific memory space (HBM or VMEM). Kernels operate on VMEM Refs, and data must be explicitly copied between HBM and VMEM. The emit_pipeline function handles pipelining (overlapping memory transfers with compute) automatically. The matmul tutorial shows real TPU performance: 91% FLOP/s utilization at 8192x8192x8192 with bf16, which is competitive with hand-optimized TPU kernels. Pallas is the recommended way to write custom TPU kernels, just as Triton is for NVIDIA GPUs. CORRECTED 2026-10-05: this record previously stated Pallas was TPU-only with 'no path to NVIDIA or AMD GPUs'. That was true of the source it cited and is no longer true of Pallas. JAX documentation now carries a full Mosaic GPU backend (jax.experimental.pallas.mosaic_gpu) with its own programming model, memory spaces and pipelining guides, and a fully specified Blackwell matmul guide showing 69.44% TensorCore utilisation - 109.6% of cuBLAS and 109.3% of CUTLASS - after six optimisation steps from a 37.62% basic kernel. No separate GPU record was added, because this IS Pallas: a second 'pallas-gpu' record would fragment one entity across two backends and make the TPU comparison harder, not easier. See source w7k-pallas-mosaic-gpu. The GPU programming model is NOT the TPU one, and that is the substantive difference: a Pallas 'thread' on GPU is a warpgroup (4 warps = 128 CUDA threads) running in lockstep, with GMEM/SMEM/TMEM memory spaces and TMEM staging on Blackwell, where the TPU backend uses HBM/VMEM with explicit copies. AMD remains unsupported - there is no ROCm backend, so 'no AMD path' is still accurate and is the surviving half of the old claim. The honest comparison for a GPU team is therefore no longer Pallas-versus-Triton: on Blackwell, documented Pallas Mosaic GPU beats cuBLAS and CUTLASS on this benchmark. Note JAX's own caveat on those numbers - 'don't believe matmul benchmarks if they don't specify input data distribution', with iid-normal fp16 described as one of the slower distributions - and that the GPU API is still under jax.experimental, so stability is the real argument for Triton over Pallas on GPU, not performance.