1.2 Overview of AI Software Stacks

Overview of AI Software Stacks #

Modern AI systems sit on top of a deep stack: Python model code, graph capture, compiler intermediate representations (IR), kernel libraries, device runtimes, and finally, accelerator machine code. When we talk about teaching LLMs to write system kernels for AI accelerators, we need to understand where those kernels live in this stack, what existing compilers already do well, and where AI-generated kernels can make things more efficient.

This section builds a mental model of that stack. We will use PyTorch, TorchDynamo, TorchInductor, Triton, XLA, HLO, CUDA, CUTLASS, cuBLAS, FBGEMM, and related systems as concrete examples. The goal is not to memorize every compiler component, but to understand the interfaces between them: where a model becomes a computation graph, where a graph becomes a compiler IR, and where hand-written or agent-written kernels become necessary.

Custom kernels for longer tail cases #

Most ML compilers cover the common cases. For example, PyTorch’s torch.compile() stack can capture many models, fuse common tensor operations, and generate efficient GPU kernels, XLA can compile large graph computations, and TVM can auto-schedule dense kernels. Vendor libraries like cuBLAS, cuDNN, etc., provide extremely optimized primitives which act as reusable building blocks that compilers may reuse.

However, ML workloads have a “long tail”. This refers to the large number of uncommon, irregular, model-specific, or hardware-specific patterns that are not handled optimally by general-purpose compiler. Examples include tensor layouts, dynamic control flow, custom normalization variants, quantization/dequantization patterns, MoE routing, gather/scatter-heavy operators, dynamic shapes, heterogeneous fusion patterns, and accelerator-specific memory movement.

ML compilers are very good at common, regular patterns. For instance, a compiler may do very well on the common transformer path:

x = x + residual 
x = layernorm(x) 
y = x @ W 
y = gelu(y)

but struggles when the computation becomes:

tokens = route_tokens_to_experts(x, topk_indices) 
expert_outputs = grouped_expert_gemm(tokens, expert_weights) 
out = combine_expert_outputs(expert_outputs, routing_weights)

The first pattern is mostly regular dense computation, whereas the second involves dynamic routing and irregular indexing, which is exactly the kind of regime where specialized kernels remain important. Hence, it becomes important to understand torch.compile() to identify what it can already do well, which we will look at next.

PyTorch eager mode #

Before discussing torch.compile(), it is important to understand PyTorch eager mode. In normal PyTorch eager execution, when you write: y = model(x), Python executes the model line-by-line. Each operation dispatches immediately. For example,

y = torch.matmul(x, W)
z = torch.relu(y + b)

In eager mode, this typically means: torch.matmul dispatches to an implementation (eg., from cuBLAS), and so do y + b and torch.relu which get dispatch to their own elementwise kernels. There is no ahead-of-time graph construction. Note that PyTorch still uses highly optimized libraries such as cuBLAS, depending on the operator and backend; the key point is that eager mode executes operations one at a time.

A simplified eager-mode stack lowering looks like:

PyTorch Python code 
        ↓ 
ATen operator dispatch 
        ↓ 
Backend kernel implementation 
        ├─ CUDA kernel 
        ├─ cuBLAS / cuDNN call 
        ├─ CPU kernel └─ other backend-specific implementation 
       ...
        ↓ 
Device execution

An ATen operator is one of PyTorch’s core tensor operations – they are the fundamental operations that PyTorch dispatches internally. For example, when you write y = torch.relu(x + b), PyTorch internally sees something closer to: aten::add, aten::relu. The Aten operator is the logical operation, and is mapped to a backend kernel, which is the actual implementation for a device. To “write a kernel for an ATen operator” means implementing the backend code that executes the operator on a particular device (e.g., a CPU implementation, a CUDA implementation, a quantized CPU implementation, a TRN implementation, etc.).

Default torch.compile() stack #

When we call compiled_model = torch.compile(model), PyTorch enters a different execution path. By default, torch.compile() uses TorchInductor as its backend. A simplified torch.compile() stack lowering would look like:

PyTorch Python model
    ↓
TorchDynamo
    ↓
FX graph (usually ATen-level)
        ↓
AOTAutograd (mainly for training)
        ↓
TorchInductor
   ├─ generated Triton kernels
   │     ↓
   │   Triton IR
   │     ↓
   │   LLVM IR
   │     ↓
   │    PTX
   │     ↓
   │   SASS (executed on GPU)
   │
   ├─ external library calls
   │     ├─ cuBLAS / cuBLASLt
   │     ├─ cuDNN
   │     └─ FlashAttention-style kernels
   │
   ├─ ATen fallback kernels 
   └─ custom backend kernels

In what follows, we describe each level of this stack briefly.

  1. TorchDynamo: TorchDynamo captures Python execution into graphs. It does not replace Python with a completely static graph compiler at the outset. Instead, Python still runs line-by-line (as in eager mode), but TorchDynamo observes and intercepts PyTorch operations. The regions that can be represented as graphs are captured. If Dynamo encounters unsupported Python behavior, data-dependent control-flow, or operations it cannot safely capture, it may introduce a graph break. A graph break means that compilation stops for that region, and execution may fall back to eager mode before resuming capture later.

  2. FX graph: The captured graph is represented as an FX graph. FX is a Python-friendly graph representation used widely in PyTorch tooling. An FX graph is often composed of ATen operators. For example, a Python expression like y = torch.relu(x + b) may become an FX graph with nodes corresponding to ATen-level operations such as aten::add and aten::relu. FX is higher-level and more readable than traditional compiler IRs. It is still close to PyTorch semantics and is therefore easier to inspect, transform, and debug. Contrary to FX graphs, TorchScript was the older PyTorch graph compiler, which was not as readable.

  3. AOTAutograd: This is mostly relevant for training. In eager PyTorch, autograd works dynamically – in forward pass, PyTorch builds a tape of operations and saves tensors needed for backward. This is flexible, but not ideal for a compiler. AOTAutograd rewrites autograd into explicit forward and backward graphs. Conceptually, it transforms training into: ForwardGraph(inputs, params) --> outputs, saved_tensors and BackwardGraph(grad_outputs, saved_tensors, params) --> gradients. This gives the compiler explicit graphs for both forward and backward computation. TorchInductor (discussed next) can then optimize and compile these graphs.

  4. TorchInductor: TorchInductor is the default backend for torch.compile(). It takes the FX graph and decides how to lower it. Its main responsibilities include analyzing the FX/ATen graph, identifying fusion opportunities, generating code for fused regions, calling external libraries for large primitives, and producing executable code. For many elementwise and reduction-heavy subgraphs, Inductor generates Triton kernels. For large GEMMs, convolutions, and other vendor-optimized primitives, it often calls external libraries like cuBLAS or cuDNN.

    TorchInductor can be especially effective in fusing regular, compiler-friendly patterns (e.g., pointwise and elementwise fusion). For instance, the operation out = torch.relu(x * a + b) in eager mode might be multiple kernels like multiply kernel, add kernel and relu kernel, but TorchInductor can often fuse this into one Triton kernel: out[i] = relu(x[i] * a + b) – which improves performance by reducing kernel launches and avoiding intermediate writes to GPU memory. Similarly, Inductor can fuse broadcasted add and multiply operations into one kernel if the memory access pattern is regular. It can also remove or simplify redundant views, reshapes, and transposes. These are not always fusion in the strictest sense, but they can eliminate unnecessary work and enable better fusion later.

    However, TorchInductor has limits and often struggles when fusion crosses the boundary of highly-optimized library primitives, irregular execution patterns, or hardware-specific schedules. Some examples are as follows:

    • Arbitrary epilogue logic a computation such as y = x @ W; out = silu(y + b) may be lowered into a cuBLAS GEMM followed by a separate Triton kernel for bias and activation; however, because cuBLAS is an opaque black box that TorchInductor cannot freely modify to inject arbitrary epilogue logic, fused GEMM epilogues often require a custom kernel.

    • Online fusion Similarly, Flash-Attention style full attention fusion requires avoiding materialization of the attention matrix through careful tiling, online softmax, shared-memory use, and warp-level scheduling, which are usually not discovered automatically.

    • Dynamic control flow includes Python branches or loops that depend on runtime values:

         if x.shape[1] > 1024: 
       y = long_context_path(x) 
   else: 
       y = short_context_path(x)

      Dynamo may handle some shape-dependent branching with guards and recompilation, but arbitrary data-dependent control can cause graph breaks.

    • Mixture-of-Experts routing: MoE workloads involve top-k routing, token dispatch to experts, grouped GEMMs, gathering outputs back, and possible cross-device communication. This is difficult because it combines irregular indexing, variable expert loads, dynamic shapes, many small or grouped GEMMs, and often distributed communication. MoE remains a canonical example where custom kernels remain important.

    • Quantization and heterogeneous fusion: Quantized inference adds further complexity through mixed dtypes, scale handling, dequantization/requantization, packed layouts, and formats such as FP8, INT8, INT4, MX, or blockwise scaling.

    Due to these limitations of generic ML compilers, specialized kernels are often required.

  5. Triton: Triton is a Python-based language and compiler for writing GPU kernels. In the torch.compile() stack, TorchInductor relies heavily on Triton as a code generation target for many fused kernels. When a user writes Triton kernels directly, they essentially start at the Triton-level and skip TorchDynamo, FX, AOTAutograd, and TorchInductor’s fusion analysis. Triton is often a good fit for custom RMSNorm, fused elementwise chains, packing/unpacking kernels, layout transforms, etc. However, for peak-performance GEMM, for example, one may still need CUDA, CUTLASS, cuBLASLt, or specialized libraries.

  6. LLVM, CUDA, PTX, SASS: At the bottom of the NVIDIA GPU stack, we encounter lower-level compilation artifacts. LLVM IR is a low-level compiler intermediate-representation. It is target-independent compared to PTX (described later). Many compilers lower into LLVM IR because LLVM provides mature optimization passes and code generation infrastructure (e.g., dead code elimination, instruction simplification, eliminating duplicate computations, loop optimizations, etc.). On the other hand, a CUDA kernel is a GPU function written in CUDA C++, which is typically compiled using nvcc and lowered through NVIDIA’s CUDA compilation toolchain toward PTX and/or architecture-specific GPU code.

    PTX is NVIDIA’s virtual assembly language – it is GPU-specific, but not yet the final machine code for a particular GPU architecture. PTX can be thought of as a portable GPU assembly representation that can be compiled for different NVIDIA generations, such as sm_80, sm_90, or later targets. It is further lowered into SASS (Streaming Assembler), which is the actual GPU machine code. This is the final architecture-specific machine code executed by the GPU’s streaming multiprocessors. In practice, PTX may be compiled to SASS ahead-of-time by NVIDIA’s tools or just-in-time by the NVIDIA driver when the program runs on a particular GPU.

Vendor libraries: cuBLAS, cuDNN, cuBLASLt, and CUTLASS #

As one would obviously expect, not all kernels are generated by TorchInductor or go through Triton. Many operations still fall back to highly-optimized vendor libraries, some of which are described below:

  1. cuBLAS: This is NVIDIA’s production BLAS library, providing highly optimized dense linear algebra routines, especially GEMM. When PyTorch or Inductor needs to do a large matrix multiplication, it often calls cuBLAS under the hood, rather than generating a GEMM kernel from scratch. cuBLAS kernels are shipped as precompiled binaries, and may include native SASS for specific GPU architectures and PTX for compatibility. At runtime, cuBLAS chooses an appropriate kernel variant based on shapes, datatypes, layouts, and hardware. This is why TorchInductor cannot easily fuse arbitrary code inside a cuBLAS GEMM, as it is an opaque library call boundary.

  2. cuBLASLt: Contrary to above, it is a more flexible matmul API, and provides more control over layouts, algorithm selection, and some epilogue fusion options. If one wants GEMM plus bias or activation fusion while still using an NVIDIA library, cuBLASLt may be more appropriate than classical cuBLAS.

  3. cuDNN: cuDNN is NVIDIA’s deep neural network primitives library. It is commonly used for convolutions, pooling, normalization, and related DNN operations. In a compiler stack, cuDNN sits at the same general layer as cuBLAS: an optimized library backend called by higher-level frameworks.

  4. CUTLASS: CUTLASS is not a runtime library in quite the same sense as cuBLAS. It is an open-source C++ template library for building high-performance GEMM, convolution, and related kernels. It is used when one wants high performance GEMM-like kernels with custom epilogues, layouts, datatypes, or fusion behavior. Modern CUTLASS, especially CUTLASS 3.x, uses CuTe as a core internal abstraction for describing tensor layouts, tiling, and thread/data partitioning. CuTe is especially useful because high-performance GEMM and attention kernels are largely about layout and schedule. The math is simple; the hard part is mapping that math efficiently onto registers, shared memory, warps, and tensor cores. Since CuTe is a C++ template DSL, it is not lowered like Triton. CuTe C++ templates are lowered to CUTLASS CUDA C++ kernel code, which is subsequently lowered further by nvcc to PTX and SASS.

Other libraries: FBGEMM, FBGEMM_GPU, DeepGEMM #

These libraries live at the kernel/library layer, similar to cuBLAS, cuDNN, CUTLASS, etc. FBGEMM (from Facebook), was historically associated with CPU inference, quantized GEMM, and recommendation workloads. On the PyTorch side, FBGEMM can act as an ATen backend kernel provider. FBGEMM_GPU also has GPU components. Similarly, DeepGEMM (from DeepSeek-AI) also has specialized GEMM kernels for a variety of modern transformer shapes.

TVM Relay and TIR #

Apache TVM (Tensor Virtual Machine) is another compiler stack for ML. The reason we should care, even if we already discussed the PyTorch stack, is that TVM gives useful vocabulary for explaining why kernel-writing agents are needed. A graph compiler can decide that some ops should be fused, but then someone or something still has to produce a good low-level schedule: tiling, vectorization, memory layout, shared-memory staging, thread mapping, and so on. TVM’s Relay-to-TIR split makes that distinction very concrete. In addition to being useful for compiler research, TVM is still used in edge and embedded deployment, along with custom/emerging accelerators. TVM has multiple IRs, two of which are useful to compare with PyTorch and XLA (discussed later).

  1. TVM Relay is a higher-level graph IR. It represents tensor computations such as matmul, relu, conv, reshape, and is used for graph-level optimization, operator fusion decisions, layout rewriting, constant folding, and other transformations.

  2. TensorIR (TIR): On the other hand, TIR is lower level. It represents how to compute tensor programs using loops, memory accesses, tiling, vectorization, and parallelization.

A simplified TVM lowering path would be:

Model graph --> Relay --> TIR --> Target code

For example, suppose the model does Y = relu(X @ W + b). At the Relay level, it may look like a graph of high-level tensor ops (specifying what to compute: dense matmul, add bias, apply ReLU):

def @main(%X, %W, %b) {
  %0 = nn.dense(%X, %W)
  %1 = add(%0, %b)
  %2 = nn.relu(%1)
  return %2
}

Relay can reason that these operators form a producer-consumer chain and decided to fuse them into one fused primitive fusion/computation region. After relay fusion, TVM may rewrite the above into something like:

def @main(%X, %W, %b) {
  %fused = fn (%p0, %p1, %p2, Primitive=1) {
    %0 = nn.dense(%p0, %p1)
    %1 = nn.bias_add(%0, %p2)
    %2 = nn.relu(%1)
    return %2
  }

  %fused(%X, %W, %b)
}

However, that decision alone does not specify how the fused computation should run on hardware. When further lowered to TIR, the same computation becomes explicit loop nests over the output matrix etc., as follows:

@T.prim_func
def fused_dense_bias_relu(X, W, b, Y):
    for i, j, k in T.grid(M, N, K):
        with T.block("matmul"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            Y[vi, vj] += X[vi, vk] * W[vj, vk]

    for i, j in T.grid(M, N):
        with T.block("bias_relu"):
            vi, vj = T.axis.remap("SS", [i, j])
            Y[vi, vj] = T.max(Y[vi, vj] + b[vj], 0.0)

This says how to compute it: loop over rows, columns, and reduction dimension. A compiler schedule might then transform the TIR into something more hardware-aware (e.g., tiling i/j/k mapping loops to CUDA blocks/threads, etc.).

TVM’s Relay-to-TIR split makes the distinction between graph-level and tiling-level abstractions explicitly concrete. In the PyTorch stack, the corresponding distinction is less obvious, but still exists. For instance, one can think the FX/ATen graph as analogous to TVM Relay, and TorchInductor/Triton kernels/Other library calls, as playing the role of TIR.

MLIR and dialects #

We saw above that the PyTorch and TVM stacks show the same basic pattern: a high-level model is progressively lowered through intermediate representations until it becomes executable kernels. MLIR (Multi-Level Intermediate Representation) provides a broader compiler framework for expressing this kind of multi-level lowering, where each abstraction level is represented by a different “dialect”.

For kernel-writing agents, MLIR is relevant because it makes the optimization hierarchy explicit: an agent may reason effectively at one level of abstraction, such as tensor graphs or loop schedules, but struggle at another, such as hardware-specific instruction selection or memory-space management. Dialects therefore provide both structure and degrees of freedom: they let us ask where an agent should interfere in the lowering stack, what information should be exposed at that level, etc.

An example of MLIR dialects for Y = relu(A @ B) is as follows:

  1. High-level tensor/graph dialect: At the top, the compiler may represent the computation as something which is still close to the model meaning:

    %0 = stablehlo.dot_general %A, %B
%1 = stablehlo.maximum %0, %zero
return %1

  2. Structured linear algebra dialect: Then it may lower into something like the linalg dialect:

    %0 = linalg.matmul ins(%A, %B : tensor<MxKxf32>, tensor<KxNxf32>)
              outs(%C : tensor<MxNxf32>)

%1 = linalg.generic ins(%0 : tensor<MxNxf32>)
                    outs(%Y : tensor<MxNxf32>) {
^bb0(%x: f32, %out: f32):
    %z = arith.constant 0.0 : f32
    %r = arith.maximumf %x, %z : f32
    linalg.yield %r : f32
}

    The structure is now clear: linalg.matmul = structured matrix multiplication and linalg.generic = elementwise ReLU. This is still not raw loops, but it is more schedule-friendly.

  3. Loop + memory dialects: Next, the compiler may lower to explicit loops and memory buffers using dialects:

    scf.for %i = 0 to %M {
scf.for %j = 0 to %N {
    %sum = arith.constant 0.0 : f32

    scf.for %k = 0 to %K {
    %a = memref.load %A[%i, %k]
    %b = memref.load %B[%k, %j]
    %prod = arith.mulf %a, %b
    %sum_next = arith.addf %sum, %prod
    }

    %zero = arith.constant 0.0 : f32
    %relu = arith.maximumf %sum, %zero
    memref.store %relu, %Y[%i, %j]
}
}

    At this level, the computation is described in terms of loops, loads, stores, scalar arithmetic, and memory buffers. Now, the compiler can reason about tiling, loop interchange, vectorization, memory reuse, and buffer allocation.

  4. GPU/LLVM-level lowering: Near the bottom, the compiler may lower into GPU or LLVM-level dialects:

    gpu.launch blocks(%bx, %by, %bz) threads(%tx, %ty, %tz) {
// compute tile of Y using GPU threads
// load A/B tiles
// multiply-accumulate
// apply ReLU
// store result
}

    Here, LLVM IR refers to a low-level, compiler-friendly, representation with explicit control flow, scalar arithmetic, memory operations, and function calls. Subsequently, for GPUs, it is further lowered toward PTX and SASS. The key motivation in MLIR is that each dialect exposes the right details for that stage.

XLA, HLO, StableHLO, and PyTorch/XLA #

So far, we have looked at two compiler views of AI workloads: The PyTorch torch.compile() stack shows how a Python model can be captured into an FX graph and lowered by TorchInductor into Triton kernels, vendor-library calls, or backend fallbacks. TVM shows a more explicit separation between graph-level optimization in Relay and low-level scheduling in TIR.

XLA introduces a third important view: whole-graph compilation for accelerators. This matters because not every AI accelerator is programmed by writing CUDA-style kernels. Google TPUs, for example, are commonly reached from PyTorch through PyTorch/XLA, which connects PyTorch models to the XLA compiler and XLA devices (TPUs). AWS Trainium similarly uses the AWS Neuron compiler stack, where the Neuron graph compiler can transform models from frameworks such as PyTorch or XLA HLO into optimized executable code. In these settings, the compiler interface is often not “write a kernel directly”, but rather “lower a model or graph into a tensor-level IR that the accelerator compiler can optimize.” This is where HLO and Stable HLO become important.

HLO is XLA’s tensor-level compiler representation, while StableHLO is a portable operation set intended to sit between ML frameworks and ML compilers. For kernel-writing agents, this broadens the question: the agent may not always emit Triton or CUDA; it may instead need to rewrite PyTorch code so that it lowers cleanly to HLO, expose fusion opportunities to the accelerator compiler, or reason about where a lower-level custom kernel is actually needed.

A simplified PyTorch/XLA path is:

PyTorch model 
    ↓ 
PyTorch/XLA lazy capture 
    ↓ 
XLA Op IR / lazy graph 
    ↓ 
HLO / StableHLO 
    ↓ 
XLA compiler 
    ↓ 
XLA runtime interface (often PJRT) 
    ↓ 
Device runtime 
    ↓ 
TPU / accelerator execution

We next briefly overview each step in this lowering stack. As an illustrative example, suppose we write:

def f(x, w, b):
    y = x @ w
    y = y + b
    y = torch.relu(y)
    return y

Lazy graph capture and XLA Op IR: On a normal CUDA tensor in eager mode, PyTorch may execute this as separate operations: matmul, add and ReLU. Moreover, TorchDynamo tracing in torch.compile() will run Python normally and capture graph regions as Python executes. On the other hand, PyTorch/XLA lazy capture will record the operations instead of executing them immediately. Only at a synchronization point, XLA will compile and execute the graph. Synchronization points include printing a tensor value, calling .item(), moving a tensor to CPU, explicit xm.mark_step(). Lazy capture exposes larger fusion opportunities, reducing the launch overhead. The frontend recording format of the lazy graph is also referred to as the XLA Op IR – it represents the pending tensor operations, and is not yet the final normalized compiler representation.

HLO and StableHLO: The next level, i.e., HLO, stands for High-level Optimizer. It is XLA’s main tensor-level compiler IR. HLO is lower-level than PyTorch code, but still much higher-level than LLVM IR, PTX or machine code. Conceptually, the HLO for the above example may look like:

%0 = dot(%x, %w)
%1 = broadcast(%b)
%2 = add(%0, %1)
%3 = maximum(%2, 0)
return %3

This is where XLA can run graph-level optimizations such as fusion, algebraic simplification, common subexpression elimination, layout assignment, buffer assignment, and backend-specific rewrites. For instance, XLA may decide:

matmul remains a target-specific dot/GEMM primitive
add + relu may be fused
layouts may be chosen to avoid unnecessary transposes
buffers may be reused to reduce memory usage

On the same note, StableHLO is a standardized, portable version of HLO. It exists because internal HLO can evolve over time, while frameworks and compiler backends need a stable interchange format. Once a model is lowered into Stable HLO, different accelerator compilers can consume that representation.

Note: Sometimes people say XLA IR loosely. It may refer to either the frontend lazy graph or the HLO itself.

XLA runtime interface and Device runtime: After XLA compiles an HLO program, the result is a device executable. Something then has to manage that executable: load it, cache it, bind input/output buffers, launch it, and collect results. This is done by the XLA runtime interface (often called PJRT: Portable JIT Runtime). OpenXLA describes PJRT as a uniform device API: frameworks call PJRT, and device-specific implementations sit behind that interface. In other words, the framework does not need to know all hardware details; it talks to a PJRT client/plugin for any accelerator. The device runtime is below PJRT, which is hardware-facing and actually knows how to operate the accelerator.

This separation is what lets frameworks and compilers target multiple devices through a common runtime interface, while still allowing each hardware vendor or backend to implement the low-level details differently. For kernel-writing agents, this matters because an agent optimizing an XLA-style accelerator stack may need to reason not only about HLO fusion, but also about runtime behavior: avoiding unnecessary sync points, avoiding repeated recompilation, keeping tensors resident on device, and respecting the execution model of PJRT/device runtimes.

Summary and Where Agents Fit #

Across every system in this section — PyTorch torch.compile(), TVM, MLIR, and XLA — the same pattern recurs: a model description is progressively lowered through a series of intermediate representations, from a high-level graph, to a tensor-level IR, to a schedule/kernel, to a virtual ISA, and finally to machine code driven by a runtime. The names differ per stack (FX vs. Relay vs. StableHLO at the graph level; Triton IR vs. TIR vs. linalg/scf at the schedule level), but the layering is the same. A useful unifying mental model is as follows:

Model source           PyTorch nn.Module / model.forward
       ↓
Graph capture          TorchDynamo FX graph · PyTorch/XLA lazy graph · (TorchScript, older)
       ↓
Tensor compiler IR     Inductor IR · HLO / StableHLO · TVM Relay · MLIR tensor/linalg
       ↓
Kernel / schedule IR   Triton IR · TVM TIR · MLIR scf/memref/gpu · CUDA C++ / CUTLASS / CuTe
       ↓
Low-level IR / ISA      LLVM IR · PTX
       ↓
Machine code           SASS (NVIDIA) · TPU executable · Trainium executable
       ↓
Runtime                CUDA driver · PJRT / XLA runtime · Neuron runtime · device runtime
       ↓
Hardware               GPU / TPU / Trainium / other AI accelerator

The central takeaway is a division of labor: compilers own the regular, common-case path (dense transformer blocks, pointwise/reduction fusion, standard GEMMs), while hand-written or agent-written kernels exist for the long tail — fused GEMM epilogues, FlashAttention-style online softmax, MoE routing, quantized/heterogeneous fusion, and accelerator-specific memory movement — where general-purpose fusion does not find a good schedule.

A kernel-writing agent therefore does not live at a single level. Given a slow region, the right action might be to:

  • leave it to TorchInductor,
  • rewrite PyTorch source to expose a fusion the compiler can already take,
  • emit a Triton kernel (elementwise/reduction/layout-bound work),
  • drop to CUDA / CUTLASS / cuBLASLt for GEMM epilogues, tensor-core scheduling, or attention,
  • select an existing specialized kernel (FlashAttention, FBGEMM, DeepGEMM),
  • or lower the graph to HLO/StableHLO for a TPU or Trainium backend.

A strong agentic kernel-optimization system should understand this stack well enough to choose the correct intervention point. Choosing correctly requires reasoning across the whole stack at once — model semantics, operator graphs, compiler IRs, memory layouts, hardware execution, and the runtime. This is precisely why post-training, reasoning, and agentic optimization matter for kernel generation. The goal is not to replace compilers but to complement them: compilers for the common case, tuned libraries for known-hard primitives, and LLM-driven synthesis or search for the long tail that still escapes optimal compilation.