Compilers & Kernels

What happens between your model code and the silicon: the compiler that splits one program across a device mesh, decides which operations fuse, assigns every buffer an address, and emits the collectives that move tensors between chips.

Sub-topics

• XLA — The compiler and runtime under JAX and PyTorch/XLA: HLO, the pass pipeline, fusion, layout, buffers, SPMD, and the backends.
• Pallas — JAX’s kernel language and the libraries built on it: BlockSpecs and a grid in Python, lowered to Mosaic on a TPU and to Triton or Mosaic GPU on a GPU.

Explainers

• The SPMD Partitioner: One Program Across a Device Mesh — You write a program as if one giant machine ran it, annotate a few tensors with how they are split across a mesh of chips, and the compiler rewrites it into N identical per-device programs — inserting every all-gather, reduce-scatter, all-to-all, and halo exchange needed to keep the maths correct. We trace how those annotations propagate through the graph, where each collective is born (and which resharding forces it), the one windowed op that needs halo exchange, and the case the partitioner cannot solve on its own — drawn, computed, and animated.
• Anatomy of a FlashAttention Kernel — Plain attention is not slow because it does too much math — it is slow because it writes an N×N score matrix out to memory and reads it back, twice. FlashAttention never writes that matrix at all: it walks the keys in tiles, keeps a running max and denominator so softmax stays exact block by block, and recomputes what it needs in the backward pass instead of storing it. We compute the materialization bill, step through the online-softmax recurrence on a real tiny attention, do the IO complexity honestly, and cover what FA2 and FA3 actually changed — drawn, computed, and animated.
• Inside XLA: ~200 Passes and the Fusion Decision — A traced program enters XLA as a graph of abstract operations and leaves as a handful of GPU kernels. In between, a couple hundred passes rewrite it — and one of them, fusion, decides which operations share a kernel by putting a stopwatch on every merge. We walk the ordered pass pipeline, the memory round-trip fusion deletes, the priority cost model (time_unfused − time_fused) and its greedy queue, the vetoes that refuse a merge, buffer assignment, codegen, and the recompilation you pay for all of it — drawn, computed, and animated.
• The Loop That Runs Every XLA Pass: hlo_pass_pipeline.cc, Line by Line — A compiler pass is small and well behaved: it takes a program, makes one kind of improvement, and says whether it found anything. Running a few hundred of them back to back is a different job, and this is the 334-line C++ file that does it, read end to end. The two flags that skip passes, their four-shape grammar, and the two occurrence counters they resolve against. The invariant checkers between passes, why they only run after a pass that changed something, and the single rule that makes them safe. The debug mode that hashes the whole module twice to catch a pass that lied about its own return value. The dump condition with a special case hiding inside it, where the catch-all pattern writes fewer files than a narrow one. And the written ledger every pass appends a row to, opened before the skip decision so its numbering survives a bisect. The signature exhibit is the loop itself, ported line for line over a seven-pass pipeline, so you can watch a flag reshape a run.
• The PJRT Plugin Contract: pjrt_api.cc, Line by Line — A vendor ships a shared library for an accelerator the framework has never been compiled against, and it runs your model anyway. The entire arrangement is one exported C symbol and two integers, and this is the 209-line file where the two sides meet. We read all of it: the global map from a device-type string to a function table, the lowercase that is the whole matching rule, the write-once registration that refuses to let two packages claim one name, and the loader itself, which dlopens a library, asks for a single symbol called GetPjrtApi, and never closes the file again because the table lives inside it. Then the handshake. Where the plugin gets its version number from, why major must match exactly and minor only has a floor, the environment variable that switches the rule from generous to strict, and the value that reads as encouraging while doing the opposite. The signature exhibit runs the real comparison order: set the version the plugin was built against, set the environment variable, and read back the exact error string absl::StrCat assembles, trailing-parenthesis quirk included.
• HLO Module Anatomy: What XLA Holds While It Compiles — An HLO dump is an object graph printed. This walks that graph in the XLA source at a pinned commit: an HloModule that owns a vector of computations and points at one of them as the entry, an HloComputation whose parameters and return value are both just instructions in its own body, and an HloInstruction that owns nothing beyond an opcode, a shape and a list of pointers out to operands and back from users. Then the two things that sit beside the graph and decide most of what it costs: the shape that carries a layout, where minor_to_major turns an index into an address and a disagreement turns into a copy you never wrote; and the schedule, an optional total order added late to hold peak memory down. It closes on the verifier that runs between all two hundred passes, in the order it runs its checks, and the refusals worth knowing by name.
• Reading an XLA Dump: What the Compiler Writes, and How to Read It — Someone asks you to attach the HLO dump. You set one flag, point a directory at it, and get back a pile of files with names like module_0000.jit_step.0007.simplification.after_algsimp.before_reshape-mover.txt. This reads that directory end to end at a pinned commit of openxla/xla: the four inference rules one flag triggers before anything is written, the filename grammar and what every segment of it means, the two named bookends of a compile and the four buffer-assignment reports that ride with the second one, the per-pass files and the rule that quietly drops any pass that changed nothing, and the flags that make the dump larger, smaller, or unreadable. Then the two tools that take a dumped file and do something with it without a framework anywhere in sight: hlo-opt, which re-runs named passes and prints any stage of the compile, and run_hlo_module, which executes the module and checks it against the interpreter. The signature exhibit assembles the actual directory listing from the real inference rules, one flag at a time.
• How a StableHLO Module Becomes HLO: mlir_to_hlo.cc, Line by Line — Your framework does not hand the compiler a graph. It hands it a versioned document written in a dialect called StableHLO, and 434 lines of C++ decide what that document is allowed to say. We read all of them. The includes, which split into the two vocabularies this file sits between. The thirteen MLIR passes a module walks through before anything called HLO exists: seven from a Shardy pipeline it pulls in unconditionally, six added here by name, ending with the one that copies captured constants into control-flow regions because XLA has no implicit capture. The Shardy-loses-to-GSPMD fallback at the top of the function, and the block sixty lines later that silently reads the flag it just switched off. The parser, and the one error message that names three different version problems because it genuinely cannot tell them apart. Then the write side, which is the longer half: the version clamp that keeps the lower of what a plugin asked for and what this build can produce, the allow-list that refuses to serialize any operation nobody has promised to read, and the twelve-week compatibility window written as a requirement rather than a number. The signature exhibit runs the real branches: pick a direction, flip the flags, and every pass your module walks through appears in order, at the line that adds it.
• Layout Assignment: Where the Copy in Your HLO Dump Comes From — A copy instruction shows up in your dump between two lines you wrote next to each other, and the only thing that changed across it is the little list in braces. This reads the pass that put it there, at a pinned commit of openxla/xla. What a layout is, which is one permutation of dimension numbers read from the fastest-moving end. The three kinds of constraint the pass collects, and why exactly one of them can be settled by inserting a copy. The switch statement that answers, for all 134 opcodes, whether an operation may hold one arrangement out and a different one in, and the 89 that may not, which is what lets a constraint travel across most of a graph for free. The deque with two ends, the three propagators hanging off it, and the cap of two rewrites that stops it oscillating. The derivation at a transpose, whose stated goal is to make the transpose move no bytes at all, and the copy that lands one edge above it as a result. Then the priority arithmetic that settles a collision, the row-major default handed to whatever nobody argued about, the tuple rule that sends a constraint further upstream instead, and the moment inside AssignLayouts where the copy is finally created. The signature exhibit runs the real propagation loop over a four-instruction module, one pop at a time, with every layout on screen produced by the same arithmetic the pass uses.
• Buffer Assignment: How Every Value Gets an Address — A compiled XLA program never allocates. It receives a handful of slabs from the runtime and every kernel was compiled with its offsets inside those slabs already baked in. This walks the pass that picks them, at a pinned commit of openxla/xla: the values that must share memory because the language says so, the schedule that turns a dependency graph into a clock so that a live range can be two integers, the closed-interval interference test and the one endpoint case it makes an exception for, the colour that is a memory space rather than a graph colouring, the four kinds of value that never reach the packing, and the biggest-first best-fit heap that packs everything left. Then the nine ordered reasons an allocation refuses a buffer, the in-place update whose copy you never see until something reads the old value, and the four files a dump writes with the answer in them. The signature exhibit packs a real schedule with the real rules and lets you drag a live range until the sharing breaks.
• What Shares a GPU Kernel: priority_fusion.cc, Line by Line — A GPU kernel launch costs about a microsecond of nothing happening, plus a round trip through memory for its inputs and another for its outputs. So the largest single decision an ML compiler makes is which operations get to share a kernel, and on the GPU backend that decision is 1,419 lines of C++ we read end to end. The type choice the whole pass turns on: a priority is an absl::Duration, the wall-clock time this merge is estimated to save, which is what lets an ordered map double as a priority queue and two infinities put bitcasts first and constants last with no special case in the ordering. The correction to the usual summary: the queue holds producers, not edges, and a producer is scored against every one of its consumers at once, so the cost of being duplicated into three kernels is always in the number. The incremental machinery that keeps the pass from being quadratic, where a re-score is the old score plus the delta from the new consumers minus the runtimes of the departed ones. The Triton path that is tried first and the elemental checks that run when it declines. The thirteen refusal strings, quoted verbatim, that are what you actually read in a fusion dump. And the three flags that let you watch it decide on your own model, including the compiler fuel that bisects to the exact merge that changed your numbers. The signature exhibit runs the real ordering rules over a small graph: press pop, watch a producer that scored below zero turn profitable because a neighbour was absorbed somewhere else.
• Instruction Fusion Legality: instruction_fusion.cc, Line by Line — Merging two operations into one kernel deletes a write and a read of a whole tensor, which on this hardware is the difference that matters. So the compiler wants to fuse everything, and this 1,248-line C++ file is where it works out what it may. The unit of decision is one graph edge: a producer, a consumer that takes it as an operand, and one boolean. Read end to end: the hand-written classification of all 134 opcodes into cheap and expensive, in a switch with no default case; the two opcodes that may always be copied and why copying them lowers memory traffic; the global pre-pass that bans a producer from duplication unless every consumer will swallow it; the reverse-post-order queue and the thirty-line comment about the duplicate clone it exists to prevent; and the quarter of the file that is about correctness rather than speed, where a slice read out of a buffer meets an update written back into it. The signature exhibit is the six-gate decision ladder ported from the file, run over one small module, so you can watch an edge get refused and see which predicate did it.
• IFRT Arrays: pjrt_array.cc, Line by Line — A jax.Array spread over 512 chips is one Python object, and underneath it is a list of ordinary single-device buffers plus a description of how they tile the whole. This 696-line C++ file is the class that holds those two things together, and the walk reads all of it. The validator that runs before any array exists, which matches buffers to devices strictly by position and compares only the devices this process can address. The six construction paths, one of which quietly skips the validator. The three copy semantics that the header describes as different and that one function implements identically, with the TODO admitting it. Disassembly, which hands back one array per shard without moving a byte. The read back to the host, which refuses anything but a single shard unless the array is replicated. And a copy path where the same invariant is re-derived against a different device list. The signature exhibit runs the validator itself: pick a sharding, perturb one buffer, and watch which of the five checks refuses it and with which string.
• The CPU Thunk Executor: thunk_executor.cc, Line by Line — A GPU gets a driver queue to submit work to. A CPU gets nothing, so the compiled program arrives as a flat list of runtime actions called thunks, and something has to work out what may run at the same time. This 841-line C++ file is that something, read end to end. Each unit declares which buffers it reads and which it writes, and every dependency in the program is derived from those declarations and nothing else, by a conflict test where read-after-read is free. Then a transitive reduction deletes the edges other edges already imply and hands each node a priority equal to how much it unblocks. At run time it is one atomic counter per node and a ready queue: a thunk finishes, decrements its successors, and any counter that hits zero names work that can start now. We read the two edge kinds and why sharing a collective communicator is a weaker constraint than sharing memory; the three thresholds that quietly make a small program single-threaded; the recursive halving that spreads forty ready thunks over four threads without forty queue pushes; and the completion count over sink nodes that makes success and failure end the same way. The signature exhibit is a real ten-thunk graph you step one loop iteration at a time, with all three ready-queue classes.
• The GPU Codegen Path: From a Fused HLO to a Kernel — When the optimizer stops, your model is a few hundred fusions and a handful of library calls, and none of them is machine code yet. This reads the two decisions that finish the job, at a pinned commit of openxla/xla. First the kind: a string an upstream pass stamped on the backend config, or, when there is no stamp, ten questions asked about the fusion roots in a fixed order until one answers yes, with a tenth that always does. Then the class the switch builds from that kind, including the two extra questions the loop arm asks before it settles, one of which turns a copy into a memcpy and skips the kernel entirely. Then the three families behind those classes: six emitters that write a module in an MLIR dialect XLA defined for the purpose and lower it through forty-one passes to LLVM; the Triton path, which builds one module and hands it to another compiler; and the library paths, where cuBLAS arrives as a custom call and cuDNN arrives as a fusion, for reasons that follow from what each library accepts. Then the autotuner: the only pass that compiles and runs real kernels while it is still compiling, the four ways it can get a config, the sentences it uses to refuse, the two-microsecond window that lets a thriftier kernel win over a faster one, the clustering that treats correctness as agreement rather than truth, and the two-tier cache with a key caveat the source states out loud. The signature exhibit runs the real dispatch function over seven candidate fusions and shows which question fires and which class gets built.
• XLA Collectives: One Guest List, Then One Order — A collective is the one instruction in a compiled program that cannot be executed by looking at its own operands, because it needs values from machines that do not share memory with this one. This reads the whole subsystem at a pinned commit of openxla/xla, from both sides of the line. On the compiler side: the list of integers an instruction carries, the two optional fields that decide whether those integers are replica ids, partition ids or positions in a flat enumeration of the device grid, and the switch with four arms that multiplies one written group out against a mesh into the sets of global device numbers the runtime will actually open communicators for. Then the flattened id and the device number, which walk the same grid in opposite directions and disagree on any mesh with more than one of each. On the runtime side: ranks, communicators, cliques and the factory that makes them, all defined without naming a vendor; the future that means launched rather than finished, and means two different things on GPU and CPU; the rendezvous that agrees on membership before any bytes move; and the three separate places where the answer to how do we avoid a deadlock is the same answer. The signature exhibit runs the real group arithmetic and the real verifier checks over a mesh you dial.
• Writing Pallas Kernels in Colab: What Interpret Mode Proves, and What It Cannot — One keyword argument runs a TPU kernel on a laptop, and almost every confusing Pallas afternoon comes from not knowing what that run checked. This reads the machinery at a pinned commit of jax-ml/jax. There are three interpreters behind one argument, not one: a generic scan over the grid that models the arithmetic and removes memory entirely, a TPU mode that simulates HBM, VMEM, DMAs and semaphores and carries a happens-before race detector, and a real chip, the only rung with a compiler that can refuse you. What each rung catches and what it structurally cannot, drawn as a walk. Which import names which chip, with the memory-space tables side by side and the SMEM collision that means scalar memory on one and shared memory on the other. What interpret=True literally executes: a while loop over the grid, dynamic slices for blocks, outputs and scratch pre-filled with NaN. The padding the interpreter adds before the loop starts, computed live from the real function. Why the grid promises a count and not an order, and the one parameter that randomizes it. Then the case worth committing to: an index map one block past the end, which raises nothing, returns finite numbers, and quietly duplicates half the output. The debugging surface as a picker, one question in and one switch out. And the six-cell notebook shape, including the reset cell that exists because the interpreter keeps its simulated memory after an exception on purpose.
• pallas_call, Line by Line: One Call, Three Calling Conventions — Every Pallas kernel on every backend enters JAX through one 1381-line Python file. This walks it whole, at a pinned commit of jax-ml/jax. The file declares a single primitive, then spends the rest of itself teaching JAX what that primitive means under each transform: an abstract evaluation that decides the output types and refuses a non-manual mesh, a lowering rule that swaps its own body for a scan when you ask for interpret mode, a batching rule that answers vmap by growing the grid rather than looping, and a public function whose real work is turning a pile of keyword arguments into one GridMapping. The signature exhibit runs the real slot arithmetic out of that file: dial the grid, the scalar prefetch operands, the inputs, the outputs and the scratch, and watch the three orderings the primitive, the kernel and the index map each demand.
• The Mosaic TPU Pipeline, Line by Line — A Pallas kernel on a TPU computes from VMEM, a scratchpad of on-chip memory small enough that the array you passed in does not fit. So every block has to be copied in from HBM before the core can touch it, and copied back out afterwards, and if the core waits for either of those the matrix unit sits idle for most of the schedule. One file decides when each copy is issued, when the core is allowed to wait for it, and how many copies of a block VMEM holds at once. This walks that file at a pinned commit of jax-ml/jax: the buffered ref and its four independent cursors, the semaphore per slot, the four boolean predicates that are the entire schedule, the prologue that primes the buffers stage by stage, the loop body whose line ordering is the only reason anything overlaps, and the epilogue that waits for the last send nobody else waited for. The signature exhibit runs the real predicates over a grid you dial and prints the resulting schedule, slot by slot.
• Pallas on a TPU: From a Jaxpr to the Mosaic Dialect — A Pallas kernel is a Python function, but by the time a TPU sees it the Python is long gone: what the backend receives is a jaxpr, and what it emits is an MLIR module in the Mosaic dialect wrapped inside one HLO custom call. This reads that translation at a pinned commit of jax-ml/jax. The type mapping first, where a ref becomes a memref carrying a memory-space attribute and a value becomes a vector, and the four-step chain that turns a BlockSpec with no memory space into the literal string vmem. Then the window parameters that decide whether a block is copied in for you or left in HBM for your own DMA, and the two allocators behind run_scoped, one calling alloca and one calling a semaphore allocator. Then the table: 135 primitives carry a TPU rule at this commit, computed from the real registrations and filterable by core type and by the MLIR dialect each rule builds into. Then the two doors to the MXU, one of them a single tpu.matmul and the other five primitives that name an accumulator and an MXU index. It closes on refusal: the dtypes the lowering will not take, the block shapes it rejects before lowering starts, and the one message that means no rule exists at all.
• BlockSpec and the Grid: How Pallas Cuts an Array Into Kernel-Sized Pieces — A Pallas kernel never sees your array. It sees one block of it, chosen by a Python function you wrote and a loop the compiler ran, and this page reads the arithmetic that connects the two at a pinned commit of jax-ml/jax. The grid is a loop nest and the kernel is its body, so the kernel runs prod(grid) times. A BlockSpec is a block shape plus an index map, and the index map returns block indices rather than element indices: the start of the slice is the block size times the number your map returned, which is one match statement in core.py and nothing more. From there the page walks the five kinds of block dimension and what each one asks your map to return, the three defaults and why omitting a BlockSpec entirely does not mean the whole array but an unblocked ref in whichever memory the backend chose, the two memory-space vocabularies and where DEFAULT lands on a TPU, and the conversion that turns your closure into a traced jaxpr along with the three rules that conversion enforces. Then the honest parts: a block shape that does not divide the array runs anyway, on a full-sized block, reading padding whose values the documentation tells you to assume are garbage, discarded on output and not on input; the grid runs with the last axis fastest on the interpreter and in an order you may not assume on a chip; and a TPU wants the last two dimensions of your block to be multiples of 8 and 128. The signature exhibit ports the real block-index arithmetic, the ceiling division, the start indices and the padding rule, so you can dial an array shape against a block shape and read the slice the kernel would receive.
• Pallas on a GPU: Two Backends, and What Each One Lets You Write — One pallas_call on an NVIDIA card reaches one of two compilers, and they do not accept the same kernel body. Each backend is, quite literally, a dictionary from JAX primitive to lowering function, and a primitive with no entry raises before a single instruction is emitted. This reads both dictionaries at a pinned commit of jax-ml/jax: 101 primitives on the Triton side, 117 on the Mosaic GPU side, 59 in common, and a matmul that lowers on one and has no rule at all on the other. Then it reads a production kernel library at a pinned commit of google/tokamax to see what a team does with that: a Triton attention kernel in one file, a Mosaic GPU dispatcher that owns no kernel, and two kernel bodies forked by GPU generation because the accumulator moved out of the registers on Blackwell. The signature exhibit checks one softmax body against the real rule tables, line by line.
• Flash Attention on a TPU, Line by Line — The flash attention algorithm fits on a napkin. The Pallas kernel that runs it on a TPU is 1715 lines: four kernel bodies, three separate launches, eleven block sizes, and a four-line diagonal test consulted from ten places. This walks the whole file at a pinned commit of jax-ml/jax.
• Splash Attention: When the Mask Stops Being Arithmetic and Becomes the Loop — Flash attention treats a causal mask as a value: compute the whole block of logits, then overwrite half of them with negative infinity. Splash attention treats the same mask as a schedule. Before the kernel is ever traced, a Python pass cuts the mask into blocks the size of one grid step, labels each block empty, partial or full, and turns that label grid into three small integer arrays that live in TPU scalar memory. The kernel then reads its own loop structure out of them: an empty block is not computed and its data is not fetched, a full block skips the masking entirely, and for a banded mask the grid itself is rebuilt narrower. The signature exhibit and the dial both run the real classification out of the pinned file, so the block counts on this page are the counts the compiler would produce.
• The tokamax Op and Its Autotuner, Line by Line — A fast kernel is never one kernel. It is a family of implementations and a family of tilings, and something has to choose between them on every call. Tokamax, the fused-kernel library Google builds on Pallas, answers that in one 701-line Python file. This walks it whole at a pinned commit: the frozen Op dataclass and the two hooks a backend overrides, the twelve lines that turn a call into a dispatch key made only of argument shapes and dtypes, the five-step ladder that decides whether you get an explicit config, a cached measurement, an unmeasured guess or an error, the measurement loop that fills the cache in, and the payload the file writes into the compiled program so a captured production trace can be tuned on a machine that never runs your model. The signature exhibit is that real key, opened against the library’s own shipped cache: pick a ragged-dot workload and a chip, and read every cached implementation with its winning tiling and its measured median.
• Attention for sm90, Line by Line — A production flash attention kernel for the H100, read whole at a pinned commit of google/tokamax. The algorithm is the one everybody knows; what a real kernel adds on top of it is six hundred lines of scheduling. One thread block launches three warpgroups and each branches immediately into a different role: two do the arithmetic on 232 registers a thread, and the third drops to 40 and does nothing but issue transfers, running a stage ahead of the maths. Automatic synchronisation is switched off at launch, so every handshake in the file is one of sixteen numbered hardware barriers placed by hand, and running out of them is a real failure mode. The page walks the constants, the shared-memory plan, both configuration heuristics, both warpgroup roles and the launch, then steps through one query tile: which slot each warpgroup is holding, how far ahead the loader has run, and where the causal mask starts biting. A second exhibit runs the file own shared-memory estimate so you can watch a tiling stop fitting in 227 KiB.