Pallas, XProf and a design review

Chapter 9 ended at the edge of what XLA lets you control. This chapter covers the last two tools: writing the kernel yourself in Pallas, and measuring what actually ran with XProf. It closes with a design review that uses every chapter.

Pallas on TPU

Pallas is JAX's kernel language. A kernel is a Python function over Refs, mutable views of memory, invoked once per point of a grid. BlockSpecs say which block of each array each grid point sees [PL-grid]. On TPU, Pallas lowers through Mosaic, an MLIR-based compiler. The result is embedded in the XLA program as a custom call, scheduled alongside XLA's own fusions [PL-tpu].

The tiled matmul below runs, and is checked against jnp.dot, in scripts/ch10_pallas_matmul.py (interpret mode on CPU; drop interpret=True on TPU):

def matmul_kernel(x_ref, w_ref, o_ref, acc_ref):
    # Grid axes (i, j, k): output tile (i, j), contraction step k (innermost).
    @pl.when(pl.program_id(2) == 0)
    def _():
        acc_ref[...] = jnp.zeros_like(acc_ref)

    acc_ref[...] += jnp.dot(x_ref[...], w_ref[...], preferred_element_type=jnp.float32)

    @pl.when(pl.program_id(2) == pl.num_programs(2) - 1)
    def _():
        o_ref[...] = acc_ref[...].astype(o_ref.dtype)

pl.pallas_call(
    matmul_kernel,
    grid=(m // bm, n // bn, k // bk),
    in_specs=[pl.BlockSpec((bm, bk), lambda i, j, kk: (i, kk)),
              pl.BlockSpec((bk, bn), lambda i, j, kk: (kk, j))],
    out_specs=pl.BlockSpec((bm, bn), lambda i, j, kk: (i, j)),
    out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
    scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],
    compiler_params=pltpu.CompilerParams(
        dimension_semantics=("parallel", "parallel", "arbitrary")),
)(x, w)

What each part does on a TPU [PL-tpu] [PL-pipe]:

• The grid runs sequentially, in lexicographic order. There are no concurrent thread blocks. Grid points execute one after another on a TensorCore. Putting the reduction axis last means the output block stays put across consecutive steps, and Pallas skips re-copying a block whose index did not change.
• BlockSpecs drive the DMAs. Arrays passed to pallas_call live in HBM. For each grid step, Pallas copies the blocks the index maps select into VMEM, double-buffered by default: the copy for step n+1 runs while step n computes. pipeline_mode=pl.Buffered(n) deepens the buffer.
• The (8,128) rule. The last two dimensions of every block must be divisible by 8 and 128, or equal the array's full dimensions, because of Chapter 2's tiling.
• Memory spaces.
  • VMEM is the default for kernel refs, and pltpu.VMEM scratch persists across grid steps; here it holds the fp32 accumulator.
  • SMEM holds scalars and indices.
  • pl.ANY leaves an array in HBM for manual DMA.
• dimension_semantics. Marking grid axes "parallel" lets Pallas split them across the two TensorCores of a megacore chip; the docs say this applies to v4 and v5p. "arbitrary" marks the reduction axis.

Tile size is arithmetic intensity

Each grid step streams a $b_m \times b_k$ block of x and a $b_k \times b_n$ block of w from HBM and does $2 b_m b_k b_n$ FLOPs. Its arithmetic intensity is about $b_m b_n / (b_m + b_n)$, independent of $b_k$ [SB1]. VMEM bounds how big the tiles can be.

Computed by scripts/ch10_pallas_matmul.py. VMEM counts double-buffered input and output blocks and the fp32 accumulator. The compiler's scoped-VMEM limit per kernel, set by xla_tpu_scoped_vmem_limit_kib, can be lower than the physical size.

Beyond the matmul

• Scalar prefetch and block sparsity. Index maps can read small integer arrays prefetched into SMEM. A block table then decides which KV page each grid step loads, and that is how paged attention is written [PL-sparse].
• Remote DMA. pltpu.make_async_remote_copy sends a VMEM buffer to another chip's VMEM over ICI, with semaphores for completion. The distributed Pallas guide builds ring AllGather, AllReduce and bidirectional ReduceScatter from it, which is how a collective matmul is hand-written [PL-dist].
• SparseCore and per-core kernels. pl.kernel with core and subcore meshes targets individual TensorCores and SparseCore subcores [PL-core].

When to write one. Write a Pallas kernel when XLA cannot express the access pattern, as with ragged or paged attention, MoE grouped matmuls, quantized matmuls with unusual scaling and collective matmuls with custom overlap. Do not write one when the problem is a missing fusion or a bad sharding; those are fixed with constraints and flags (Chapter 9). Libraries such as Tokamax and tpu-inference's kernels/ directory already cover the common cases [K-tokamax].

Capturing a profile

From JAX [XP-jax]:

with jax.profiler.trace("/tmp/tb"):
    y = step(params, batch)
    y.block_until_ready()
# then: xprof --logdir /tmp/tb   (or TensorBoard's Profile tab)

jax.profiler.start_server(9999) lets you capture on demand from the XProf UI instead. XProf's overhead on TPU is under 1% [XP-home].

From tpu-inference [R-tpuinf-prof]:

• examples/tpu_profiling.py profiles a chosen model, input length, output length and batch. The guide suggests batch 1 with a long input for prefill, and batch 256 with one-token input and output for decode.
• PHASED_PROFILING_DIR captures prefill-heavy, decode-heavy and mixed batches automatically from a running server.
• USE_JAX_PROFILER_SERVER=True enables on-demand capture.

Reading the Trace Viewer

The Trace Viewer shows one timeline per TensorCore [XP-trace] [SB9].

• The row that matters is XLA Ops: the operations that actually ran, named by their HLO names. Rows built from jax.named_scope and the Python stack are approximations.
• All cores run the same program, so TPU:0 is usually enough.
• Repeated blocks are layers.
• Keys: W and S zoom, A and D pan. Clicking an op shows its HLO and source line.

Two signatures to recognize:

• A synchronous collective is a long bar labelled all-gather, reduce-scatter or all-reduce on the XLA Ops row with nothing beside it: the TensorCore is waiting.
• An overlapped collective appears as all-gather-start ... all-gather-done (or collective-permute-start/done), with compute fusions between the pair [X-async].

When the first signature appears where you expected the second, the usual causes are:

• The collective's result is needed immediately, so there is nothing to overlap it with.
• The latency-hiding scheduler is off, or was constrained by memory.
• XLA chose a synchronous form because of a sharding mismatch: often a large unexpected AllGather, the "compiler tickling" case of Chapter 9.

Reading an op

An op in the trace or a dump looks like this [SB9]:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)} fusion(bf16[32,32,8192]{...} %fusion.32), kind=kCustom, calls=%all-reduce-scatter.3

From left to right:

• The name.
• The output type and shape.
• The layout: {2,1,0} lists dimensions from most minor to most major.
• The tiling T(8,128)(2,1): (8,128) tiles with bf16 pairs packed into 32-bit words [X-tiled].
• The memory space: S(1) is VMEM; no S, or S(0), is HBM.
• The op, its operands, and attributes. A fusion with calls=%all-reduce-scatter.N is a ReduceScatter. On TPU, a fusion contains at most one matmul plus the elementwise work around it.

The other tools

• HLO Op Profile: time per op, grouped by category, with achieved FLOP and bandwidth utilization. Use it to find the few ops worth attention [XP-opprof].
• Roofline Model: plots every op, up to the top 1,000, by arithmetic intensity against achieved FLOP/s, with HBM and VMEM rooflines (and CMEM on v4). An op on the slanted part of the roofline is memory-bound; on the flat part, compute-bound. An op well below both lines is waiting on something else, usually latency or communication [XP-roof].
• Memory Viewer: HBM allocation over program order and the peak. Use it to debug out-of-memory errors and see what is live at the peak [XP-mem].
• Graph Viewer: the HLO graph around an op, with shardings.
• Megascale Viewer: DCN traffic for multislice jobs.

Design review

The capstone uses every chapter. Given a model, a chip, a latency target and a cost goal, produce the mesh, quantization, batch, expected step time and the bottleneck in one sentence. Try it with the calculator in Chapter 7 before reading the model answer.

Brief. Serve LLaMA-3-70B on Ironwood. Prompts are 8k tokens; conversations grow to 32k tokens. Time to first token ≤ 500 ms; time per output token ≤ 30 ms. Maximize tokens per second per chip.