Sakana AI builds an AI agent to write and optimize CUDA kernels

Writing PyTorch is easy, but squeezing every ounce of performance out of a GPU usually requires getting down into the weeds of CUDA.

That low-level work demands deep hardware expertise and endless manual tweaking. Sakana AI wants to automate the headache away, unveiling an end-to-end agentic framework that turns PyTorch code into high-performance CUDA kernels.

Taking test-time compute to the GPU metal

While frontier LLMs have shown impressive reasoning when writing high-level code, optimizing low-level systems code has largely remained out of reach.

Existing benchmarks for GPU kernels have also been riddled with exploitable loopholes and lacked diverse test conditions, making it hard to trust reported gains.

To establish a firmer foundation, Sakana AI introduced robust-kbench, a new benchmark designed to rigorously assess kernel performance and functional correctness across varied setups.

On top of this benchmark sits an automated pipeline. A frontier language model takes standard PyTorch code and translates it into an equivalent CUDA implementation before kicking off an iterative optimization loop.

Evolutionary optimization that beats native PyTorch

Raw translation is rarely optimal, so Sakana AI deployed a specialized evolutionary meta-generation procedure tailored specifically for CUDA.

To prevent the system from getting stuck on broken code, LLM-based verifiers step in to evaluate correctness and filter out faulty kernels early, saving valuable hardware verification time.

In evaluations on robust-kbench, kernels produced by this automated setup consistently beat PyTorch's native implementations in practical tasks across both forward and backward passes.

Impressively, the pipeline learned to fuse operations and deploy complex runtime optimization strategies without human intervention.

Why it matters

For machine learning engineers and systems researchers, hand-writing custom CUDA kernels is one of the most tedious bottlenecks in model scaling. An agent that can turn standard PyTorch into optimized GPU code opens the door to rapid experimentation without waiting on scarce performance engineers.

Across the broader industry, GPU performance optimization has historically been guarded by specialized talent. Sakana AI demonstrates that combining test-time search with evolutionary prompting can bridge the gap between high-level frameworks and bare-metal performance.

We think Sakana AI made a smart call by putting verifiers ahead of hardware execution. Running broken kernels directly on GPUs is slow and crash-prone, so filtering out hallucinations early makes the search loop practical. That said, it remains to be seen how well these techniques scale to irregular, non-standard architectures outside controlled benchmark suites.