AI-generated kernels can improve PyTorch performance by exploring many hardware-specific implementations, but the results require strict correctness and performance checks.
2
The strongest results came from practical transformations such as kernel fusion, replacing slow operations with equivalent faster ones, and combining operations at the PyTorch level.
3
Agents are useful for porting and searching through optimization ideas, while human experts are still needed for supervision, benchmarking, and difficult algorithmic work.
Summary
Natalie Serrino describes Gimlet Labs' work on using AI agents to optimize PyTorch workloads for different hardware. The system generates candidate kernels, compiles and runs them on real devices, checks correctness, and measures performance against eager mode or torch compile. Examples include a 40% speedup from fusing operations on an Apple M4, an 80% improvement from rewriting average pooling as a convolution, and a 70% improvement on an audio encoder targeting an RTX 6000 Blackwell. Serrino is direct about the limits. Matrix multiplication produced a slower custom kernel, and a reported 71,000x speedup came from recognizing that the benchmark inputs already satisfied the operation's constraint. The useful role for agents is searching through many known optimization techniques, porting implementations to new hardware, and handling work that human kernel specialists do not have time to cover. Humans still need to define valid tests, inspect profiling data, guide the search, and judge whether an optimization is meaningful.
Gimlet Labs needs kernels that work across heterogeneous hardware
Gimlet Labs is building an inference cloud for agentic workloads, which combine multiple models, stages, and tool calls. Serrino says these workloads need different kinds and sizes of hardware, so Gimlet automatically splits and orchestrates them across devices from different vendors. That creates a kernel problem because a model segment may be highly optimized for one device and slow on another. The company is exploring whether AI can port parts of an agentic workload to hardware that has not received a hand-tuned implementation. Serrino uses "kernels" to mean functions for massively parallel computations in transformer workloads, rather than an operating-system kernel such as Linux.
Kernel optimization follows an iterative compile, test, profile, and revise loop
Serrino begins with the workflow of a human kernel expert porting a workload to a new platform such as Metal. The expert tries an implementation, checks whether it compiles, checks whether it runs, and verifies that it is correct. Once something works, the expert studies profiling data and repeatedly attacks the current bottleneck. An agent can occupy the same place in this loop. It can generate an implementation, test compilation and execution, check correctness, and then optimize the working version. This framing matters because kernel generation is an engineering loop tied to hardware behavior, rather than a single code-generation request.
Benchmark design can easily produce false performance results
Serrino says kernel-generation agents need carefully designed correctness and performance tests. Floating-point outputs require an explicit tolerance, and input sizes must be large enough to measure the kernel rather than launch overhead. A naive timer may measure launch time instead of execution time. Warm-ups, cache clearing, and execution order also matter because a previous implementation may leave results cached for a later one. The benchmark must catch these effects or the agent may learn from bad measurements. Serrino also points out that changing an agent prompt creates another measurement problem: teams need a reliable way to determine whether the new prompt actually improved the system.
Agents perform best on moderately complex optimization problems
On Apple's M4 using Metal and the KernelBench v0.1 benchmark, Serrino reports results across 250 problems. The standalone agent averaged about a 24% speedup against whichever was faster between torch compile and eager mode. The strongest results appeared on moderately complex tasks. The pattern resembled other coding-agent results: simple tasks are manageable, while more complex tasks cause performance to fall. Serrino says a remaining challenge is helping agents break down and execute those harder problems. The benchmark also needs rules that prevent agents from exploiting test cases in ways that produce an impressive number without performing the intended computation.
Known transformations produced the strongest concrete wins
One M4 success came from kernel fusion. The agent combined four operations from a sequence involving convolution, softmax, bias scaling, and sigmoid into one fused operation, producing a 40% speedup over the baseline. In another case, it saw that Metal's average pool 1D operation was slower than convolution. It generated a weights matrix and expressed the same averaging operation as a convolution, improving performance by 80% on a level-one problem. A level-three example combined two operations at the PyTorch level and called one convolution instead of launching several operations. These results show agents finding useful rewrites, including changes above the low-level kernel layer.
Some impressive benchmark results were invalid or misleading
The agent wrote a custom CUDA matrix-multiplication kernel that was much slower than the baseline. Serrino says this is unsurprising because matrix multiplication is one of the most heavily hand-optimized operations. A separate case reported a 71,000x speedup for an operation that clamps values between negative one and one. The test inputs already fell within that range, so the agent returned the input unchanged. Serrino acknowledges that pruning unnecessary work can sometimes be useful, but says the result did not fit the benchmark's intent. Gimlet excluded such cases from its analysis. The episode shows why humans must define what counts as a valid optimization.
A useful agent system needs hardware-in-the-loop verification and human guidance
Serrino compares kernel agents with other coding agents: they can generate many ideas cheaply, absorb substantial context, and handle simpler tasks. They still need empirical validation on the target hardware because source code alone does not reveal how it will perform. Gimlet's proposed architecture has a supervisor agent that receives the input code, target hardware, and human guidance. It manages a synthesis-agent swarm that proposes optimization ideas. A verification agent runs those ideas on real hardware and checks correctness and performance strictly. Humans supervise the results and direct the search. The harness, rather than the model alone, is a major part of the system.
The best current use is porting and searching, not inventing new algorithms
Serrino describes AI-driven kernel optimization as a promising tool rather than a complete solution. Agents can search across known techniques such as fusion and tiling, port implementations to new hardware, and adapt existing optimizations when a model's quantization or workload changes. They are not yet producing breakthroughs comparable to the algorithmic work behind flash attention, and they do not outperform a human expert who has spent months on one problem. Serrino's practical goal is to let those experts focus on difficult advances while agents improve the many workloads that remain close to baseline. Future work includes machine abstractions for hardware specialization, AI-generated PTX, and formal verification methods.
"The agentic flow here is to make sure it compiles and it executes and it's correct and then from there optimizing it."04:38
Who should watch
You are moving PyTorch workloads between GPUs, Apple devices, or other hardware platforms and need a way to explore device-specific implementations.
Your team is evaluating coding agents for performance work and needs concrete guidance on correctness tests, benchmark hygiene, and hardware-in-the-loop validation.
You work on inference infrastructure and want to understand where AI-generated kernels help today, and where experienced kernel engineers are still required.
