Code: PeppaKing8/nanogpt-jax · Benchmarks: TPU v6e-16
If you have trained a neural network on a TPU, you have probably been told: write clean JAX, let XLA handle it. That is genuinely good advice. XLA is a strong optimizing compiler — it fuses element-wise ops, tiles matrix multiplies, and schedules data movement across the memory hierarchy without any programmer intervention. For a large fraction of deep learning workloads, it is effectively the ceiling.
But "effectively the ceiling" and "actually the ceiling" are different things. Some algorithms carry structure that a compiler cannot infer from the tensor operations alone. Causal attention is a canonical example: the upper triangle of every attention score block is semantically zero and need not be computed — but XLA sees a QKᵀ followed by a softmax, not a loop that can skip half its iterations. Sparse attention patterns, ring-buffer KV caches, and block-structured recurrences have the same property. For these, there is performance left on the table.
Pallas is JAX's answer to this gap. It allows programmers to write custom kernels — in JAX — that take direct control over which memory tiles are loaded, in what order, and how partial results accumulate. This post is a guided tour of that capability, using a nanoGPT-style transformer as the testbed. We will cover the TPU hardware background that motivates kernels in the first place, the mechanics of writing a Pallas attention kernel end-to-end (including the often-overlooked backward pass), and what happens when promising ideas fail to survive from op-level benchmarks to full model training. By the end, the code is released and the 1.57× training speedup on v6e-16 is fully reproducible.
Understanding when a custom kernel helps requires understanding what the hardware and compiler already do for you — and exactly where their visibility ends.
A TPU v6e is organized around three components. The Matrix Multiply Unit (MXU) is a systolic array that executes dense 128×128 bf16 matrix multiplies at very high throughput. The Vector Processing Unit (VPU) handles element-wise operations, reductions, and transcendentals. High-Bandwidth Memory (HBM) holds all tensors that do not fit on chip — model weights, activations, gradients — and is accessed at several hundred GB/s. That bandwidth, while impressive in absolute terms, is a fraction of what the MXU can consume when fully utilized on a dense matmul. This means the TPU's true bottleneck for most operations is memory bandwidth, not arithmetic throughput. Any algorithm that avoids unnecessary trips to HBM wins.
There is also a small, fast on-chip memory (VMEM in Pallas/Mosaic terminology) that sits between HBM and the compute units. Loading a tile from HBM into VMEM and computing on it multiple times — rather than reloading it for each use — is the central optimization pattern for custom kernels.
When you run a JAX function through jit, JAX traces it into an XLA computation graph (HLO), which the XLA compiler then optimizes and lowers to TPU machine code. XLA is excellent at operator fusion: rather than writing every intermediate tensor to HBM and reading it back, it chains compatible operations into a single pass. A linear layer followed by a layer norm, for instance, is typically executed as one fused kernel — the linear output never touches HBM between the two ops. XLA also chooses tensor layouts, determines tiling strategies for matrix multiplications, and schedules collective communications in multi-device programs.
What XLA cannot do is reason about algorithmic loop structure that is implicit in the tensor expression. Given a QKᵀ contraction, XLA sees a dense matmul. It does not know that a triangular causal mask will zero out approximately half the result, that those zeros could have been avoided entirely by not computing them, or that an online recurrence over K/V blocks could replace the entire materialized matrix. Expressing that knowledge requires a kernel.
Pallas is JAX's framework for writing custom kernels that bypass XLA's automatic tiling. A Pallas kernel defines a grid of independent programs. Each program receives its position in the grid via pl.program_id(axis), uses BlockSpec to declare which slice of each input tensor it should work on, and calls pl.load / pl.store to move tiles between HBM and VMEM. The whole kernel is invoked through pl.pallas_call, which integrates cleanly with JAX's JIT, autodiff, and sharding infrastructure. On TPU, Pallas programs are compiled by Mosaic, JAX's TPU-specific lowering framework.
The key distinction from just writing JAX code is control: a Pallas kernel decides which tiles to load (enabling sparse skipping), in what order (enabling cache-friendly access patterns), and how to accumulate partial results (enabling online algorithms that never materialize large intermediates).
The simplest Pallas kernel to build intuition with is a tiled vector add: split the input into chunks, assign one program to each chunk, and add. The snippet below shows the full structure: a kernel function that reads two input tiles and writes one output tile, and a pallas_call wrapper that specifies the grid, output shape, and the BlockSpec mapping each program ID to its tile.
import jax
import jax.numpy as jnp
import jax.experimental.pallas as pl
def add_kernel(x_ref, y_ref, out_ref, *, block):
pid = pl.program_id(0) # which tile am I?
offs = pid * block + jnp.arange(block) # my slice of indices
x = pl.load(x_ref, offs) # load from HBM into VMEM
y = pl.load(y_ref, offs)
pl.store(out_ref, offs, x + y) # store result back
def add(x, y, block=128):
n = x.shape[0]
return pl.pallas_call(
add_kernel,
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
grid=(jax.lax.cdiv(n, block),),
compiler_params=dict(mosaic=dict(dimension_semantics=("parallel",))),
)(x, y, block=block)The dimension_semantics=("parallel",) hint tells Mosaic that all programs in this grid are independent and can run in any order — the TPU equivalent of a CUDA thread block with no inter-warp dependencies. For attention, the grid will be (batch, heads, query_tiles), and the programs will not all be independent (the online softmax accumulates across K/V tiles within one query program), but the basic structure is the same.
Standard attention computes softmax(QKᵀ / √d) · V, materializing a full S×S score matrix in the process. For a sequence length of 8192, this matrix is 256 MB per head in float32 — far too large to fit in VMEM, and expensive to move through HBM even once. FlashAttention solves this with a tiled online algorithm: instead of materializing the full score matrix, it processes attention in blocks of queries and keys, maintaining a numerically stable running softmax state (m, l, acc) across K/V iterations. The output is computed block-by-block and never requires the full S×S matrix to exist anywhere.
This algorithm is a natural fit for Pallas. The outer loop over query tiles maps to the Pallas grid dimension. Each program owns one query tile, iterates over K/V tiles in a loop (the inner loop, explicitly written by the programmer), and accumulates its running state. For causal attention, the inner loop can terminate early: any K/V tile whose key positions all lie in the future of the query positions can be skipped entirely, without loading K or V from HBM at all.
The update(m, l, acc, s, V) call implements the standard online softmax recurrence: it takes the current running max m, normalizer l, and accumulator acc, computes the new per-block max m′ = max(m, max(s)), rescales the old state to account for the change in scale, and adds the new contribution. After all K/V blocks have been processed, dividing the accumulator by the normalizer yields the exact softmax-weighted sum — numerically equivalent to the materialized version, but computed without ever forming the S×S matrix.
The two implementation details that most affected throughput were the K/V BlockSpec (an incorrect tile shape can double HBM traffic for a mathematically identical kernel) and the causal block skip (the kernel must avoid loading and dotting future K/V tiles entirely, not just mask their contributions after computing them).
def causal_flash_fwd_kernel(q_ref, k_ref, v_ref, o_ref, m_ref, l_ref,
*, block_q, block_k):
b, h, q_block = pl.program_id(0), pl.program_id(1), pl.program_id(2)
q = load_q_block(q_ref, b, h, q_block, block_q)
acc = jnp.zeros([block_q, head_dim], jnp.float32)
m = jnp.full([block_q], -jnp.inf, jnp.float32)
l = jnp.zeros([block_q], jnp.float32)
for k_block in range(0, q_block + 1): # causal skip
k = load_k_block(k_ref, b, h, k_block, block_k)
v = load_v_block(v_ref, b, h, k_block, block_k)
scores = dot(q, k.T) * scale + causal_mask(q_block, k_block)
m, l, acc = online_softmax_update(m, l, acc, scores, v)
store(o_ref, b, h, q_block, acc / l[:, None])
store(m_ref, b, h, q_block, m)
store(l_ref, b, h, q_block, l) # saved for backwardA tempting shortcut is to write only the forward kernel and rely on JAX's automatic differentiation for the backward pass. This does not work well. JAX's autodiff will differentiate through the Pallas kernel by treating it as an opaque function and applying a finite-difference or checkpointing strategy, but the resulting backward computation typically does not respect the memory discipline of the forward. For attention specifically, a naive autodiff backward will materialize the full S×S attention probability matrix — exactly the thing the forward kernel was designed to avoid.
The native backward recomputes the attention probabilities tile-by-tile from the saved forward statistics (m, l), then accumulates the gradients for Q, K, and V in the same tiled fashion as the forward. This requires implementing a custom VJP (via jax.custom_vjp), which is more work but is the only way to preserve the HBM-efficiency of the full kernel.
This is not a Pallas quirk — it is the same reason FlashAttention-2 ships a hand-written backward kernel. The lesson generalizes: for any memory-efficient forward algorithm, the backward must be written to match it. Until a native backward exists, op-level forward benchmarks are not evidence that a kernel will help training.
A strong op-level benchmark does not automatically transfer to a full model. The gap here came from the preprocessing boundary. Before each attention call, the model applies Q/K RMSNorm, RoPE positional encoding, and a layout transposition from (B,S,H,D) to (B,H,S,D). If the Pallas attention kernel consumes (B,H,S,D) tensors but the rest of the model produces (B,S,H,D), that transposition happens on every attention block in every layer of every training step. At 12 layers, this overhead is not negligible.
The solution is to fuse the preprocessing into a second Pallas kernel: a single program that performs Q/K RMSNorm, applies RoPE, and lays out the result in the format the attention kernel expects — with a matching custom VJP for the backward. Only when both kernels are active does the end-to-end speedup fully materialize. The configuration switches are attention.impl=ours (Pallas attention fwd/bwd) and attention.preprocess=pallas (Pallas preprocessing fwd/bwd).
| Path | W&B run | Median step | Speedup |
|---|---|---|---|
| XLA baseline | zu8z0sdp | 242.81 ms | 1.00× |
| Pallas attn + Pallas preprocess | mgzf30mi | 154.45 ms | 1.57× |
Dense causal attention reduces HBM traffic by avoiding materialization of the score matrix, but it still computes every lower-triangular block. Sparse attention patterns go further: they reduce actual floating-point work by skipping attention computation for token pairs that are definitionally irrelevant. Pallas is particularly well-suited for this because the skip decision is made before any K/V tile is loaded — there is no padding, no masked tensor, and no wasted compute on elements that will be zeroed out.
StreamingLLM is an efficient long-context attention pattern built on two observations. First, the first few tokens of a sequence act as "attention sinks" — tokens that receive disproportionate attention mass regardless of content, likely because the softmax must put its probability mass somewhere. Second, local context (recent tokens) is usually sufficient for the next-token prediction task. StreamingLLM therefore attends only to the first K sink tokens and the most recent W tokens, ignoring the middle of the sequence entirely. For a query at position i, the attended key positions are:
$$j < K \quad \text{or} \quad \max(K,\, i - W + 1) \le j \le i.$$
In the Pallas kernel, this translates to a simple block-level skip: before loading any K/V tile, the kernel checks whether the tile's key range falls in the sink region, the window region, or the middle. Only tiles in the first two regions are loaded and computed. The middle tiles are skipped without any HBM read at all. For decoding, a ring-buffer KV cache keeps sink tokens in fixed slots while recent tokens rotate through window slots, bounding cache memory to O(K + W) regardless of total generated length.
At S=2048, W=512, gains over dense Pallas are modest because the sparse ratio is small. At S=8192, W=512, the forward is approximately 1.4× faster and backward 1.8× faster than dense causal Pallas at the op level. Model-level results are pending: sparse training changes the attention distribution, so the meaningful comparison is speed at matched perplexity, not speed alone.
PyTorch's FlexAttention introduced a clean API for custom attention sparsity: the user writes a boolean function mask_mod(b, h, q, k), the system precomputes a block-level sparsity mask, and a compiled kernel skips all inactive blocks. Our JAX/Pallas implementation follows the same design. The user provides a mask function; build_block_mask evaluates it at tile granularity and produces a block-level boolean tensor; the attention kernel iterates only over tiles the mask marks active, skipping the rest before loading K/V or computing any dot products. Within active tiles, fine-grained causal masking is still applied at the token level.
def local_causal_mask(b, h, q, k, window=512):
# True if key k should be attended by query q
return (k <= q) & (k >= q - window + 1)
block_mask = build_block_mask(
local_causal_mask, batch, heads, seqlen, block_q, block_k
)
out = flex_pallas_attention(q, k, v, block_mask, causal=True)These numbers are striking, but the caveat deserves emphasis: the baseline here is explicit masked JAX attention (compute everything, then mask), which is a particularly weak starting point. The honest next comparison is against dense Pallas at matched quality. That experiment requires running sparse training to convergence, which is marked as ongoing work.
The most educational part of the project was discovering where custom kernels do not help. Several ideas that appeared promising at the op level disappeared entirely when measured end-to-end. The failures tend to fall into three categories.
The most tempting early idea was to fuse the QKV projection with attention — after all, Q, K, and V are computed immediately before attention is applied, so the fusion seems natural. The fused kernel consumed the pre-projection activations directly and output the attention result in one pass. In op-level experiments at small batch sizes, it showed marginal improvement. At model scope, it measured 154.34 ms/step, essentially tied with the 154.45 ms from the unfused Pallas path.
The reason, in retrospect, is that the QKV projection is a large dense GEMM that XLA already handles well. Fusing it does not eliminate the layout transposition that follows — and that transposition, sitting at the boundary between the projection output format and the attention kernel's expected format, is where the actual cost lives. No amount of projection fusion removes a boundary that lives downstream of the projection. The correct kernel to write is the one that eliminates the transposition, not the one adjacent to it in the graph.
Several kernels looked excellent in isolation and disappeared in whole-model training. The fused vocabulary cross-entropy kernel is the starkest example. The final logits tensor — shaped (B·S, V) with vocabulary size in the tens of thousands — is large, and a kernel that avoids materializing it fully should in principle save a lot of HBM traffic. In the single-op benchmark it did. In the full distributed model, the savings were absorbed by the resharding and all-to-all collective communications required to move the tensor between devices in the post-attention, pre-softmax layout — overhead invisible at single-op scope but substantial at model scope.
The general rule is: an op-level benchmark measures one operation with fixed inputs. A model-level benchmark measures that operation plus everything that must happen before and after it to keep the rest of the model running correctly. A kernel can introduce new costs (layout changes, device transfers, broken fusions) that only appear in the second measurement.
W4A16 quantization — storing weights in 4-bit integers and dequantizing at use time — dramatically reduces the memory footprint of large models. It is a legitimate technique for fitting larger models on a given amount of HBM. As a throughput path on a TPU, however, it performed worse than standard bf16 on the tested shapes. The reason is architectural: TPU MXU is optimized for bf16 matrix multiplies and executes them at near-peak throughput. Dequantizing 4-bit weights back to bf16 before the matmul adds a non-trivial overhead, and without native int4 compute support, the bf16 matmul path is simply faster than the dequant-then-multiply path at the sequence lengths we tested.
The central lesson from this project is not that Pallas is powerful — it is — but that it is powerful in a narrow and specific way. The productive targets are operations where the algorithm has structure the compiler cannot infer: causal and sparse attention, long-context memory layouts, decoding ring buffers, and any component where the researcher knows a sharper algorithm than the one implied by the tensor expression. Everything else — dense GEMM, standard layer norm, optimizer steps — is better left to XLA.
For dense causal language modeling at the scales tested (12-layer nanoGPT, v6e-16, S=2048), a Pallas attention kernel with a native backward and a matching preprocessing kernel yields a real 1.57× end-to-end training speedup. For sparse attention, the op-level results are substantially larger, and the open problem has shifted to the right one: maintaining training quality while reducing floating-point work. The code is at PeppaKing8/nanogpt-jax. Key files: pallas_attention.py (dense kernel), sparse_attention.py (StreamingLLM and FlexAttention-style), model.py (dispatch via attention.impl and attention.preprocess).
python -m nanogpt_jax.train \
--config configs/smoke.yaml \
--override attention.impl=ours \
--override attention.preprocess=pallas \
--override training.num_iterations=3