Down to the Metal: Profiling JAX & Pallas on TPU

23 Oct 2026

Stop trusting wall-clock timers on TPUs! In the third episode of the Bielik on TPU series, we open up the hardware and compiler toolbox to see what actually happens on the TPU v5e chip when running native JAX and custom Pallas kernels for the Bielik 1.5B LLM.

In the first two episodes, we built a memory-bound RMSNorm kernel and an autotuned bfloat16 Matmul kernel on the MXU - but we measured both using Python wall-clock time. Today, we go down to the metal with headless JAX trace profiling and post-optimization XLA High-Level Optimizer (HLO) inspection to separate Python dispatch overhead from pure on-chip execution.

Key highlights covered:

• Headless TPU profiling: capturing jax.profiler.trace outputs and parsing Chrome Trace Event Format (trace.json.gz) directly in Python
• JAX Native vs. Pallas on RMSNorm: how XLA emits two loop fusions with a buffer copy while our autotuned Pallas kernel runs as a single fused custom-call
• Solving the Episode 1 mystery (Host vs. Device gap): discovering why small tensors (512x512) suffered a 13.8x host dispatch overhead
• Fixing the Pallas profiler blind spot: why XLA reports 0 bytes accessed and 0 FLOPs for opaque Pallas custom-calls, and how to unlock true hardware roofline metrics with pl.CostEstimate and pl.estimate_cost at zero runtime overhead
• Hardware-level proof of Matmul autotuning: measuring pure on-chip MXU throughput - jumping from 12.37 TFLOPS (fixed 128 block) to 162.25 TFLOPS (autotuned Pallas), compared against XLA’s 187.23 TFLOPS
• Peeking into the compiler brain (hlo.txt): inspecting post-optimization XLA HLO and the structural boundary between XLA fusion and Pallas custom-calls (Matmul + Residual)

Resources:

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