Down to the Metal: Profiling JAX & Pallas on TPU

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

Writing TPU Kernels in Pallas: Matmul on the MXU

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

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

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

Bielik Anatomy 🧠 Ep.7 - Assembling a Full LLM from Custom Triton Kernels

Seventh episode of the Bielik Anatomy series. It’s moment of truth time - every custom Triton kernel built throughout the series (RMSNorm, MatMul, RoPE, Flash Attention, and SwiGLU) is now wired together into a fully functional Bielik 1.5B Instruct model architecture.

The episode walks through constructing the complete decoder layer step by step, tackles the unexpected challenge of handling bias in linear and activation layers, loads the official pretrained weights from HuggingFace safetensors, and spins up an interactive chat interface to talk with Bielik live.

Key engineering highlights covered:

• Bielik 1.5B architecture overview and why biases forced a change in kernel strategy
• Custom 2D-grid Embedding kernel vs. native PyTorch
• Zero-cost MatMul Bias fusion via tl.constexpr compile-time flags - no branching overhead
• Fused SwiGLU with Bias computed entirely in GPU registers
• Decoder Layer assembly integrating GQA, RoPE, and Flash Attention
• Autoregressive sampling mechanics and full model construction
• Benchmarks on RTX 4060 Ti: a 28% speedup in Time To First Token (TTFT) - plus a deep dive into the memory bandwidth bottleneck that limits throughput during generation without a KV Cache

Resources:

• GitHub Repo

Bielik Anatomy 🧠 Ep.6 - SwiGLU Kernel Fusion: 30% Faster with Triton on RTX 4060

Sixth episode of the Bielik Anatomy series. Eliminating redundant read/write cycles to GPU global memory can unlock massive performance gains. Benchmarked on an RTX 4060, a custom Triton kernel delivers a 15-30% speedup in TFLOPS compared to standard PyTorch and torch.compile.

The episode walks through the Triton code line by line: managing L2 cache swizzling for maximum block reuse, designing dual accumulators in registers for parallel matrix multiplications, and applying the SiLU (Swish) activation directly on-chip before writing back to HBM - fusing the entire SwiGLU operation into a single kernel pass.

Resources:

• GitHub Repo