Writing TPU Kernels in Pallas: Matmul on the MXU

27 Sep 2026

From 12 TFLOPS to matching native XLA at 140 TFLOPS: How do you build an efficient Matrix Multiplication (Matmul) kernel for Google TPUs using JAX and Pallas?

In the second episode of the Bielik on TPU series, we leave memory-bound kernels behind and dive straight into compute-bound workloads on the TPU’s Matrix Execution Unit (MXU). We walk through mapping LLM projections (QKV, Output, LM Head), writing the reduction loops in Pallas with @pl.when guards, managing on-chip VMEM accumulator scratchpads, and building an automated grid autotuner to unlock peak hardware throughput.

Key highlights covered:

• Matmul in LLM architectures and mapping Bielik projections to the MXU
• Baseline tiled Matmul kernel in Pallas with reduction loops and @pl.when boundary guards
• Why naive Pallas hits a 12 TFLOPS wall on TPU v5e
• Hardware alignment (8/128 constraints): meeting Mosaic’s strict alignment requirements to avoid hardware stalls
• VMEM budgeting: pre-filtering scratchpads under the 15MB limit to prevent compile-time out-of-memory errors on TPU v5e / v5 lite
• Hand-rolling an autotuner for Pallas & Mosaic: tuning tile geometry to saturate the MXU and crack 140 TFLOPS

Resources:

• GitHub Repo
• Previous Episode (Ep.1 - RMSNorm)
• Bielik Anatomy (GPU Triton Series)
• JAX Documentation
• Pallas / TPU Mosaic Guide