GPU, kernels & serving · the compiler track

AI Compilers — From Graph to Kernel

A linearized rebuild of the modern machine-learning compiler: how a high-level tensor program — the y = gelu(x @ w + b) you wrote — becomes a handful of fused, hardware-specialized GPU kernels, derived as one chain of decisions where each pass is forced by an inefficiency the last one left behind.

This is the systems sibling of the GPU Kernels and Triton tracks. Those teach you to hand-write a fast kernel; this teaches you how a compiler generates them automatically — what torch.compile, XLA, TVM, and TensorRT actually do between your Python and the SASS the GPU runs. It is a hub: where kernels, serving, or training go deeper, it links out instead of repeating.

The one idea the whole track turns on

A neural network is a mathematical graph; a GPU is a bandwidth-and-FLOPs machine. Between what you wrote (a hardware-agnostic graph of tensor ops) and what the hardware wants (a few large, fused, specialized kernels) sits a wide semantic gap, and naive eager execution leaves the chip mostly idle. A compiler closes that gap by lowering the graph through levels of representation, and every pass it runs answers one question: how do I preserve the math exactly while moving fewer bytes and launching fewer, more specialized kernels? Each level of IR exists only because the level above it could not express the next optimization.

Who this is for

You can read Python and PyTorch, you know what a tensor op and a GPU kernel are (or you've read a few lessons of the GPU track), and you want to stop treating torch.compile/XLA as magic. No compiler-course background assumed — SSA, IR, passes, scheduling, and codegen are all built up from first principles. By the end you should be able to look at any pass in any ML stack and say which inefficiency forced it, and sketch the pipeline yourself.

The spine — one forced chain

Read the diagram left to right. The top row is the compiler's five stages; below, each lesson is a wall, and the arrow to the next is the idea it's forced to introduce. Nothing is here "because compilers have it" — each link exists because, without it, the chain breaks.

The lessons

The semantic gap

The premise: one line gelu(x @ w + b) is, eagerly, 3+ kernels and 3+ HBM round trips — the hardware wanted one fused kernel. The five stages of any ML compiler, and the wall→fix→new-wall spine of the whole track. Interactive: a gap meter — stretch an op chain and watch eager waste grow while the fused cost stays flat.

Part I · The gap

The eager baseline and its ceiling

Price what every later pass must beat: per-op dispatch (~µs), kernel launch (~µs), and the dominant HBM round trip (bytes / bandwidth). The two ceilings — launch-bound (many tiny ops) and bandwidth-bound (each op re-reads HBM) — are whole-program problems. Interactive: an eager cost profiler that flips between the two regimes.

Part I · The gap

Capturing the graph

You can't optimize what you can't see. Three capture strategies and their trade: tracing (bakes in shapes, drops control flow), scripting/AST (control flow, language subset), and bytecode analysis (TorchDynamo → FX, with guards and graph breaks). Interactive: a capture simulator — watch tracing silently drop a branch while Dynamo guards and breaks.

Part II · Front end

The intermediate representation

To rewrite a graph safely you need rules: SSA / value semantics, a typed dataflow DAG (shape, dtype, device in the type), and multiple levels of IR — high (framework ops) → mid (linalg) → low (vector) — because no one abstraction hosts every optimization. Lowering trades portability for control. Interactive: a lowering ladder for softmax through four IR levels.

Part II · Front end

Rewrites & canonicalization

The pass / pattern-rewrite engine — match a subgraph, replace with an equivalent, iterate to a fixed point — plus the staples: constant folding, CSE, DCE, algebraic identities, canonicalization. The hard constraint: preserve semantics (mind float non-associativity). The frontier: equality saturation / e-graphs. Interactive: a rewrite playground that shrinks a graph to its fixed point.

Part III · Middle end

Operator fusion — the central win

The highest-leverage pass. Merge a region into one kernel that loads from HBM once, computes in registers/SRAM, writes once — turning ~2×len round trips into ~2. Vertical, horizontal, and matmul-epilogue fusion; FlashAttention as algorithm-level fusion; and when not to fuse (register pressure, compute-bound GEMMs). Interactive: a fusion-grouping DAG with a register-pressure meter that breaks on over-fusion.

Part III · Middle end

Memory planning

Fusion and the graph spawn many intermediate buffers; allocate each naively and peak memory is the sum, not the max. Liveness analysis, buffer reuse by graph coloring, in-place ops, buffer donation/aliasing, and arena allocation collapse peak toward max-concurrent. Interactive: a memory planner — lifetime bars on a timeline, toggling reuse and in-place.

Part III · Middle end

Layout & data formats

A buffer's physical shape decides coalescing, tensor-core eligibility, and transpose count. Pick formats (NCHW vs NHWC, tiled/blocked for tensor cores, packed/padded) and propagate them to cancel adjacent transposes — a global assignment problem where a locally good layout can be globally terrible. Interactive: a layout chooser that inserts or cancels transposes.

Part III · Middle end

Scheduling — the loop nest

The graph says what to compute; the same math has astronomically many loop nests, and the choice swings perf 10–100×. Halide's algorithm/schedule split and the knobs — tile, reorder, vectorize, unroll, parallelize, compute_at, stage to SRAM — plus the polyhedral model in a paragraph. Interactive: a schedule explorer where no-tiling sinks and tiles-too-big spill.

Part IV · Back end

Code generation

A schedule is just a plan; codegen emits real instructions. Two paths: LLVM IR → NVVM → PTX → SASS (max peak, locked, verbose) or a tile DSL like Triton as the backend (fewer lines, portable) — why Inductor chose Triton. Host launch code vs device kernel; the code cache. Interactive: a codegen target picker trading peak perf against portability.

Part IV · Back end

Autotuning & cost models

The schedule space (tiles × warps × stages × …) is enormous and shape-dependent; no human picks it per kernel. Benchmark candidates on device (Triton @autotune, AutoTVM) or prune with a learned cost model (Ansor). The core trade: compile-time search vs run-time speedup, amortized over a million runs. Interactive: an autotune search where a cost model finds the optimum with far fewer measurements.

Part IV · Back end

Dynamic shapes & control flow

Every pass so far assumed static shapes; serving has variable batch, growing KV cache, and data-dependent branches — the defining hard problem. Symbolic shapes, guards + recompilation, bucketing/padding (waste FLOPs for cache hits), and control-flow ops (cond/scan) instead of graph breaks. Interactive: a specialization cache — compile storm vs bucketing vs one symbolic compile.

Part V · The hard parts

Autodiff as a compiler pass

Training needs a backward graph no one wrote. Reverse-mode AD as a graph-to-graph transform builds a joint forward+backward graph the same passes then optimize (AOTAutograd). Two enablers: functionalization (kill mutation so rewrites are sound) and the save-vs-recompute partitioner — activation checkpointing as a compiler min-cut. Interactive: a recompute partitioner trading memory against backward FLOPs.

Part V · The hard parts

Distributed compilation

The graph exceeds one device, and placing hundreds by hand is intractable. Make partitioning a pass: annotate a few tensors with sharding specs over a device mesh, let the compiler propagate shardings (GSPMD/Shardy/DTensor), insert collectives where layouts disagree, and overlap comm with compute. Interactive: a sharding propagator — one all-reduce vs all-to-all-everywhere.

Part V · The hard parts

Operator coverage — decomposition & the long tail

Every pass since codegen assumed the compiler can generate every op — but a framework has thousands, and the thousandth has no kernel. Decompose the huge op set to a small closed core (primTorch / Core ATen) so it's O(ops + backends), not O(ops × backends); for the rest, register custom ops as opaque (with fake/meta kernels) or fall back to a vendor library — anything but a graph break. Interactive: one unknown op mid-chain, and how each handling choice saves or shatters fusion.

Part VI · Shipping the compiled model

Quantization & precision lowering

The biggest remaining attack on bytes-moved isn't a better kernel — it's spending precision. Insert quantize/dequantize nodes, pick scales by calibration (PTQ) or learning (QAT), and fuse the dequant into the matmul epilogue so you compute in int8/fp8. The twist on the track's "preserve semantics" motif: this is the one pass that's deliberately not bit-exact — it preserves the math on an accuracy budget. Plus fp8, int4 weight-only, and 2:4 sparsity. Interactive: a precision dial with an int4 accuracy cliff.

Part VI · Shipping the compiled model

The runtime & executor

The compiler emits an inert artifact — kernels plus a memory plan. Something must run it: streams & async dispatch, the caching allocator that realizes the lesson-06 plan without malloc on the hot path, and CUDA-graph capture/replay that collapses N launches into one driver call — finally killing the launch-overhead ceiling from lesson 01. The catch: graphs need stable shapes, so they fight dynamic shapes. Plus the AOT path (torch.export → AOTInductor). Interactive: a launch-overhead killer that breaks under dynamic shapes.

Part VI · Shipping the compiled model

Debugging compiled models

Fusion reassociates floats, quantization is approximate by design, graph-replay caches addresses — every pass that promised "preserve semantics" can change results or silently recompile. Tell acceptable float drift from a real bug (rtol/atol vs an fp32 reference), bisect with the minifier, read the graph-break/recompile/guard logs, and navigate the determinism-vs-speed trade (TF32, accumulation dtype, atomics). Interactive: a divergence detective that triages drift from a planted bug.

Part VI · Shipping the compiled model

The real stacks

The same chain ships under a dozen names. A Rosetta stone lining up torch.compile (Dynamo → AOTAutograd → Inductor → Triton), XLA/HLO, TVM/Ansor, the MLIR ecosystem, and TensorRT against the five stages — same pipeline, different IRs and bets (JIT vs AOT, training vs inference, open vs vendor). Interactive: a stack Rosetta that lights up each component per stage.

Part VII · Synthesis

Capstone — one block, source to kernel

Follow y = gelu(x @ w + b) with a dynamic batch and its backward from Python through every stage — capture, IR, rewrite, autodiff, decompose, fuse, quantize, plan, lay out, schedule, codegen, autotune, runtime/CUDA-graph, guard, shard — watching the scoreboard collapse (~5 kernels → 1–2, ~10× bytes → ~2×). Then the frontiers: megakernels, ML-for-compilers, the long tail. Interactive: an end-to-end scoreboard you build pass by pass.

Part VII · Synthesis

How to use this

Read it in order. Every lesson opens on the exact wall the previous one hit and derives the next idea as the forced fix. Skipping breaks the chain.
Touch every knob. Each lesson has one widget that turns its core trade into a slider with a configuration that breaks. Drive it until you've seen the failure, not just the happy path.
Open the code. The IR dumps, FX graphs, and generated Triton are real — read them, don't skim them.
Follow the cross-links. This track is a hub. Where another series goes deeper — kernels, serving, the training stack — it links out instead of repeating.

Companion tracks

This track is the "how the kernels get generated" sibling of the hand-written-kernel tracks. It links into what it overlaps rather than duplicating it: kernels → GPU Kernels and Triton; the training stack & its memory/parallelism math → CS336; serving & dynamic-shape inference → vLLM and SGLang; compression & precision → Knowledge Distillation.