Markdown

Agentic CUDA kernel generation¶

Scope: an agent harness that takes an unmodified PyTorch model, generates candidate CUDA kernels for the operators it actually executes, validates them against eager outputs, and integrates the winners back into the execution path. This page covers the search structure (MCTS with progressive widening and a minimisation UCT), the validation and repair loop, the guarded-selection rule, and the Amdahl arithmetic that decides whether any of it matters for your end-to-end latency. It is the concrete-system companion to AI-assisted performance optimization; the manual craft it automates is in CUTLASS GEMM, kernel fusion, and Triton.

The numpy block below is executed and asserted in this page (Python 3.11, numpy 2.x). It implements the published search rules and computes end-to-end speedups from the paper's own per-operator speedup and runtime-share table. The per-operator measurements are the paper's, taken on an NVIDIA DGX Spark with a GB10 GPU; the end-to-end figures are our Amdahl computation from those measurements, which the paper does not report directly, and they are labelled as such. Nothing here was run on a GPU by this knowledge base.

What it is¶

Kernel Forge is an open-source agentic harness with a specific design goal: to close the gap between optimising a kernel in isolation and improving a real model. The authors' complaint about prior work is precise, and it is the reason the system is shaped the way it is. Existing tools are "largely evaluated on randomly generated tensors and isolated kernels, emit standalone CUDA code that developers must manually reintegrate, mostly target only LLM PyTorch models, and offer limited support for inspecting and debugging results."

The pipeline:

Capture. Given a PyTorch model, its weights, and example inputs, capture the supported operator calls that execute during real inference, group them into workload-specific variants, and build an operator card recording the actual shapes and inputs each operator sees.
Generate. An LLM proposes a CUDA implementation for a selected operator, conditioned on a parent kernel and its lineage.
Validate. Compile with nvcc; compilation errors go back to the generator through a bounded repair loop. Launch with the operator card's stored inputs; runtime failures (launch errors, crashes, illegal memory accesses) also go to the repair loop. Surviving candidates are checked against PyTorch eager outputs.
Search. Valid candidates become nodes in an MCTS tree. Measured latency is recorded, and the best latency in a node's subtree is propagated up the ancestor chain. The tree "retains both faster and slower valid candidates because a slower candidate may still provide a useful starting point for later revisions."
Integrate. Winners are automatically integrated back into the model's execution path, with guarded selection: use the generated kernel where it improves measured operator latency, retain PyTorch eager where it does not.

Reported at 50 optimisation iterations per kernel on four models (ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, Qwen 3.5 35B-A3B): 14 kernels beat PyTorch eager, with 1.52x on adaptive_avgpool2d, 1.70x on group_norm, 2.83x on softmax in Gemma 4 E2B, and 1.54x on softmax in Qwen 3.5 35B-A3B.

Why the design choices matter¶

Optimising against real shapes beats optimising against random tensors. In real workloads kernels run under "model-specific conditions, including tensor shapes, activation distributions, neighboring operators, and memory behavior that may be absent from isolated-kernel benchmarks." The operator card is what carries that context into the search.
A tree beats a chain. Linear refinement commits to one line of attack. MCTS keeps several alive, and keeping slower valid candidates matters because a slow-but-correct kernel is often a better base for the next revision than a fast one that is structurally stuck.
The repair loop makes generation usable. Most first-draft generated CUDA does not compile. Feeding nvcc errors back is the difference between a 5% yield and a workable one.
Guarded fallback converts a risky optimiser into a safe one. This is the load-bearing safety property, and the executed block quantifies exactly how load-bearing.
Integration in place is the actual product. Emitting standalone CUDA that a human must wire back in is where most of the engineering time goes.

When to use it (and when not)¶

Use it when you have a model with a meaningful share of runtime in native, non-vendor-optimised operators (elementwise ops, normalisations, pooling, activations) and you want that share reduced without hand-writing CUDA.

Do not use it when:

Your runtime is dominated by GEMM and convolution. These are already vendor-tuned. The published results show generated kernels lose to cuBLAS and cuDNN paths every time, and often by a lot.
You cannot afford guarded selection. Without it, the same measurements produce a catastrophic regression. See below.
You need a correctness guarantee beyond output matching. Validation compares against eager outputs on the captured inputs. That is a strong smoke test, not a proof, and it does not cover shapes the capture run never saw.
You expect the headline speedups to be end-to-end. They are per-operator. The arithmetic below is the part that decides whether they matter.

Architecture¶

flowchart TB
  M["unmodified PyTorch model<br/>+ weights + example inputs"] --> CAP["capture executed operators<br/>build operator cards<br/>(real shapes, real inputs)"]
  CAP --> SEL["MCTS controller:<br/>select parent kernel + lineage"]
  SEL --> GEN["LLM generator:<br/>candidate CUDA source"]
  GEN --> NVCC{"nvcc compiles?"}
  NVCC -->|"no"| REP["bounded repair loop<br/>(compiler error fed back)"]
  REP --> GEN
  NVCC -->|"yes"| RUN{"launches cleanly?"}
  RUN -->|"no: crash / illegal access"| REP
  RUN -->|"yes"| CHK{"matches PyTorch eager<br/>on operator-card inputs?"}
  CHK -->|"no"| REJ["reject"]
  CHK -->|"yes"| NODE["insert node with measured latency"]
  NODE --> BACK["propagate best subtree latency<br/>up the ancestor chain"]
  BACK --> SEL
  NODE --> GUARD{"beats eager latency?"}
  GUARD -->|"yes"| USE["integrate generated kernel"]
  GUARD -->|"no"| FALL["retain PyTorch eager"]

How the search and the economics work, validated¶

import numpy as np

# ============== 1. minimisation UCT: lower score wins =======================
def uct_min(latency, n_child, n_parent, C):
    """score(c|p) = L_c - C * sqrt(log N_p / N_c). Lower is selected."""
    return latency - C * np.sqrt(np.log(n_parent) / n_child)


C = 0.5
N_P = 400
# Two children: one fast and well-explored, one slower and barely tried.
explored = uct_min(latency=1.00, n_child=200, n_parent=N_P, C=C)
fresh = uct_min(latency=1.12, n_child=2, n_parent=N_P, C=C)
assert fresh < explored                       # exploration bonus wins early
# As the fresh child accumulates visits its bonus decays and the exploit wins.
assert uct_min(1.12, 200, N_P, C) > explored
visits = [2, 5, 20, 100, 400]
scores = [uct_min(1.12, v, N_P, C) for v in visits]
assert scores == sorted(scores)               # monotonically less attractive
# C=0 is pure exploitation: the lowest measured latency always wins.
assert uct_min(1.00, 1, N_P, 0.0) < uct_min(1.12, 1, N_P, 0.0)
# The exploration bonus is BOUNDED at C*sqrt(log N_p), reached when N_c = 1.
max_bonus = C * np.sqrt(np.log(N_P))
assert abs((1.0 - uct_min(1.0, 1, N_P, C)) - max_bonus) < 1e-12
assert abs(max_bonus - 1.22387) < 1e-5
# So a candidate whose latency is worse than that bound is never selected, no
# matter how under-explored it is. C therefore has to be scaled to the units of
# L: with latencies in milliseconds, C=0.5 is effectively pure exploitation.
assert uct_min(9.99, 1, N_P, C) > uct_min(1.00, 10_000, N_P, C)
assert uct_min(1.00 + max_bonus * 0.9, 1, N_P, C) < uct_min(1.00, 10_000, N_P, C)
assert uct_min(1.00 + max_bonus * 1.1, 1, N_P, C) > uct_min(1.00, 10_000, N_P, C)


# ============== 2. progressive widening and the alpha anneal ================
def alpha(root_visits, a0=0.5, a1=0.3, anneal_over=1000):
    t = min(root_visits, anneal_over) / anneal_over
    return a0 + (a1 - a0) * t


def max_children(node_visits, root_visits):
    return int(np.floor(node_visits ** alpha(root_visits)))


assert abs(alpha(0) - 0.5) < 1e-12
assert abs(alpha(1000) - 0.3) < 1e-12
assert abs(alpha(2000) - 0.3) < 1e-12          # clamped after the anneal window
assert abs(alpha(500) - 0.4) < 1e-12
# Early search branches broadly; late search reuses existing branches.
assert max_children(100, root_visits=0) == 10
assert max_children(100, root_visits=1000) == 3
assert max_children(100, 0) > max_children(100, 1000)
# A node must be visited before it can widen at all.
assert max_children(1, 0) == 1
# Widening is sublinear in visits either way, which is what bounds the fan-out.
assert max_children(10_000, 0) < 10_000 and max_children(10_000, 1000) < 10_000


# ============== 3. guarded selection, and why it is not optional ============
def guarded(speedup):
    """Keep the generated kernel only when it beats eager; else fall back."""
    return max(1.0, speedup)


def amdahl(shares_speedups, guard=True):
    """shares in percent of captured operator-region time."""
    covered = sum(s for s, _ in shares_speedups)
    assert covered <= 100.0 + 1e-9
    new = sum(s / (guarded(v) if guard else v) for s, v in shares_speedups)
    return 100.0 / (new + (100.0 - covered))


# --- ResNet-50 at opt50 (paper's reported per-operator numbers) -------------
RESNET = [(13.71, 1.257),   # tensor add        (generated win)
          (11.90, 1.004),   # relu              (generated win, marginal)
          (1.63, 1.212),    # max pool          (generated win)
          (0.12, 1.515),    # adaptive avgpool  (the 1.52x headline)
          (53.35, 0.999)]   # conv2d            (vendor-backed, no win)
r_guarded = amdahl(RESNET)
assert 1.03 < r_guarded < 1.04, r_guarded
# The 1.52x headline operator contributes almost nothing: it is 0.12% of runtime.
solo = amdahl([(0.12, 1.515)])
assert solo < 1.0005, solo
assert (r_guarded - 1) / (solo - 1) > 70       # the boring 13.71% add dominates

# --- Stable Diffusion 3.5 Medium: the losses are on the big operators -------
SD_LOSSES = [(53.40, 0.495),   # linear
             (27.12, 0.021),   # scaled-dot-product attention: ~48x SLOWER
             (2.55, 0.299)]    # conv2d
SD_WINS_SHARE = 10.48          # group_norm + layer_norm + SiLU, combined
assert min(v for _, v in SD_LOSSES) == 0.021
assert 1 / 0.021 > 47                          # a 47x regression, not a rounding

# With the guard, the three winning operators bound end-to-end gain. Their
# individual shares are not broken out, so bracket the blended rate.
best = amdahl(SD_LOSSES + [(SD_WINS_SHARE, 1.699)])
worst = amdahl(SD_LOSSES + [(SD_WINS_SHARE, 1.049)])
assert worst < best
assert best < 1.05, best                       # even the optimistic bound is <5%
assert worst < 1.006, worst

# WITHOUT the guard the same measurements make the model an order of magnitude
# slower, because the 27% operator regressed by ~48x.
unguarded = amdahl(SD_LOSSES + [(SD_WINS_SHARE, 1.3)], guard=False)
assert unguarded < 0.08, unguarded
assert 1 / unguarded > 12                      # >12x end-to-end SLOWDOWN
assert best / unguarded > 13

print("all Kernel Forge assertions passed")
print("  ResNet-50 end-to-end with guard:", f"{r_guarded:.4f}x")
print("  ResNet-50 from the 1.52x headline op alone:", f"{solo:.5f}x")
print("  SD3.5 end-to-end with guard, bracket:", f"{worst:.4f}x - {best:.4f}x")
print("  SD3.5 end-to-end WITHOUT guard:", f"{unguarded:.4f}x",
      f"({1 / unguarded:.1f}x slower)")
print("  max children at 100 visits, early / late:",
      max_children(100, 0), "/", max_children(100, 1000))

Executed output:

all Kernel Forge assertions passed
  ResNet-50 end-to-end with guard: 1.0328x
  ResNet-50 from the 1.52x headline op alone: 1.00041x
  SD3.5 end-to-end with guard, bracket: 1.0049x - 1.0451x
  SD3.5 end-to-end WITHOUT guard: 0.0703x (14.2x slower)
  max children at 100 visits, early / late: 10 / 3

The search rules, and the one constant you must scale¶

Selection combines progressive widening with a minimisation UCT. Widening allows a new child while |children(x)| <= floor(N_x ^ alpha(R)), with alpha annealed from 0.5 to 0.3 over the first 1000 root visits, "allowing broader branching early in search and more reuse of existing branches later." Block 2 makes that concrete: a node with 100 visits may hold 10 children early and only 3 late.

The selection score is L_c - C * sqrt(log N_p / N_c), and because it is a minimisation the exploration term is subtracted. Block 1 surfaces a consequence that is easy to miss when porting a UCT from a win-rate setting: the exploration bonus is bounded at C * sqrt(log N_p), so a candidate whose latency is worse than that bound will never be selected however under-explored it is. In a standard UCT the reward is a normalised value in [0, 1] and a C around 1 is sensible. Here L is a latency, so C has to be scaled to the units and magnitude of the latencies you are comparing. With latencies in milliseconds and C = 0.5, the search is effectively pure exploitation and the tree degenerates toward a chain, which is the thing MCTS was introduced to avoid. Normalise latencies (for example, against the eager baseline) before tuning C.

The honest end-to-end picture¶

The per-operator table is the most useful thing in the paper, and it tells a consistent story: generated kernels win on cheap native operators and lose on vendor-backed ones.

ResNet-50 at 50 iterations:

Operator	Speedup vs eager	Share of captured operator time
tensor add	1.257x	13.71%
relu	1.004x	11.90%
max pool	1.212x	1.63%
adaptive avgpool	1.515x	0.12%
conv2d	0.999x	53.35%
batch norm	0.992x	not reported
linear	0.903x	not reported

Stable Diffusion 3.5 Medium shows "the same pattern more sharply": group_norm at 1.699x, layer_norm at 1.049x and SiLU at 1.052x together account for only 10.48% of captured operator time, while linear reaches 0.495x at 53.40%, scaled-dot-product attention reaches 0.021x at 27.12%, and conv2d reaches 0.299x at 2.55%.

Block 3 turns those into end-to-end numbers with Amdahl's law and guarded selection. Three results:

The headline operator is nearly irrelevant. adaptive_avgpool2d's 1.52x is real, but it is 0.12% of runtime, worth 1.00041x end-to-end. The unglamorous 13.71% tensor-add at 1.257x contributes more than seventy times as much. When you read a per-operator speedup, the first question is always the runtime share.
ResNet-50 lands at about 1.033x end-to-end, from the reported operators, with everything else falling back to eager. A 3.3% improvement for a fully automated pass over an unmodified model is a genuinely useful result. It is not the 1.52x the abstract leads with.
Stable Diffusion 3.5 brackets between 1.005x and 1.045x (the three winners' individual shares are not broken out, so the block computes both ends). Even the optimistic end is under 5%.

And the result that should govern how you deploy this: without guarded selection the same measurements make SD 3.5 about 14x slower end-to-end, because a 27% operator regressed by a factor of 48. The guard is not a nice-to-have or a tuning refinement. It is the only thing standing between this pipeline and a catastrophic regression, and any reimplementation that omits it is unsafe by construction.

How to use it: a practical protocol¶

Profile first. Get the runtime share of every operator before you start. If your top three operators are GEMM, convolution and attention, the expected payoff is small and you should know that in advance.
Capture with representative inputs. The operator card records real shapes. Capture with the shapes you actually serve, not a toy batch.
Target the native operators. Elementwise, normalisation, pooling and activation kernels are where the wins are. Vendor-backed operators are where the losses are.
Normalise latency before tuning the UCT constant. See above.
Keep guarded selection on, always. Then measure end-to-end, not per-operator.
Re-validate on shapes the capture missed. Correctness was established against eager on the captured inputs only.

How to maintain it¶

Pin the toolchain. A generated kernel is validated against a specific nvcc, a specific driver, and a specific GPU. The published results are on a DGX Spark with a GB10 GPU; nothing transfers to a different architecture without re-measurement.
Re-run the guard after any upgrade. A PyTorch or cuDNN update can move the eager baseline, which can silently flip a guarded decision from "use generated" to "should have fallen back".
Store the tree, not just the winner. Slower valid candidates are useful revision bases, which is exactly why the search keeps them.
Treat generated CUDA as third-party code. It is machine-written, compiled, and executed on your hardware. Review, sandbox, and version it accordingly.

Failure modes¶

No guarded fallback. The dominant risk, quantified above at roughly 14x end-to-end slowdown on one of the published workloads.
Reading per-operator speedups as end-to-end. Amdahl's law is unforgiving and the headline operators here have tiny runtime shares.
Chasing vendor-backed operators. cuBLAS and cuDNN paths were not beaten once in the published results.
UCT constant on the wrong scale. Turns the tree search into a linear chain, silently.
Correctness by output-matching only. Passing on captured inputs does not cover unseen shapes, edge values, or numerical accumulation differences that only show up over many steps.
Repair-loop budget exhaustion mistaken for capability. A bounded repair loop that keeps hitting its bound means the generator cannot write valid CUDA for that operator; more iterations will not fix it.
Generalising across GPU architectures. A kernel tuned for one GPU's occupancy and memory hierarchy is a starting point elsewhere, not a result.

References¶

Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels: https://arxiv.org/abs/2607.24762
KernelBench (isolated-kernel benchmark the paper positions against): https://arxiv.org/abs/2502.10517
UCT / bandit-based Monte Carlo planning: https://link.springer.com/chapter/10.1007/11871842_29
Progressive widening in MCTS: https://hal.science/hal-00542673