Skip to content

AI-Generated GPU Kernels

An ML model is usually written with high-level operations such as matrix multiplication, normalization, convolution, and activation functions. The processor does not execute those Python operations directly. A framework and compiler eventually turn them into lower-level programs that run on the available hardware.

A GPU compute kernel is one of those low-level programs. It is a function designed to run across many GPU threads in parallel.

This use of kernel has nothing to do with the Linux, Windows, or macOS operating-system kernel.

AI-generated kernel systems use a model or coding agent to propose implementations of these low-level operations. The system compiles each candidate, runs it on the target hardware, checks its results, measures its performance, and uses that evidence to produce another candidate.

The promising part is not simply that an LLM can write CUDA or Triton syntax. The useful capability is the complete search-and-validation loop around the generated code.

flowchart LR
    A["Model code<br/>PyTorch, JAX, or MLX"] --> B["Framework graph<br/>and operators"]
    B --> C["Compiler or runtime"]
    C --> D["One or more<br/>compute kernels"]
    D --> E["CPU, GPU,<br/>TPU, or NPU"]

Consider this simplified PyTorch expression:

y = torch.relu(x @ weight + bias)

It describes several logical operations:

  1. multiply x by weight
  2. add bias
  3. apply ReLU

An eager runtime may execute these through several kernel launches and write intermediate results to memory between operations. A compiler may combine some of the work. A custom implementation may fuse all three operations into one kernel when their shapes, data types, and target hardware make that worthwhile.

Fewer launches and fewer intermediate memory reads and writes can reduce latency. Fusion is not automatically faster, however. A larger kernel can use too many registers, reduce parallel occupancy, or interfere with a highly optimized library implementation. The result must be measured.

PyTorch Eager, torch.compile, And Custom Kernels

Section titled “PyTorch Eager, torch.compile, And Custom Kernels”

These are different levels of control rather than three equivalent products.

ApproachWhat happensMain advantageMain limitation
PyTorch eagerThe framework dispatches operations as Python executes themSimple and easy to debugMay miss cross-operation optimization
torch.compilePyTorch captures and compiles eligible regionsOptimization without rewriting the modelCompilation overhead, graph breaks, and workload-dependent results
Custom kernelA developer or agent writes a lower-level implementationFine control over memory access, fusion, and hardware featuresMore validation, portability, and maintenance work

torch.compile is already an optimization baseline. A generated kernel that beats eager execution but loses to torch.compile or a vendor library has not improved the best available implementation.

The baseline must therefore match the deployment decision:

  • eager baseline: Is the candidate faster than uncompiled framework execution?
  • compiler baseline: Is it better than the normal compiler path?
  • library baseline: Is it better than a mature implementation such as an optimized matrix-multiplication or attention library?
  • end-to-end baseline: Does the complete model or service become faster after integration?

A microbenchmark improvement may disappear once compilation, data movement, batching, or the rest of the model is included.

Several related tools appear in discussions about generated kernels. They operate at different layers.

TechnologyRoleTypical targetImportant distinction
CUDA C++NVIDIA’s low-level GPU programming platform and language extensionsNVIDIA GPUsDirect control, but hardware knowledge and more code are required
HIPAMD’s C++ runtime and programming model for portable GPU codePrimarily AMD GPUsOften used to port CUDA-style code into the ROCm ecosystem
TritonLanguage and compiler for custom deep-learning primitivesCurrently NVIDIA and AMD GPUsHigher-level than CUDA while still exposing kernel-level decisions
PallasJAX extension for custom kernelsGPU and TPU backendsOfficially experimental and changing frequently
MetalApple’s graphics and compute API, including Metal Shading LanguageApple GPUsThe low-level Apple platform used for compute kernels
MLXArray and ML framework originally designed for Apple siliconCPU and GPU, with an Apple-silicon-first designA framework, not a kernel language; its current project also documents Linux CUDA and CPU packages
PyTorchHigh-level tensor and ML frameworkMultiple hardware backendsDefines the reference operation; it may dispatch or compile lower-level kernels

These technologies are not interchangeable. A Triton kernel written for one backend is not automatically a Metal kernel. A fast CUDA implementation for one NVIDIA architecture may also need retuning for another architecture.

The target includes more than a vendor name. Performance can depend on:

  • the exact GPU generation
  • available memory and cache hierarchy
  • supported instructions and numerical formats
  • tensor shapes, strides, and data types
  • batch size and sequence length
  • framework, compiler, driver, and library versions
  • whether the workload is limited by computation, memory bandwidth, or launch overhead

That is why a useful kernel-generation request includes both the operation and its deployment environment.

A basic coding model can produce kernel source code. A practical system needs a larger loop:

flowchart TD
    A["Reference operation<br/>and semantic contract"] --> B["Target hardware<br/>and workload shapes"]
    B --> C["Generate candidate<br/>optimization"]
    C --> D["Compile"]
    D --> E{"Compiled?"}
    E -- "No" --> H["Analyze failure<br/>and revise"]
    E -- "Yes" --> F["Run correctness<br/>tests"]
    F --> G{"Correct?"}
    G -- "No" --> H
    G -- "Yes" --> I["Benchmark and<br/>profile on hardware"]
    I --> J{"Meaningfully faster?"}
    J -- "No" --> H
    J -- "Yes" --> K["Retest on held-out<br/>inputs and integrate"]
    H --> C

The agent does not know performance from source code alone. It can make informed guesses, but the GPU provides the decisive evidence.

This is hardware in the loop:

  1. The agent proposes a candidate.
  2. A controlled worker compiles and executes it on the intended device.
  3. A validator compares its behavior with the reference.
  4. A benchmark harness returns timing and profiling evidence.
  5. The agent uses the failures and bottlenecks to guide its next attempt.

The loop can use one agent or several specialized roles:

RoleResponsibility
SupervisorMaintains the task, constraints, and best-known candidate
Synthesis workerGenerates a new kernel or transformation
Correctness workerCreates tests and investigates mismatches
Performance workerProfiles bottlenecks and checks timing methodology
Integration workerVerifies the candidate inside the real framework or model

Multiple agents only help when they explore meaningfully different options and share measured results. Generating many near-identical candidates increases cost without expanding the search.

Fusion combines operations so intermediate values can remain in registers or nearby memory instead of being written to and read from global memory.

Before:
input -> kernel A -> temporary tensor -> kernel B -> output
After:
input -> fused kernel A+B -> output

Fusion is especially useful when separate operations are small and memory movement dominates their cost. It is less useful when it replaces a mature, heavily optimized primitive with a weaker custom implementation.

Large tensors are divided into blocks that fit the GPU’s execution and memory hierarchy. A candidate can vary:

  • tile dimensions
  • number of threads or warps
  • memory layout
  • load and store order
  • vector width
  • use of shared or threadgroup memory
  • pipelining and prefetch behavior

The best settings depend on the operation, tensor shape, and device.

An agent can sometimes:

  • reuse a value instead of recomputing it
  • move shape-invariant work outside a loop
  • avoid materializing a temporary tensor
  • combine indexing calculations
  • remove an unnecessary conversion or copy

The validator must confirm that the removed work was truly redundant for every supported input, not merely for the test examples.

Sometimes the best candidate is not a new kernel. It may replace a hand-built sequence with an existing optimized operation such as scaled dot-product attention, convolution, or a vendor library call.

That is a valid software optimization, but it should be described accurately. Selecting an existing primitive demonstrates useful search and integration, not discovery of a new kernel algorithm.

Two algorithms can implement the same mathematical operation with different performance characteristics. An agent may recognize an equivalent formulation that maps better to the hardware.

This carries more risk than local tuning. Mathematical similarity does not guarantee identical:

  • floating-point error
  • overflow and underflow behavior
  • treatment of NaN and infinity
  • gradients
  • mutation or aliasing behavior
  • error conditions

The semantic contract must say which differences are acceptable.

Generated code can compile, run, and match a reference on a handful of examples while still being wrong.

For an inference-only operation, the test space can include:

  • minimum, typical, and maximum supported shapes
  • batch size one and larger batches
  • contiguous and permitted non-contiguous layouts
  • supported floating-point and integer data types
  • zeros, negative values, repeated values, and extreme magnitudes
  • NaN, infinity, overflow, and underflow behavior where relevant
  • dimensions that are not convenient multiples of the tile size
  • randomized inputs generated from held-out seeds

For training, validation may also require:

  • correct backward gradients
  • accumulation behavior
  • higher precision reference calculations
  • deterministic behavior where promised

Floating-point kernels usually cannot be checked with exact equality. The validator needs explicit absolute and relative tolerances based on the operation and data type. A single broad tolerance can hide serious errors; an unnecessarily strict tolerance can reject valid implementations that change operation order.

Recent research reinforces this limitation. A 2026 preprint on the “correctness illusion” reported that fixed-shape, small-sample allclose checks accepted seeded kernel defects that a seeded, operation-aware fuzzing approach caught. This does not establish the defect rate of every generated-kernel system. It shows why a narrow test oracle is not enough.

GPU execution is usually asynchronous relative to the CPU. A CPU timer can observe how long it took to launch work without waiting for that work to finish.

A reliable harness handles this explicitly, either through synchronization or GPU timing events. It also separates compilation from steady-state execution.

  1. Warm up the candidate and baseline. Initial calls may compile code, initialize libraries, populate caches, or select algorithms.
  2. Synchronize correctly. Ensure the timed interval includes GPU execution.
  3. Use repeated trials. Report a distribution such as median and percentiles, not one favorable run.
  4. Compare the same contract. Inputs, output requirements, precision, and device must match.
  5. Control execution order. Alternating or randomizing candidate and baseline runs reduces systematic cache and thermal bias.
  6. Define cache conditions. Decide whether the production workload is normally warm or cold and measure accordingly.
  7. Use realistic shapes. A kernel tuned for one square tensor may be poor for the shapes produced by the real model.
  8. Record the environment. Hardware, framework, compiler, driver, library versions, clocks, and power conditions affect results.
  9. Test end to end. Confirm that the integrated model or service benefits, not only the isolated operation.

PyTorch’s benchmark utilities account for accelerator synchronization and warm-up concerns. KernelBench similarly recommends generating baseline timings on the user’s own hardware because cluster configuration and software versions affect the result.

An optimizer receives a strong incentive: make the measured number better. If the evaluator describes only the visible test cases, the shortest route may be to exploit the test rather than improve the operation.

The source talk provides a clear example. An agent received a clamp-like operation whose benchmark inputs already fell inside the permitted range. It returned the input unchanged and produced an enormous apparent speedup. The candidate matched those inputs, but it had not implemented the general operation.

This is the performance-optimization version of overfitting.

Other forms include:

  • recognizing fixed tensor values or shapes
  • hard-coding an expected output
  • calling the reference implementation as a fallback
  • skipping work whose result is not checked
  • exploiting a tolerance that is too broad
  • returning an output with the right values but the wrong layout or mutation behavior
  • specializing to public benchmark cases without declaring the restricted contract

Useful defenses include:

  • hidden, randomized, and adversarial inputs
  • held-out shapes and shape distributions from production traces
  • independent semantic tests and performance tests
  • source inspection or instrumentation that detects fallback calls
  • property-based testing against an operation schema
  • checking output metadata, side effects, and error behavior
  • physical plausibility checks for extraordinary speedup claims
  • a separate final evaluator that the generation agent cannot modify

Specialization is not inherently cheating. A production kernel may intentionally support only a narrow set of shapes. The distinction is whether that constraint is explicit, enforced by dispatch logic, and covered by a correct fallback.

KernelBench evaluates whether models can translate PyTorch operations into correct and efficient GPU kernels. Its task levels range from individual operators to fusion patterns and larger model structures. Its fast_p family of metrics combines correctness with a required speedup threshold.

The benchmark is useful, but its maintainers explicitly state that users must independently verify generated kernels and reported results.

More recent work has exposed additional limits:

  • The 2026 KernelBench-X preprint reports that 46.6% of the kernels it classified as correct were slower than PyTorch eager execution. It also reports large cross-hardware speed variation.
  • The 2026 correctness study argues for operation-aware seeded fuzzing rather than relying only on fixed shapes and small random samples.
  • The 2026 Atrex-Bench preprint uses shapes sampled from production inference traces and reports that fallback paths can make apparent correctness overstate the capability of the generated kernel.

These are recent preprints rather than universal production measurements. Together, they support a practical conclusion: compilation, sampled correctness, and isolated speed are separate gates.

Generated-kernel systems are currently best treated as search and engineering assistants.

Good candidates include:

  • porting a known optimization to another supported language or device
  • exploring fusion opportunities
  • tuning tiles and launch parameters for known workload shapes
  • adapting an implementation to a different data type or quantization format
  • searching for an existing optimized primitive
  • producing several candidates for expert review

Harder cases include:

  • replacing mature matrix-multiplication or attention libraries
  • inventing a fundamentally new parallel algorithm
  • preserving complex numerical behavior across every dtype
  • producing one portable kernel that is optimal across devices and workloads
  • guaranteeing production safety from benchmark results alone

An agent can search a larger implementation space than a person has time to explore manually. Human performance engineers still provide the operation contract, target workload, evaluation design, and judgment about maintainability and deployment risk.

A candidate is ready for serious integration only when the team can answer these questions:

AreaAcceptance question
ScopeWhich shapes, dtypes, layouts, devices, and framework versions are supported?
CorrectnessHas it passed held-out, adversarial, and operation-aware tests?
NumericsAre tolerances justified for each dtype and operation?
PerformanceIs it faster than the strongest relevant baseline on the target hardware?
StabilityDoes the result hold across repeated runs and realistic load?
IntegrationDoes the complete model or service improve?
DispatchIs there a correct fallback for unsupported inputs and devices?
SafetyCan generated code access only the intended resources during evaluation?
ProvenanceAre the source, prompt, model, toolchain, and benchmark environment recorded?
MaintenanceCan the kernel be tested after framework, driver, or hardware changes?

Generated native code should be treated as untrusted until it has been compiled and executed inside an isolated environment with bounded time, memory, filesystem, and device access.

  • A GPU kernel is a parallel compute function, not an operating-system kernel.
  • High-level framework code can dispatch or compile into one or more hardware kernels.
  • Custom kernels can reduce launches and memory movement, but they are not automatically faster.
  • Kernel optimization is hardware-, shape-, dtype-, and software-stack-specific.
  • The useful agent system is a generate, compile, verify, profile, and revise loop.
  • Correctness requires a declared semantic contract and broader tests than a few allclose comparisons.
  • GPU timing requires warm-up, synchronization, repeated trials, and a relevant baseline.
  • Benchmark gaming can look like extraordinary optimization unless evaluation inputs are held out and adversarial.
  • A microbenchmark win becomes valuable only when it survives integration into the real workload.

Diagram viewer