AI Compilers

Writing optimal CUDA kernels by hand is powerful but doesn’t scale — every new operation, every shape change, every hardware generation requires re-optimization. AI compilers automate the tedious parts (tiling, fusion, scheduling) while letting experts focus on algorithmic innovation.

---

1. Why Compilers?

Approach	Pros	Cons
Hand-written CUDA	Maximum performance, full control	Weeks per kernel, hardware-specific, hard to maintain
cuBLAS/cuDNN	Near-optimal for standard ops	No custom fusion, limited to supported ops
AI Compiler	Custom fusion, auto-tuning, portable	80-95% of hand-tuned, compilation overhead

The 80/20 rule: Compilers handle the 80% of operations that are “standard optimizable” (element-wise, reduction, simple GEMM variants). The 20% that need custom algorithms (FlashAttention, custom routing) still require hand-written kernels.

---

2. Triton: Python GPU Kernels

Triton lets you write GPU kernels in Python with block-level abstractions — no shared memory management, no warp-level thinking, no bank conflict debugging.

How It Works

import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)  # block index
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < N

    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

# Launch
add_kernel[(N // 256,)](x, y, out, N, BLOCK=256)

What Triton Handles Automatically

Manual in CUDA	Automatic in Triton
Shared memory allocation	Compiler decides when to use SRAM
Bank conflict avoidance	Compiler inserts padding/swizzle
Memory coalescing	Compiler reorders accesses
Warp-level primitives	Block-level abstraction hides warps
Register allocation	Compiler manages registers
Tiling loop structure	You define block size, compiler tiles

Triton for Fused Kernels

Where Triton shines — custom fused operations:

@triton.jit
def fused_residual_layernorm(x_ptr, res_ptr, gamma_ptr, beta_ptr, out_ptr,
                              N, eps, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N

    x = tl.load(x_ptr + row * N + cols, mask=mask)
    res = tl.load(res_ptr + row * N + cols, mask=mask)
    gamma = tl.load(gamma_ptr + cols, mask=mask)
    beta = tl.load(beta_ptr + cols, mask=mask)

    # Fused: residual + layernorm
    h = x + res
    mean = tl.sum(h, axis=0) / N
    var = tl.sum((h - mean) * (h - mean), axis=0) / N
    norm = (h - mean) / tl.sqrt(var + eps)
    out = norm * gamma + beta

    tl.store(out_ptr + row * N + cols, out, mask=mask)

Performance

Triton typically achieves 80-95% of hand-tuned CUDA for standard patterns. For GEMM specifically, Triton’s auto-tuner can reach 90%+ of cuBLAS by searching tile sizes, number of pipeline stages, and warp configurations.

---

3. torch.compile (TorchInductor)

PyTorch’s built-in compiler — captures the Python computation graph and generates optimized code.

How It Works

@torch.compile(mode="reduce-overhead")
def forward(x, weight, bias):
    h = x @ weight + bias
    h = torch.nn.functional.gelu(h)
    h = torch.nn.functional.dropout(h, p=0.1)
    return h

Under the hood:

1. TorchDynamo: captures Python bytecode into a graph (handles control flow, dynamic shapes)
2. TorchInductor: compiles the graph into optimized Triton kernels (or C++/CUDA)
3. Auto-fusion: identifies element-wise operation chains and fuses them
4. CUDA Graphs: wraps the compiled graph for reduced launch overhead

Compilation Modes

Mode	Effect	Use case
default	Balanced optimization	General training
reduce-overhead	CUDA Graphs for minimal launch overhead	Inference, repetitive training steps
max-autotune	Search over multiple kernel implementations	Final deployment optimization

What torch.compile Handles Well

• ✅ Element-wise fusion (residual + activation + dropout)
• ✅ Pointwise operations with broadcasting
• ✅ Simple reductions (sum, mean, norm)
• ✅ Standard GEMM (delegates to cuBLAS)

What It Doesn’t Handle

• ❌ Custom attention patterns (use FlashAttention directly)
• ❌ Dynamic shapes that change every iteration (recompilation each time)
• ❌ Complex control flow dependent on tensor values
• ❌ Operations not in the supported op list

---

4. XLA (Accelerated Linear Algebra)

Google’s compiler, primarily used with JAX and TensorFlow. Compiles entire computation graphs into optimized HLO (High-Level Optimizer) IR.

Key Differences from TorchInductor

	TorchInductor	XLA
Primary framework	PyTorch	JAX, TensorFlow
Compilation unit	Individual ops / subgraphs	Entire program
Target hardware	NVIDIA GPUs	GPUs, TPUs, CPUs
Fusion strategy	Op-by-op	Whole-program optimization
Dynamic shapes	Supported (with recompilation)	Requires static shapes (JIT traces)

XLA’s whole-program optimization can find fusion opportunities that op-level compilers miss — but requires static shapes and JAX’s functional programming style.

---

5. TVM (Tensor Virtual Machine)

An open-source compiler framework for deploying ML models across diverse hardware:

Feature	What it does
Auto-scheduling	Automatically finds optimal loop tiling, parallelization, vectorization
Hardware targets	CUDA, Metal, Vulkan, CPU (LLVM), custom accelerators
Model import	ONNX, TensorFlow, PyTorch models
Use case	Edge deployment, cross-platform optimization

TVM is less commonly used for NVIDIA GPU training (where Triton/torch.compile dominate) but important for edge deployment where you need to compile models for diverse hardware.

---

6. When to Use What

Scenario	Best tool
Standard training loop	torch.compile (zero effort, good speedup)
Custom fused kernel	Triton (Python-level GPU programming)
Maximum GEMM performance	cuBLAS (still the gold standard)
Custom attention variant	Hand-written CUDA or Triton
JAX / TPU training	XLA
Edge deployment	TVM or TensorRT
Production inference	TensorRT-LLM (NVIDIA’s serving compiler)

The Typical AI Infra Stack

Application (Python)
  → torch.compile (auto-fuses element-wise ops)
    → Triton (custom fused kernels where needed)
      → cuBLAS/cuDNN (standard GEMM/conv)
        → FlashAttention (hand-written CUDA for attention)

Most codebases use all of these together — different tools for different operations.

---

Self-Check

• Explain why AI compilers exist (hand-written CUDA doesn’t scale)
• Write a simple Triton kernel and explain what the compiler handles automatically
• Explain what torch.compile does under the hood (Dynamo → Inductor → Triton)
• Explain the three torch.compile modes and when to use each
• Explain when to use Triton vs cuBLAS vs hand-written CUDA
• Explain XLA’s whole-program optimization advantage over op-level compilers

---

References

• Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations
• torch.compile Tutorial
• XLA: Optimizing Compiler for Machine Learning
• TVM: An Automated End-to-End Optimizing Compiler