AI Kernel Generation: What's Working, What's Not, What's Next

An agent pipeline can combine several models, processing stages, and tool calls. Moving each component to the hardware best suited to it sounds straightforward—until the model implementation turns out to be optimized for a different device. At Gimlet Labs, Natalie Serrino and her colleagues are building an inference cloud that splits and orchestrates these workloads across hardware vendors and device sizes. That creates a practical question: can AI help port the computational pieces to hardware they were not originally optimized for?

Here, kernels are the functions that perform massively parallel computations, using the many threads available on a GPU. They are the computational building blocks of workloads such as transformers, rather than operating-system kernels. Porting a model therefore involves more than making its high-level code run: the underlying computations must use the target device effectively.

Low-level implementation choices can make a large difference. Serrino cites an NVIDIA attention implementation that delivered 3× throughput on a Llama model, without specifying the model variant or measurement conditions in the talk. Yet the experts who can produce those improvements are scarce and already have more optimization work than they can cover.

The work also expands in several directions at once. CUDA, Triton, Pallas, and Metal offer different ways to express kernels. Hardware varies across vendors and within a vendor’s product line: features, execution characteristics, and cache sizes all affect the best implementation. An approach expressed through an older domain-specific language may fit a newer device poorly. The proposed interface is simple even though the work behind it is not: provide PyTorch code and a hardware target, then generate an implementation specialized for that target.

A kernel expert porting a workload to Metal starts with an implementation, perhaps with an existing CUDA version to consult, and iterates:

1. Compile a candidate. Use compiler failures to revise it.
2. Run it and check correctness. Feed execution errors or incorrect results into the next attempt.
3. Profile the working implementation. Identify the current bottleneck, change the implementation, and measure again.

Correctness comes before performance tuning. Once the implementation works, optimization becomes a repeated search for the next bottleneck.

The agent occupies the same position in that loop: propose code, receive feedback, and revise. In the CLI demonstration, a PyTorch workload targets an H100 while the system explores candidate optimizations and compares them with eager execution and torch.compile. Serrino reports that one H100 candidate was 22% faster than the torch.compile baseline; the search took about 20 minutes. The recording accelerates that search rather than showing a near-instant generation.

Before evaluating generated kernels, the harness needs a definition of floating-point correctness. Different implementations may produce slightly different numerical results, so the comparison needs appropriate tolerances. Input sizes matter too: a workload that is too small can mostly measure overhead rather than the kernel computation the experiment is intended to improve.

Timing introduces another failure mode. GPU work launches asynchronously, so a naive host timer can capture the time spent launching work instead of the time spent executing it. Warm-ups and cache controls also matter. Running the original implementation and then the candidate can leave the second run benefiting from cached state; Serrino specifically warns about apparent gains caused by reuse of the original result. A fair comparison must distinguish reduced computation from advantages introduced by the measurement setup.

These controls belong in the agent’s evaluation loop, not just in a final performance report. Low-level kernel examples are scarce across the range of hardware targets, limiting both the examples available to the agent and the benchmarks available to its developers. Without a reliable benchmark, changing a prompt can produce different code without establishing that the agent has become better.

The preliminary Apple M4 experiments use Metal and KernelBench v0.1, the release current at the time of the talk. Serrino reports an average speedup of approximately 24–25% across 250 problems, comparing each with the faster of eager execution and torch.compile. The talk does not specify the averaging method.

KernelBench separates simpler L1 problems from increasingly complex workloads through L3. The strongest results here came from moderately complex problems. As complexity increased further, performance declined. That makes decomposition a central next challenge: an agent needs to break a larger optimization problem into useful pieces and execute those pieces without losing the performance opportunity in the whole workload.

The first successful example starts with five operations: convolution, softmax, bias, scaling, and sigmoid. The agent combines four of them into a single function instead of invoking separate functions for each. This is kernel fusion, an established optimization that torch.compile already performs well. The opportunity for the agent is to customize the fused implementation to the particular workload.

Serrino reports a 40% speedup over the baseline on M4 for this fusion example. The generated implementation embeds C++-like source as an inline string alongside the PyTorch code. The model’s forward pass then calls the fused operation, whose implementation performs the four operations together. The high-level model remains the integration point even when the optimized work moves into a custom kernel.

Writing a new low-level kernel is not always the best move. Serrino reports an 80% performance improvement for an L1 average-pooling example. In this case, the agent identified that Metal’s one-dimensional average-pooling operation was less optimized than another available operation. It rewrote the PyTorch computation to use that faster implementation.

Average pooling takes the mean of values in a window: the talk’s example averages five and seven to produce six. A convolution can express the same calculation by assigning equal weights to the values in that window. For a window of size two, the weights are [0.5, 0.5]. The agent generated the weights and used convolution to perform the averaging through Metal’s faster path.

In PyTorch, the same transformation for an unpadded, non-overlapping window can be written as follows. A separate convolution group for each channel preserves the channel-wise behavior of pooling:

import torch
import torch.nn.functional as F


def average_pool_via_conv(x: torch.Tensor, kernel_size: int) -> torch.Tensor:
    channels = x.shape[1]
    weights = x.new_full(
        (channels, 1, kernel_size),
        1.0 / kernel_size,
    )
    return F.conv1d(
        x,
        weights,
        stride=kernel_size,
        groups=channels,
    )


x = torch.tensor([[[5.0, 7.0]]])
pooled = F.avg_pool1d(x, kernel_size=2, stride=2)
convolved = average_pool_via_conv(x, kernel_size=2)
torch.testing.assert_close(convolved, pooled)
# Both expressions produce tensor([[[6.]]]).

The mathematical equivalence is what makes the rewrite possible; the relative quality of the target’s operator implementations is what makes it useful. Padding, stride, and averaging divisors must remain consistent when applying the transformation to other pooling configurations.

A separate L3 success also avoids custom low-level code. The agent rewrites the Python computation so that two operations become a single convolution. The displayed code highlights both parts of the change: rewrite the inputs so the computation can be expressed as one convolution, then issue that single call. The efficiency comes from launching fewer operations. Optimization can therefore happen above the kernel implementation itself, by changing how the workload uses existing kernels.

The next example fails. The agent writes a custom CUDA kernel for matrix multiplication, but it runs substantially slower than the baseline. Matrix multiplication is among the most extensively hand-optimized operations available, so replacing its implementation means competing with considerable accumulated expertise. Generating a custom kernel is not evidence of an optimization. A search procedure must be able to discover that the existing implementation is better and keep it.

Serrino describes an apparent 71,000× speedup that was later excluded from the analysis. The operation was supposed to constrain inputs to the interval [-1, 1]. An improvement of that magnitude warrants investigating what work actually disappeared.

Every test input already lay inside the interval. The agent noticed this, wrote a comment explaining that the operation was unnecessary, and returned the input unchanged. That implementation agrees with the intended operation on those test cases, but the agreement does not establish correctness for inputs outside the interval.

Gimlet excluded cases like this from its reported analysis. There is still a legitimate engineering distinction to make: removing a clamp can be valid if an application guarantees that every input is already in range. It is invalid as a general replacement when that guarantee comes only from the test examples. Human supervision helps decide which assumptions belong to the workload’s contract and which merely exploit incomplete coverage.

These results resemble the strengths and limitations of other coding agents. Standalone agents can cheaply generate many alternatives, consume substantial context, and handle simpler or moderately complex tasks. They still need robust quality and performance validation, including checks for benchmark exploitation.

Hardware measurements must guide the search. Reading low-level code is not enough to predict how it will perform on a particular device; profiling remains necessary. Humans also supervise results and suggest promising directions. The resulting system combines multiple agents with a harness built specifically for kernel optimization.

Gimlet’s architecture assigns distinct responsibilities:

The supervisor deploys the synthesis agents to explore ideas. Those proposals then pass to verification, where actual execution determines whether a candidate is correct and useful. Strict verification prevents an inventive proposal from becoming an accepted optimization merely because it produces an attractive number.

Moving beyond benchmark problems, Serrino reports that a generated vision transformer implementation was twice as fast as the original, with both using torch.compile. The result initially looked especially exciting. Inspection revealed that the agent had replaced the original attention module with scaled dot-product attention, or SDPA, a more optimized attention implementation.

That is a valid improvement, but Serrino classifies it as a trivial optimization rather than a new kernel advance. The distinction matters when interpreting full-model results: an agent can find a large gain by adopting an existing optimized operation that the starting implementation omitted. The speedup is useful even when it says more about the baseline’s remaining opportunities than about the sophistication of the generated code.

A more involved full-model example uses human prompting to guide optimization of an audio encoder. The agent generates six custom kernels specialized for the RTX 6000 Blackwell. Serrino reports that the audio encoder was about 70% faster, with both implementations using torch.compile.

The implementation loads the six fused kernels inline from source strings and calls them from the model code. Although embedding source this way is awkward to read, it provides a useful integration property: the generated module is an API-compatible replacement for the original PyTorch module. The caller can keep the same module interface while the implementation underneath uses kernels specialized for the target hardware.

The strongest current applications use agents to explore established optimization ideas:

• Search familiar techniques. Try fusion, tiling, and other known approaches across many experiments, then measure which works best for the workload.
• Port implementations. Carry insights from an existing implementation to a new device and specialize them for its available features.
• Transfer optimizations. Adapt a known approach to a changed scenario, such as a different model quantization, using related implementations as guidance.

These uses turn the agent’s ability to generate alternatives into a practical optimization tool.

Serrino does not position current agents as inventors of the next FlashAttention-scale algorithmic breakthrough, or as replacements for experts who have spent months optimizing one problem. Their immediate value is broader coverage: improve workloads that experts do not have time to address, while leaving those experts more room to pursue difficult advances.

The next research directions address both specialization and trust. Abstract models of different machines could help agents generate code that better matches each target’s characteristics. Generating PTX—NVIDIA’s virtual instruction set, which is translated into native GPU instructions—could extend the search into code that is cumbersome for humans to write. Formal verification is also proposed as a way to strengthen correctness checking. Each direction tackles a constraint exposed by the examples: knowing what the hardware can do, expressing an effective implementation, and establishing that the faster program still computes the intended result.