Skip to content
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 captured-operator-region speedups from the paper's per-operator speedup and runtime-share table. Full-model Amdahl examples are explicitly parameterized by an illustrative captured-region fraction because the paper does not report that fraction. The measurements are from an NVIDIA DGX Spark with a GB10 GPU; 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:

  1. 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.
  2. Generate. An LLM proposes a CUDA implementation for a selected operator, conditioned on a parent kernel and its lineage.
  3. 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.
  4. 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."
  5. 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 use it

  • 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 invalid candidates actionable. Compiler and runtime failures are returned to the generator through a bounded loop. The paper reports the mechanism but does not publish a first-draft failure rate or yield.
  • 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 to use it

Validated search and economics

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."""
    if not (1 <= n_child <= n_parent):
        raise ValueError("child visits must be within [1, parent visits]")
    return latency - C * np.sqrt(np.log(n_parent) / n_child)


C = 1.0                                      # paper's reported configuration
N_P = 400
# Two feasible 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
# As the fresh child accumulates visits its exploration bonus decays.
visits = [2, 5, 20, 100, 200]
scores = [uct_min(1.12, v, N_P, C) for v in visits]
assert scores == sorted(scores)
assert scores[-1] > explored
# C=0 is pure exploitation: the lower latency wins at equal visit counts.
assert uct_min(1.00, 2, N_P, 0.0) < uct_min(1.12, 2, N_P, 0.0)

# The useful scale is the bonus difference between feasible competitors.
max_bonus = C * np.sqrt(np.log(N_P))
explored_bonus = C * np.sqrt(np.log(N_P) / 200)
fresh_margin = max_bonus - explored_bonus
assert uct_min(1.00 + 0.9 * fresh_margin, 1, N_P, C) < explored
assert uct_min(1.00 + 1.1 * fresh_margin, 1, N_P, C) > explored
try:
    uct_min(1.0, N_P + 1, N_P, C)
except ValueError:
    pass
else:
    raise AssertionError("a child cannot have more visits than its parent")


# ============== 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 widening_threshold(node_visits, root_visits):
    return int(np.floor(node_visits ** alpha(root_visits)))


def may_expand(children, node_visits, root_visits):
    """Published rule: expansion is permitted at equality."""
    return children <= widening_threshold(node_visits, 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
assert abs(alpha(500) - 0.4) < 1e-12
early = widening_threshold(100, root_visits=0)
late = widening_threshold(100, root_visits=1000)
assert (early, late) == (10, 3)
# Because the source uses <=, equality permits one more child to be generated.
assert may_expand(early, 100, 0) and not may_expand(early + 1, 100, 0)
assert may_expand(late, 100, 1000) and not may_expand(late + 1, 100, 1000)
assert widening_threshold(10_000, 0) < 10_000


# ============== 3. operator-region and full-model Amdahl arithmetic ==========
def guarded(speedup):
    """Keep generated CUDA only when it beats eager; otherwise use eager."""
    return max(1.0, speedup)


def operator_region_speedup(shares_speedups, guard=True):
    """Shares are percentages within the paper's captured operator region."""
    covered = sum(share for share, _ in shares_speedups)
    assert 0.0 <= covered <= 100.0 + 1e-9
    new = sum(share / (guarded(s) if guard else s) for share, s in shares_speedups)
    return 100.0 / (new + (100.0 - covered))


def full_model_speedup(region_speedup, region_fraction):
    """Amdahl composition when the captured region is fraction F of full latency."""
    assert region_speedup > 0.0
    assert 0.0 <= region_fraction <= 1.0
    return 1.0 / ((1.0 - region_fraction) + region_fraction / region_speedup)


RESNET = [(13.71, 1.257), (11.90, 1.004), (1.63, 1.212),
          (0.12, 1.515), (53.35, 0.999)]
r_region = operator_region_speedup(RESNET)
solo_region = operator_region_speedup([(0.12, 1.515)])
assert 1.03 < r_region < 1.04
assert solo_region < 1.0005
assert (r_region - 1) / (solo_region - 1) > 70

SD_LOSSES = [(53.40, 0.495), (27.12, 0.021), (2.55, 0.299)]
SD_WINS_SHARE = 10.48
assert 1 / min(s for _, s in SD_LOSSES) > 47
sd_region_best = operator_region_speedup(SD_LOSSES + [(SD_WINS_SHARE, 1.699)])
sd_region_worst = operator_region_speedup(SD_LOSSES + [(SD_WINS_SHARE, 1.049)])
sd_region_unguarded = operator_region_speedup(
    SD_LOSSES + [(SD_WINS_SHARE, 1.3)], guard=False)
assert sd_region_worst < sd_region_best < 1.05
assert sd_region_unguarded < 0.08
assert 1 / sd_region_unguarded > 12

# The paper does not report F. These 50% examples are illustrative, while the
# F=0 and F=1 assertions establish the Amdahl boundaries.
F_EXAMPLE = 0.50
r_full_example = full_model_speedup(r_region, F_EXAMPLE)
sd_full_worst = full_model_speedup(sd_region_worst, F_EXAMPLE)
sd_full_best = full_model_speedup(sd_region_best, F_EXAMPLE)
sd_full_unguarded = full_model_speedup(sd_region_unguarded, F_EXAMPLE)
assert full_model_speedup(r_region, 0.0) == 1.0
assert np.isclose(full_model_speedup(r_region, 1.0), r_region)
assert 1.0 < r_full_example < r_region
assert sd_region_unguarded < sd_full_unguarded < 1.0

print("all Kernel Forge assertions passed")
print("  feasible fresh-candidate latency margin:", f"{fresh_margin:.4f}")
print("  widening threshold early / late:", early, "/", late,
      "-> maximum after an allowed expansion:", early + 1, "/", late + 1)
print("  ResNet-50 captured-region speedup:", f"{r_region:.4f}x")
print("  1.52x operator alone, captured region:", f"{solo_region:.5f}x")
print("  SD3.5 captured-region guard bracket:",
      f"{sd_region_worst:.4f}x - {sd_region_best:.4f}x")
print("  SD3.5 captured region WITHOUT guard:", f"{sd_region_unguarded:.4f}x",
      f"({1 / sd_region_unguarded:.1f}x slower)")
print("  illustrative full-model results at F=50%:")
print("    ResNet-50:", f"{r_full_example:.4f}x")
print("    SD3.5 guard bracket:", f"{sd_full_worst:.4f}x - {sd_full_best:.4f}x")
print("    SD3.5 without guard:", f"{sd_full_unguarded:.4f}x",
      f"({1 / sd_full_unguarded:.1f}x slower)")

Executed output:

all Kernel Forge assertions passed
  feasible fresh-candidate latency margin: 2.2747
  widening threshold early / late: 10 / 3 -> maximum after an allowed expansion: 11 / 4
  ResNet-50 captured-region speedup: 1.0328x
  1.52x operator alone, captured region: 1.00041x
  SD3.5 captured-region guard bracket: 1.0049x - 1.0451x
  SD3.5 captured region WITHOUT guard: 0.0703x (14.2x slower)
  illustrative full-model results at F=50%:
    ResNet-50: 1.0161x
    SD3.5 guard bracket: 1.0025x - 1.0220x
    SD3.5 without guard: 0.1314x (7.6x slower)

The search rules, and the one constant you must scale

Selection combines progressive widening with a minimisation UCT. The paper permits 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. The inequality is inclusive: at 100 node visits the thresholds are 10 early and 3 late, but equality still permits one expansion, so the resulting maxima are 11 and 4. An implementation that intends a hard cap must use < instead.

The selection score is L_c - C * sqrt(log N_p / N_c); the paper uses C = 1.0. Block 1 compares only realizable states where every child visit count is at most the parent's. It shows that the selection boundary depends on the difference between the fresh and explored bonuses, not on latency units alone. Normalize latency against eager before tuning C, and reject impossible accounting states rather than using them to calibrate exploration.

Captured-region results and the missing full-model fraction

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 first computes only what the table identifies: speedup within the captured operator region.

  1. The headline operator is nearly irrelevant inside that region. adaptive_avgpool2d's 1.52x at a 0.12% region share contributes a 1.00041x region speedup. The 13.71% tensor add at 1.257x contributes more than seventy times as much.
  2. ResNet-50's reported rows imply about 1.033x for the captured region with guarded fallback. This is not a full-model latency result.
  3. Stable Diffusion 3.5's captured region brackets between about 1.005x and 1.045x with the guard. The three winners' individual shares are not reported, so the block uses the slowest and fastest winning speedups as bounds.
  4. Without the guard, the same SD 3.5 region is about 14.2x slower. The 0.021x attention result dominates the region arithmetic.

Full-model speedup requires the missing fraction F: 1 / ((1 - F) + F / S_region). The block asserts the boundaries F=0 -> 1x and F=1 -> S_region, then prints a clearly illustrative F=50% scenario. At that assumed fraction, ResNet-50 becomes about 1.016x full-model, the guarded SD result is about 1.002x to 1.022x, and the unguarded SD result is still a severe model-level regression. These are sensitivity examples, not measurements from the paper.

A second study reaches the same conclusion from the opposite direction

Kernel Forge is a fully autonomous pipeline whose safety comes from a guard applied after generation. A separate 2026 study on the MLSys 2026 FlashInfer contest asked the complementary question: what if you spend the human effort before generation instead, on the harness and the references? It ran both arms under a matched protocol on NVIDIA B200 GPUs.1

Its setup separates an evaluation harness (compilation, correctness, official-aligned timing, artifact archival) from a profile-backed optimization controller, with human-authored skills capturing operator constraints, references, profiling procedures, and promotion rules, and Codex and Claude Code agents generating candidates inside those constraints. The Full-Agent arm used an autonomous search baseline under the same protocol.

The results line up with this page's central finding. Across five operator definitions the retained Agent-Assisted artifacts reached mean-latency speedups over the supplied FlashInfer baselines of 1.62x, 18.05x, 29.68x, 1.12x, and 13.70x. Under matched final evaluation the selected Full-Agent artifacts were 1.35x to 13.25x slower than the Agent-Assisted ones on the same normalization, and critically two of them landed below the vendor baseline they were trying to beat: MoE FP8 at 0.27x and GDN Decode at 0.83x.

Read those two numbers next to the Stable Diffusion attention result above. Both studies independently produced autonomous kernels that are worse than the code they replaced, on operators that matter. Kernel Forge catches that with a guarded fallback at selection time; the harness study catches it with human-curated references and conservative promotion gates at generation time. Neither study produced a pipeline where the check was unnecessary. The paper's own conclusion is that "expert-provided optimization directions, high-quality references, and workload context remain critical for reliable AI-driven kernel optimization."

The practical reading is not that autonomy fails, since the Agent-Assisted arm is still agents writing the kernels. It is that the human contribution moved from writing CUDA to designing the environment: the constraints, the reference set, the evaluation surface, and the acceptance policy. Budget for that work explicitly, because in both studies it is what separates a speedup from a regression.

Practical protocol

  1. 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.
  2. Capture with representative inputs. The operator card records real shapes. Capture with the shapes you actually serve, not a toy batch.
  3. Target the native operators. Elementwise, normalisation, pooling and activation kernels are where the wins are. Vendor-backed operators are where the losses are.
  4. Normalise latency before tuning the UCT constant. See above.
  5. Keep guarded selection on, always. Then measure end-to-end, not per-operator.
  6. Re-validate on shapes the capture missed. Correctness was established against eager on the captured inputs only.

How to develop with it

  • Keep capture, generation, validation, benchmarking, and dispatch as separate interfaces so a model or compiler change cannot bypass correctness checks.
  • Store operator cards as versioned fixtures and add boundary shapes, non-contiguous layouts, empty dimensions, and dtype extremes before accepting a candidate.
  • Normalize latency to the eager baseline before tuning UCT, and reject visit-count states that violate 1 <= N_child <= N_parent.
  • Test the inclusive widening rule explicitly; changing <= to < changes the maximum realized fan-out by one.

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.

How to run it in production

  • Package generated kernels per GPU architecture, driver, CUDA, and PyTorch tuple; reject a package when any tuple member differs.
  • Benchmark guarded dispatch on representative full-model traffic and measure the captured-region fraction F rather than inferring it from the operator table.
  • Cap compilation and runtime repair attempts, isolate compiler processes, and retain the last error for audit.
  • Export fallback rate, numerical-rejection rate, region speedup, full-model latency, and active package identity as separate signals.

Failure modes

  • Assuming full autonomy is the goal. Two independent 2026 studies produced autonomous kernels slower than the vendor code they replaced (0.021x on SD3.5 attention here; 0.27x and 0.83x in the matched FlashInfer-contest comparison). The human work moves to harness, references and gates; it does not disappear.
  • No guarded fallback. The published SD measurements imply roughly a 14.2x slowdown inside the captured operator region; full-model damage depends on the measured region fraction.
  • Reading operator-region speedups as end-to-end. Full-model Amdahl arithmetic requires the region's fraction of total latency.
  • Chasing vendor-backed operators. cuBLAS and cuDNN paths were not beaten once in the published results.
  • UCT constant on the wrong scale. Exploration can dominate real latency differences or disappear beneath them; normalize and test the selection boundary.
  • 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
  • Shui et al., "Harness Engineering for LLM-Driven GPU Kernel Generation" (Baidu, MLSys 2026 FlashInfer contest, B200): https://arxiv.org/abs/2607.17979
  • 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

Related: AI-assisted performance optimization · Kernel fusion · CUTLASS GEMM · OpenAI Triton · PyTorch CUDA extensions · torch.compile · CUDA occupancy tuning · Roofline and arithmetic intensity · Nsight profiling workflow · PyTorch performance regression CI · Autonomous experimentation loops · Self-improving harnesses · DGX Spark · Evaluation integrity and anti-gaming

  1. Shui, Ma, Xu, Wen and Wang, "Harness Engineering for LLM-Driven GPU Kernel Generation" (arXiv 2607.17979, 2026-07-20), Baidu. MLSys 2026 FlashInfer AI Kernel Generation Contest on NVIDIA B200. Retained Agent-Assisted artifacts reached 1.62x, 18.05x, 29.68x, 1.12x and 13.70x mean-latency speedups over the supplied FlashInfer baselines across five operator definitions. Under matched final evaluation the selected Full-Agent artifacts were 1.35x to 13.25x slower than the Agent-Assisted ones, with MoE FP8 at 0.27x and GDN Decode at 0.83x falling below the supplied baseline. Not reproduced here; no B200 was available to this knowledge base. ↩