Writing Custom TPU Kernels: Can Pallas Beat XLA? (Bielik LLM on TPU v5e - RMSNorm)

13 Sep 2026

Can custom Pallas kernels on Google Cloud TPUs beat XLA’s aggressive compiler fusion?

In this first episode of the Bielik on TPU kernel series, we transition the open-source Bielik-1.5B architecture from GPU Triton kernels over to Google Cloud TPU v5e hardware using JAX and Pallas.

We break down the TPU software stack, examine how Pallas maps memory references directly to TPU tiles via BlockSpec, and build a custom RMSNorm kernel from the ground up. Finally, we run end-to-end memory bandwidth benchmarks comparing custom Pallas against native JAX across varying batch sizes and hidden dimensions.

Key highlights covered:

• Bielik-1.5B structural details (GQA, SwiGLU, linear biases)
• The TPU software stack: JAX, XLA, Pallas, and the JAX-to-Mosaic lowering pipeline
• Writing TPU kernels using Ref memory abstractions, BlockSpec, and pallas_call
• Debugging TPU kernels locally with the Pallas CPU interpret mode
• Memory-bound layer behavior and benchmarks: sequence length and hidden size sweeps on Google Cloud TPU v5e evaluating custom Pallas vs. XLA compiler fusion

Resources:

• GitHub Repo
• Bielik Anatomy (GPU Triton Series)
• JAX Documentation
• Pallas / TPU Mosaic Guide