Large language models optimize AI accelerator kernels over days
A Case Study of a Multi-Agent Kernel-Optimization Loop That Learns a New AI Accelerator
Machine LearningArtificial IntelligenceHardware ArchitectureProgramming Languages
Summary
Some AI models can write computer instructions for graphics processors well, but they struggle with newer AI hardware that looks very different. The authors show how an AI agent repeatedly improves small programs for a special AI chip called the FuriosaAI RNGD NPU by learning about the chip and its programming language on its own. The agent tests many changes carefully because running code on the real chip is slow and noisy, eventually making the programs about nine times faster than the originals. This approach creates detailed knowledge about the chip and keeps improving the programs as the software tools change.
What this means in practice
- •For hardware engineers: Automate tuning of kernel programs for new AI chips to maximize speed on scarce devices without much prior documentation.
- •For performance software teams: Use an iterative AI agent loop to optimize critical decoding kernels on novel AI accelerators for faster model inference.
Authors
Byonggon Chun
Abstract
Frontier large language models write correct CUDA kernels for most single-operator benchmark tasks, yet their accuracy collapses on accelerators whose architecture and programming language they have rarely seen. We present a case study of an LLM agent loop built for such accelerators. Our example is the FuriosaAI RNGD NPU, programmed in a recent, sparsely documented Rust DSL, where the loop optimized three Gemma-4-12B decode kernels over 12 days on a shared, scarce device. The loop builds its own target knowledge: it writes and checks a summary of the DSL documentation, measures the chip with microbenchmarks, records compiler-derived legality rules, and revisits this knowledge when the toolchain changes, rechecking rules and measurements where it can. A role-based search over per-kernel Design Trees explores whole-kernel designs and accumulates small accepted changes in the best version of each kernel. Because device time is scarce and readings are noisy, cheap checks screen new designs before any device run, and late in the run a change was accepted only after repeated comparisons with the best kernel on the device. The final kernels reach a 9.293x geomean speedup over the official baseline, each within 1.25x of the cycles needed just to move its weights. Our prompts, profiler and acceptance test are available at https://github.com/bgchun-fs/multi-agent-kernel-gen-loop.