Autonomous Kernel Synthesis: Architecting Closed-Loop LLM Autotuners for Custom Triton Operators

The Bottleneck of Manual GPU Kernel Optimization

Writing high-performance GPU kernels requires deep mechanical sympathy with underlying hardware architectures. For engineers working with modern deep learning primitives—whether crafting custom fused flash-attention variants, quantized GEMMs, or non-standard activation layers—optimizing performance in CUDA or OpenAI's Triton is traditionally an artisanal, high-friction process.

Maximizing hardware utilization across NVIDIA Hopper (H100) or Ada Lovelace architectures requires juggling register pressure, shared memory (smem) bank conflicts, warp divergence, asynchronous memory copies (cp.async), and tensor core pipelining. A human engineer typically writes a baseline kernel, measures execution latency, drops into Nsight Compute (ncu) to inspect SASS/PTX metrics, diagnoses occupancy bottlenecks, adjusts tiling configurations, and repeats.

Recent advances in LLM-driven autonomous research agents demonstrate that this profile-analyze-refactor cycle can be fully closed-looped. By combining localized LLMs with a strict compilation, testing, and profiling harness, we can build autonomous kernel discovery loops capable of achieving massive speedups over unoptimized baselines.

Architecture of a Closed-Loop Kernel Optimization System

An automated kernel generation loop is not merely an LLM writing code in a prompt loop; it is a stateful reinforcement and heuristic search architecture governed by deterministic compiler telemetry.

+-------------------------------------------------------------------------+
|                        Autonomous Tuning Loop                           |
|                                                                         |
|  +----------------+      Generates Triton      +---------------------+  |
|  |   LLM Agent    | -------------------------> |  Sanity & Correctness|  |
|  |  (Context Bank)|                            |    Fuzzing Engine   |  |
|  +----------------+                            +---------------------+  |
|          ^                                                |             |
|          | Telemetry & Error Logs                         | Pass        |
|          |                                                v             |
|  +----------------+      Compiles/Profiles     +---------------------+  |
|  | Telemetry      | <------------------------- | Benchmark & Nsight  |  |
|  | Synthesizer    |      (L1/L2, Occupancy,    | Telemetry Harvester |  |
|  +----------------+       Latency, SASS)       +---------------------+  |
+-------------------------------------------------------------------------+

1. The Specification & Baseline Phase

The agent begins with a mathematically rigid specification written in standard PyTorch (acting as ground truth) alongside an initial, naive Triton kernel implementation.

2. Correctness & Boundary Fuzzing

Before any performance profiling occurs, the candidate kernel undergoes extensive validation against random tensors, edge-case dimensions (odd batch sizes, non-power-of-two sequence lengths), and varying strides. Any compilation failure, CUDA runtime exception, or numerical drift beyond tolerance ($ ext{atol}=1 ext{e-}3, ext{rtol}=1 ext{e-}3$ for FP16/BF16) triggers immediate feedback to the agent without consuming benchmark GPU cycles.

3. Micro-Benchmarking and Hardware Telemetry Harvester

If correct, the kernel is passed to a high-precision benchmarking suite utilizing CUDA Events, with L2 cache flushing between warmup iterations to avoid warm-cache bias. For deep analysis, programmatic invocations of ncu (Nsight Compute CLI) extract crucial hardware metrics:

Compute (SM) Throughput %
DRAM & L2 Bandwidth Utilization %
Register Spill Overflows (local_load_throughput / local_store_throughput)
Warp Execution Efficiency & Issue Stall Reasons

4. Telemetry Synthesis and Dynamic Context Assembly

Raw profiling outputs contain thousands of lines of low-level data. The telemetry synthesizer compresses these metrics into actionable semantic diagnoses (e.g., "Occupancy is limited to 33.3% due to high register pressure: 64 registers/thread. Shared memory usage: 48KB/block. Issue stalls dominated by stall_long_sb (memory scoreboard)."). This condensed insight is fed back into the LLM context to drive the next code generation iteration.

Designing the Agent Feedback Loop

To make an autonomous agent perform genuine engineering instead of random walks through code permutations, the system prompt must enforce structural thinking.

Structured Prompt Payload

{
  "iteration": 4,
  "current_latency_us": 42.6,
  "baseline_latency_us": 310.2,
  "speedup_vs_baseline": "7.28x",
  "correctness": true,
  "compiler_feedback": {
    "registers_per_thread": 48,
    "shared_memory_bytes": 32768,
    "occupancy_pct": 66.7,
    "primary_stall": "stall_barrier"
  },
  "profiler_summary": "High barrier wait times indicate thread imbalance across warps during reduction step. Suggest vectorizing memory loads using tl.load with block pointers or tuning num_stages."
}

By providing the agent with precise hardware constraints, it can formulate targeted hypotheses, such as:

"If I increase BLOCK_SIZE_K from 32 to 64, I can improve arithmetic intensity in the tensor core dot product."
"Increasing num_stages from 2 to 4 enables asynchronous pipelining via TMA (Tensor Memory Accelerator) or cp.async, hiding global memory latency behind compute."
"Register allocation is spilling into local memory; refactoring intermediate accumulation tensors from 2D slices to vectorized scalar accumulators will preserve SRAM capacity."

Case Study: Optimizing a Fused LayerNorm + GELU Kernel

Let us trace how an autonomous loop iteratively refactors a fused LayerNorm + GELU operator on an A100 (SXM4 80GB) from naive execution to peak bandwidth saturation.

Iteration 0: The Naive Baseline

The baseline kernel uses a straightforward per-row threadblock assignment with basic scalar operations.

import triton
import triton.language as tl
import torch

@triton.jit
def _naive_layernorm_gelu_kernel(
    X_ptr, Y_ptr, W_ptr, B_ptr, Mean_ptr, Rstd_ptr,
    stride_x, stride_y,
    N, eps,
    BLOCK_SIZE: tl.constexpr,
):
    row_idx = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < N
    
    x = tl.load(X_ptr + row_idx * stride_x + cols, mask=mask, other=0.0)
    w = tl.load(W_ptr + cols, mask=mask, other=0.0)
    b = tl.load(B_ptr + cols, mask=mask, other=0.0)
    
    # Two-pass calculation of mean and variance
    mean = tl.sum(x, axis=0) / N
    var = tl.sum((x - mean) * (x - mean), axis=0) / N
    rstd = 1.0 / tl.sqrt(var + eps)
    
    normed = (x - mean) * rstd
    scaled = normed * w + b
    
    # Approximation of Fast GELU: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
    gelu_out = 0.5 * scaled * (1.0 + tl.sin(0.79788456 * scaled * (1.0 + 0.044715 * scaled * scaled)))
    
    tl.store(Y_ptr + row_idx * stride_y + cols, gelu_out, mask=mask)

Profiler Telemetry:

Execution Time: 184.2 μs
Memory Throughput: 112 GB/s (A100 theoretical peak: ~2,039 GB/s)
Issue: Inefficient two-pass memory reads for mean/variance, lack of Welford's algorithm, non-vectorized load transactions, and excessive trigonometric approximation overhead.

Iteration 1-3: Algorithmic Refactoring (Welford's One-Pass Algorithm)

The agent refactors the reduction step to use a parallelized Welford reduction (tl.extra.cuda or manual parallel tree reduction). This eliminates redundant SRAM spills and combines mean/variance calculations into a single spatial pass.

Iteration 4-6: Pipelining, Vectorization, and Memory Coalescing

The agent introduces vectorized pointer arithmetic, pointer alignment hints, and tunes compiler warps/stages.

@triton.jit
def _optimized_layernorm_gelu_kernel(
    X_ptr, Y_ptr, W_ptr, B_ptr,
    stride_x_row, stride_y_row,
    N: tl.constexpr, eps: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    row_id = tl.program_id(0)
    row_x_ptr = X_ptr + row_id * stride_x_row
    row_y_ptr = Y_ptr + row_id * stride_y_row

    cols = tl.arange(0, BLOCK_N)
    mask = cols < N

    # Coalesced 128-bit aligned vector loads
    x = tl.load(row_x_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    w = tl.load(W_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    b = tl.load(B_ptr + cols, mask=mask, other=0.0).to(tl.float32)

    # Numerically stable reduction
    mean = tl.sum(x, axis=0) / N
    diff = tl.where(mask, x - mean, 0.0)
    var = tl.sum(diff * diff, axis=0) / N
    rstd = 1.0 / tl.sqrt(var + eps)

    norm = diff * rstd
    out = norm * w + b

    # Native tanh approximation for GELU without register-heavy sin/poly expansions
    fast_gelu = 0.5 * out * (1.0 + tl.math.tanh(0.79788456 * (out + 0.044715 * out * out * out)))

    tl.store(row_y_ptr + cols, fast_gelu.to(tl.float16), mask=mask)

Autotuned Runner Configuration:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_N': 1024}, num_warps=4, num_stages=3),
        triton.Config({'BLOCK_N': 2048}, num_warps=8, num_stages=4),
        triton.Config({'BLOCK_N': 4096}, num_warps=16, num_stages=4),
        triton.Config({'BLOCK_N': 8192}, num_warps=32, num_stages=5),
    ],
    key=['N'],
)
def fused_ln_gelu(x, w, b, eps=1e-5):
    M, N = x.shape
    y = torch.empty_like(x)
    grid = lambda META: (M,)
    _optimized_layernorm_gelu_kernel[grid](
        x, y, w, b,
        x.stride(0), y.stride(0),
        N=N, eps=eps
    )
    return y

Final Profiler Telemetry:

Execution Time: 18.9 μs
Speedup: 9.74x over baseline
Memory Throughput: 1,740 GB/s (~85.3% of theoretical peak bandwidth)
Warp Occupancy: 88.5%

Automated Correctness and Fuzzing Framework

A critical trap in automated kernel synthesis is optimization at the expense of numeric stability. Small shifts in operation order (e.g., floating-point associative reordering) alter numerical precision.

To prevent regressions, the harness enforces a multi-tier test harness:

def verify_kernel_invariants(torch_fn, triton_fn, input_shapes, dtypes=[torch.float16, torch.bfloat16]):
    for dtype in dtypes:
        for shape in input_shapes:
            # High-variance inputs (testing near zero and dynamic range bounds)
            x = torch.randn(shape, device='cuda', dtype=dtype) * 10.0
            w = torch.randn(shape[-1], device='cuda', dtype=dtype)
            b = torch.randn(shape[-1], device='cuda', dtype=dtype)
            
            expected = torch_fn(x, w, b)
            actual = triton_fn(x, w, b)
            
            # Validate non-finite checks
            assert not torch.isnan(actual).any(), "NaN detected in output"
            assert not torch.isinf(actual).any(), "Inf detected in output"
            
            # Relative error tolerance check
            torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2)

If tolerances fail, the exact delta tensor $\Delta = |y_{\text{actual}} - y_{\text{expected}}|$ along with maximum absolute error indices are serialized into the LLM context, highlighting specifically where boundary conditions failed.

Common Failure Modes and Autonomous Mitigations

Key Takeaways for High-Throughput Autotuning

Telemetry is Context: An LLM cannot optimize CUDA or Triton code effectively on code structure alone. It requires high-fidelity compiler and hardware performance counters to make rational optimization steps.
Isolate Correctness from Performance: Never benchmark a failing or numerically drifted kernel. Failing early preserves compute and avoids misleading latency numbers caused by early kernel aborts.
Search Space Structuring: Let the LLM handle structural, mathematical, and algorithmic refactoring (e.g., changing algorithms, vectorizing, re-ordering memory passes), while delegating discrete hyperparameter tuning (block sizes, warp counts, pipelined stages) to Triton's native @triton.autotune decorator.

By uniting structured compiler feedback, rigorous mathematical fuzzing, and modern LLMs, engineering teams can build scalable, autonomous autotuning pipelines that turn routine GPU kernel optimization into an automated background process.