Compilers and Kernels

A framework program is still several contracts away from executable device code. The compiler must preserve the program's meaning for every input covered by its guards, while choosing layouts, fusion boundaries, tile shapes, library calls, and target instructions. The kernel must then honor the promised shape, stride, dtype, layout, aliasing, mutation, numerical, and device target semantics. A fast result that violates any one of those conditions is not an optimization; it is a different program.

This chapter follows the execution path introduced in Chapter 63. It begins with the cost models that guide fusion and tiling, uses FlashAttention to show how an algorithm can be redesigned around a memory hierarchy, and then traces a graph through intermediate representations to target code. The final sections separate several meanings of portability and show why generated kernels need a stronger evaluator than ordinary application code. Quantized arithmetic and its specialized matrix kernels remain the focus of Chapter 34.

A kernel implements a bounded contract

A logical operator says what a result means. A kernel is one implementation of that operator for a declared domain. On a GPU, a launch creates a grid of thread blocks; the device schedules blocks onto streaming multiprocessors, and groups of lanes execute together as warps or wavefronts. Threads in one block can cooperate through shared memory and synchronization. Registers belong to individual threads, while global device memory is visible across the grid (Nickolls et al. 2008; NVIDIA 2026).

The launch domain matters as much as the kernel body. A contiguous matrix and a transposed view can have the same logical shape but different strides. Two arguments may alias the same storage. A reduction may promise a particular accumulation dtype or tie-breaking rule. A random operator advances state, and an in-place operator changes storage visible elsewhere. A compiler may specialize a kernel for these facts, but it must either guard them at runtime or retain a correct general path.

That gives five distinct responsibilities:

Performance work begins only after those boundaries are explicit.

Roofline is a bound, not a stopwatch

The roofline model relates arithmetic work to traffic across a chosen memory boundary (Williams et al. 2009). For one kernel, define

$$I=F/Q,$$

where:

• $F$ is the chosen count of floating-point operations for the kernel;
• $Q$ is the number of bytes transferred across the measured boundary, such as high-bandwidth memory (HBM) to the chip;
• $I$ is arithmetic intensity in floating-point operations per byte.

Using a compute ceiling and a bandwidth ceiling for the same target and precision gives the following bound. Here the two terms represent the compute and traffic ceilings, respectively:

$$P_{\text{bound}}=\min\left(P_{\max}, B_{\max}I\right),$$

and therefore the algorithmic lower bound

$$T\geq \max\left(\frac{F}{P_{\max}},\frac{Q}{B_{\max}}\right).$$

where:

• $P_{\text{bound}}$ is the roofline upper bound on throughput;
• $P_{\max}$ is the relevant compute ceiling in floating-point operations per second, using the same dtype and instruction path as $F$;
• $B_{\max}$ is the achieved bandwidth ceiling in bytes per second for the relevant access pattern, preferably measured rather than copied from a data sheet;
• $T$ is execution time in seconds.

Here, the ridge intensity is $I_{\text{ridge}}=P_{\max}/B_{\max}$. Below it, reducing traffic can raise the bound; above it, reducing arithmetic or using a faster compute path can. This is a classification and an optimistic limit, not a runtime prediction. Launch overhead, dependency latency, insufficient parallelism, instruction mix, cache behavior, uncoalesced access, bank conflicts, synchronization, and occupancy can all produce a lower ceiling. Observed performance should therefore use achieved bandwidth and achieved compute rates from a profiler, not just nominal peaks.

Transformer profiles illustrate why the distinction matters. Ivanov et al. found both tensor contractions and memory-bound operators to be material parts of BERT runtime; removing data movement accelerated the measured workload, but the result did not imply that every transformer or every operator was bandwidth-bound (Ivanov et al. 2021).

The runnable counts one deliberately simple case. Here it evaluates $y=\max(ax+b,0)$ over an FP32 vector. The unfused version writes two full intermediates; the fused version keeps the scalar intermediates inside one kernel. The calculation assumes scalar $a$ and $b$, a full HBM read and write for every materialized vector, and no useful cross-kernel cache retention. It counts minimum algorithmic traffic, not measured device traffic or time.

elements = 16 * 1024 * 1024
bytes_per_element = 4  # FP32
operations_per_element = 3  # multiply, add, maximum

operations = elements * operations_per_element
unfused_bytes = 3 * 2 * elements * bytes_per_element
fused_bytes = 2 * elements * bytes_per_element

for name, traffic in [("unfused", unfused_bytes), ("fused", fused_bytes)]:
    intensity = operations / traffic
    print(f"{name:7s}: {traffic / 1e6:5.1f} MB minimum traffic, "
          f"{intensity:.3f} operations/byte")

print(f"same arithmetic: {operations} operations")
print(f"traffic reduction: {unfused_bytes / fused_bytes:.1f}x")

Fusion is a legality and resource decision

Fusion combines producer-consumer operations so an intermediate need not be materialized in global memory. It can also remove launches and expose algebraic simplifications. These gains are real only when the combined kernel remains legal and fits the target well.

A compiler first has to prove that fusion preserves dependencies and observable effects. Mutation and alias relationships can make an apparently dead intermediate visible. Random-number state, host callbacks, collectives, atomics, and exceptions constrain reordering. Reductions need special care because a new reduction order can change rounding, NaN propagation, and deterministic behavior. Even pointwise fusion can introduce fused multiply-add instructions or reassociation, so mathematical equivalence does not imply bitwise identity.

The cost decision comes next. Fusion lengthens live ranges and can increase register pressure or shared-memory demand. If values spill to local or global memory, traffic returns in a more expensive form. A larger kernel can lower occupancy, reduce scheduling freedom, duplicate work, or prevent a call to a highly tuned library. This is why over-fusion is possible. A useful compiler estimates bytes removed, launch count, parallelism, and resource use, then profiles uncertain choices rather than treating fusion as an unconditional rule (Ansel et al. 2024).

Tiling maps work onto finite resources

Tiling chooses a small block of a larger computation and reuses its operands near the processors. Here, for matrix multiplication, a thread block may load tiles of $A$ and $B$ from global memory into shared memory, synchronize, perform many multiply-accumulates into registers, and repeat along the reduction dimension. Adjacent lanes should make coalesced global-memory accesses so the hardware can serve them with few transactions.

Each tile shape is a coupled choice. Larger tiles create more reuse but consume more registers and shared memory. More resident thread blocks can hide latency, yet maximum occupancy is not itself the goal: a lower-occupancy kernel can still win if each block does more useful work. Bank conflicts, synchronization frequency, tail masks, tensor-unit alignment, and the number of pipeline stages also change the result. When shapes or architectures vary, compilers commonly use heuristics to narrow the candidates and autotuning to benchmark a finite set of schedules (Chen et al. 2018; Tillet et al. 2019; Triton Project 2026).

Autotuning results are conditional artifacts. A winning tile for one shape, dtype, stride pattern, device, driver, or library version may lose for another. The cache key must contain every fact on which the choice depends, and a kernel that mutates its inputs must restore them between trials.

FlashAttention changes the algorithm around IO

Scaled dot-product attention for one head can be written as follows; here the symbols are defined immediately below:

$$S=QK^\mathsf{T}/\sqrt{d}+C, \qquad P_{ij}=\frac{\exp(S_{ij}-m_i)}{\ell_i}, \qquad O=PV,$$

with

$$m_i=\max_{1\leq j\leq N}S_{ij}, \qquad \ell_i=\sum_{j=1}^{N}\exp(S_{ij}-m_i).$$

where:

• $Q,K\in\mathbb{R}^{N\times d}$ are the query and key matrices;
• $V\in\mathbb{R}^{N\times d_v}$ is the value matrix;
• $N$ is sequence length, $d$ is query/key head width, and $d_v$ is value width;
• $C\in(\mathbb{R}\cup\{-\infty\})^{N\times N}$ is an additive bias or mask, where $-\infty$ disallows a key;
• $S$ is the scaled score matrix, $m_i$ is the maximum of query row $i$, and $\ell_i>0$ is its normalizer;
• $P$ is the row-wise softmax matrix, $O$ is the attention output, and $\mathsf{T}$ denotes transpose;
• every query row is assumed to have at least one allowed key.

A direct implementation may materialize the complete $N\times N$ score and probability matrices in HBM. FlashAttention instead processes a query row or query tile against a sequence of key-value blocks, retaining only tiles and running statistics in fast memory (Milakov and Gimelshein 2018; Dao et al. 2022). Here, for one fixed query row $i$, initialize $m_0=-\infty$, $\ell_0=0$, and $o_0=0\in\mathbb{R}^{d_v}$. For key-value block $\mathcal{B}_t$,

$$\begin{aligned} b_t &= \max_{j\in\mathcal{B}_t} S_{ij},\\ m_t &= \max(m_{t-1},b_t),\\ \alpha_t &= \exp(m_{t-1}-m_t),\\ \ell_t &= \alpha_t\ell_{t-1} +\sum_{j\in\mathcal{B}_t}\exp(S_{ij}-m_t),\\ o_t &= \alpha_t o_{t-1} +\sum_{j\in\mathcal{B}_t}\exp(S_{ij}-m_t)v_j. \end{aligned}$$

After $T$ blocks,

$$O_i=o_T/\ell_T.$$

where:

• $\mathcal{B}_1,\ldots,\mathcal{B}_T$ partition the allowed key indices and $t$ is the block index;
• $b_t$ is the maximum score in the current key-value block;
• $m_t$ is the running maximum and $\alpha_t$ rescales earlier contributions when that maximum changes;
• $\ell_t$ is the running normalizer;
• $o_t\in\mathbb{R}^{d_v}$ is the running value-weighted numerator;
• $v_j$ is row $j$ of $V$, and $O_i$ is output row $i$.

In exact arithmetic, the recurrence maintains the same numerator and denominator as the full softmax. “Exact attention” refers to that mathematical function, not bitwise identity: tiling changes floating-point evaluation order. The algorithm still performs quadratic attention arithmetic. Its gain is that it does not materialize the complete quadratic intermediates in HBM, and its backward pass can recompute local score and probability tiles instead of saving them.

In the two-level memory model used by the original paper, the forward algorithm requires

$$H_{\text{FA}}=\Theta\left(\frac{N^2d^2}{M}\right)$$

slow-memory word transfers for $d\leq M\leq Nd$, where $H_{\text{FA}}$ is the transfer count and $M$ is fast-memory capacity in scalar words. This is an IO statement, not a FLOP or byte count; bytes also depend on the stored dtype. Later work proved a matching pointwise lower bound, within constant factors, for $M\geq d^2$ and found a different regime below that threshold (Saha and Ye 2024). It is therefore incorrect to call the original algorithm optimal for every possible SRAM size.

FlashAttention-2 improved work partitioning and reduced non-matrix operations on the A100 configurations reported in its paper (Dao 2024). FlashAttention-3 used Hopper-specific asynchronous copy and matrix instructions in its reported H100 results (Shah et al. 2024). These are useful examples of a general rule: algorithmic IO savings can transfer across targets, while the best schedule often remains target-specific.

Lowering preserves meaning while adding decisions

Compilers use multiple intermediate representations because no single level is ideal for every decision. A high-level tensor IR can still express broadcasting, reductions, effects, and symbolic dimensions from the captured graph. A lower kernel IR can express tiles, address calculations, memory spaces, barriers, and target instructions. Progressive lowering makes each new choice explicit while retaining the information needed to check it. Halide established the influential separation between an algorithm and its schedule; TVM added tensor-graph optimization and measurement-guided schedule search; MLIR provides reusable infrastructure for defining and composing such IR levels (Ragan-Kelley et al. 2013; Chen et al. 2018; Lattner et al. 2021).

The diagram is deliberately generic. A compiler may preserve a high-level operation as a library call rather than generate a kernel. Matrix multiplication, convolution, and collectives often take this path because a vendor or framework library already contains many target-specific algorithms. Other regions lower to a kernel IR, where fusion, layout, and resource choices become explicit.

PyTorch's guarded path

For torch.compile, TorchDynamo captures guarded FX graph segments from Python execution. Training adds AOTAutograd, which captures and partitions forward and backward graphs and applies decompositions and functionalization. TorchInductor then schedules loop-level IR and can emit Triton GPU kernels, C++ CPU kernels, template kernels, or calls to external operators (Ansel et al. 2024; PyTorch Contributors 2026). A graph break returns unsupported work to eager execution. A guard failure can select another compile cache entry, trigger recompilation, or force a more general path.

Dynamic shapes do not remove specialization. Symbolic dimensions and guards allow some compiled regions to cover many inputs, but values, ranks, strides, and target features may still require distinct code. Compile latency, guard failure, graph breaks, and cache growth are therefore part of the production cost model, not frontend trivia.

JAX and OpenXLA

JAX traces a specialized Python function into jaxpr. Lowering converts the program through StableHLO into XLA's optimization and code-generation pipeline. XLA performs target-dependent fusion, layout assignment, buffer planning, collective handling, and kernel or library selection before a target executable is loaded through the runtime (Frostig et al. 2018; OpenXLA Project 2026; OpenXLA Project 2026).

StableHLO defines portable tensor semantics and a compatibility window for portable artifacts. It does not freeze a physical layout, schedule, numerical accuracy across every consumer, or performance. Those decisions belong to the consumer and target. JAX also specializes on facts such as shape, dtype, and static arguments, so its compile cache needs the same operational attention as any other staged system.

Layout and autotuning are part of compilation

Logical shape does not determine physical layout. A compiler may choose which dimension is contiguous, add padding or tiling, place buffers in different memory spaces, or distribute tile elements across lanes and registers. Some transposes are metadata changes; others move the entire tensor. Library calls and matrix instructions impose layout constraints, and converting between two otherwise valid layouts can erase the gain from a fast kernel.

Schedule selection has the same conditional nature. Triton raises the programmer's unit of work from a scalar thread to a blocked program, while its compiler handles tasks such as coalescing, vectorization, shared-memory allocation, and instruction selection (Tillet et al. 2019; Triton Project 2026). The author still chooses the grid, masks, tile sizes, number of warps or stages, and sometimes target-specific mechanisms. TVM-style learned cost models and empirical autotuning search over choices; they do not make one schedule universally optimal.

Portability is not one property

Claims about a “portable kernel” need a named boundary:

PTX is an NVIDIA virtual instruction set, not a universal GPU IR. CUDA device code may be shipped as PTX, target-specific cubin, or a fat binary containing several images. PTX can be assembled ahead of time or JIT-compiled by the driver; cubin compatibility is narrower (NVIDIA 2026). An AMD backend may instead lower through AMD LLVM to an HSACO artifact, while a TPU backend emits its own executable. These branches can share high-level tensor semantics without sharing a final schedule or binary.

The same distinction applies to Triton. Its blocked programming model can reduce source duplication across supported targets, but backend coverage, available instructions, layout choices, and tuned configurations still differ. A target-specific schedule is often needed for performance portability, and that schedule must be revalidated as hardware and compilers change.

The ecosystem is part of the target

CUDA began as a programming and execution model, but applications depend on a larger contract: the driver and runtime; compiler and artifact formats; math and communication libraries; profiler and debugger behavior; framework integration; and the accumulated tests and operational knowledge around them (Nickolls et al. 2008; Chetlur et al. 2014). Alternative targets have corresponding stacks. Comparing only a kernel language or a peak arithmetic rate therefore misses much of the switching cost.

Libraries are especially important. Their contracts include supported dtypes and layouts, workspace limits, algorithm selection, determinism, stream semantics, collective ordering, ABI compatibility, and target generations. Moving a workload means finding equivalent behavior for every relied-upon contract, then validating numerical results and performance on the actual shape distribution. Source translation can reduce labor, but it does not discharge that verification.

This is also a research constraint. The “hardware lottery” describes how ideas that map well to available hardware and software receive cheaper, faster experiments (Hooker 2021). A new architecture may look weak because it lacks a fused scan, sparse operation, or layout-aware compiler path, not because its mathematical idea is weak. Kernel availability therefore affects which model designs can be evaluated credibly.

Generated kernels need a hostile evaluator

KernelBench made generated-kernel evaluation concrete with 250 PyTorch workloads and a metric that requires both correctness and speed (Ouyang et al. 2025). For task set $\mathcal{D}$, one version of that metric can be written

$$\operatorname{fast}_p = \frac{1}{|\mathcal{D}|} \sum_{i\in\mathcal{D}} \mathbf{1}\!\left[ \operatorname{correct}(k_i) \land t(k_i)<\frac{t(b_i)}{p} \right].$$

where:

• $\mathcal{D}$ is the benchmark task set and $|\mathcal{D}|$ its size;
• $k_i$ is the generated kernel for task $i$ and $b_i$ is its baseline;
• $t(\cdot)$ is synchronized execution time under the benchmark protocol;
• $p>0$ is the required speedup threshold;
• $\operatorname{correct}(k_i)$ is the benchmark's semantic verdict;
• $\mathbf{1}[\cdot]$ is one when its condition holds and zero otherwise.

Execution and profiling feedback can improve results, and multi-turn training has shown gains on fixed benchmark splits (Baronio et al. 2025). Those results do not establish general kernel reliability. A kernel can pass a few shapes while writing only part of its output, rely on narrow input distributions, mutate an input, call the reference path, or exploit allocator and synchronization state. Follow-up audits have demonstrated that broader hidden distributions and stricter baselines can reverse apparent wins (Zhang et al. 2026). This is the same evaluator problem examined in Chapter 27, now with asynchronous hardware and mutable memory added.

A production evaluator should cover at least:

1. Contract surface. Freeze supported ranks, shapes, dtype and math mode, device architecture, alignment, layout, aliasing, mutation, workspace, stream, error, and determinism behavior.
2. Hidden semantics. Use fresh random inputs and outputs, boundary shapes, zero and odd sizes, broadcast cases, noncontiguous strides, misalignment, cancellation, ties, signed zeros, NaN and infinity where supported, and independent tolerances derived from dtype and reduction length.
3. Memory and concurrency. Detect out-of-bounds access, uninitialized data, incomplete writes, races, bad synchronization, illegal aliasing, and failures under concurrent streams, graph replay, and varied allocator histories.
4. Harness integrity. Isolate the candidate; prevent reference calls, monkey-patching, hidden fallback, network or file access, and reuse of reference allocation state. Perturb every input so the output must depend on it.
5. Measurement. Separate compile time, autotuning, cold-start latency, kernel-only time, and end-to-end time. Use warm-up, device synchronization, randomized paired order, many repetitions, and a median plus useful quantiles or confidence intervals.
6. Representative comparison. Match dtype, fast-math settings, layout, workspace, graph mode, and hardware. Compare with a production baseline, not only an eager reference, over the real workload distribution.

Operating the compiler-kernel boundary

A reliable release process treats the generated artifact and its assumptions as one unit:

1. Record the graph or operator schema, guards, effect and alias model, numerical tolerance, and fallback.
2. Pin the target architecture, driver, runtime, math libraries, compiler, toolchain version, flags, and autotuning state.
3. Build a semantic gate covering eager-versus-compiled outputs, gradients when relevant, edge shapes, strides, dtypes, mutation, and concurrency.
4. Build a compiler gate covering graph breaks, guard failures, recompilations, compile cache growth, sanitizer findings, and target-code generation.
5. Build a performance gate over the representative workload distribution, measuring launch count, achieved bandwidth or compute, occupancy, registers, shared memory, spills, workspace, and peak memory.
6. Report cold-start compilation and autotuning separately from warm execution.
7. Canary the artifact with automatic rollback to the known fallback on numerical, compiler, or latency regression.
8. Re-run the gates whenever the graph, target, driver, library, compiler, or workload distribution changes.

Lower-layer constraint

The compiler cannot recover information that the framework discarded, prove aliasing facts the program did not expose, or make an unsupported target instruction exist. Conversely, a kernel can be locally fast and still lose end-to-end because its layout forces a transpose, its workspace raises memory pressure, its compile cache fragments on dynamic shapes, or its launch breaks a larger fusion region. Compiler and kernel performance must therefore be measured in the graph and service that consume them, against the bandwidth, memory, and interconnect limits from Chapter 62.

The next chapter, Chapter 65, moves one level upward in operational time: from producing executable work to keeping a long-running training system supplied, observable, and recoverable.

Further reading