Writing TPU Kernels in Pallas: Matmul on the MXU (from 12 to 140 TFLOPS)

Описание: 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 Episode 2 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.

Code Repository & Benchmarks: https://github.com/qooba/bielik-on-tpu

Timestamps:
0:00 - Introduction: Shifting to Compute-Bound on the MXU
0:42 - Matmul in LLM Architectures (Bielik Anatomy)
1:35 - Implementing the Baseline Matmul Kernel in Pallas
3:15 - Handling Tiling Boundaries & VMEM Accumulators
4:40 - First Benchmark: Why Naive Pallas Hits a 12 TFLOPS Wall
6:10 - Hand-Rolling an Autotuner for Pallas & Mosaic
7:45 - VMEM Pre-filtering & Hardware Alignment (8/128 constraints)
8:55 - Autotuned Benchmark: Cracking 140 TFLOPS on TPU v5e
10:15 - Next Episode: Deep-Dive Hardware Profiling with jax.profiler

Key Takeaways:
Hardware Alignment: Mosaic requires strict alignment to avoid hardware stalls.
VMEM Budgeting: Pre-filtering scratchpads under a 15MB limit prevents compile-time out-of-memory errors on TPU v5e / v5 lite.
Tile Geometry: Why compute logic is only half the battle on TPUs - tile geometry determines whether the MXU is saturated or starved.

If you enjoy bare-metal accelerator programming, low-level ML systems, and custom kernel engineering, consider subscribing!

#JAX #GoogleTPU #Pallas #MachineLearning #DeepLearning #LLM #GPUComputing #HighPerformanceComputing