A Case Study of a Multi-Agent Kernel-Optimization Loop That Learns a New AI Accelerator
Organizations: Fluidstack
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.
Figures & tables
| Aspect | AccelOpt [ 4 ] | Autocomp [ 3 ] | EvoKernel [ 2 ] | Ours |
|---|---|---|---|---|
| Target | AWS Trainium (NKI), also H100 (Triton) | Gemmini, AWS Trainium, NVIDIA L40S | NPU (Ascend C), also CUDA | FuriosaAI RNGD NPU (furiosa-opt Rust DSL), shared and scarce |
| Target knowledge | NKI basics and a programming guide in the prompt; lessons from slow-fast pairs | ISA in the prompt (dropped for the GPU), a menu of optimizations | Data-scarce NPU language; memory seeded with API templates | Built by agents (DSL summary, microbenchmarks, legality rules); rechecked after a toolchain change |
| Unit of search | A candidate kernel in the beam | One optimization per step: plan, then implement | A first draft, then refinements | Whole-kernel designs and their variants; one head per kernel |
| Search | Beam search: keep 6, 144 new candidates per round, 16 rounds | Beam search (6), menu dropout, LLM ensembling | Memory-based reinforcement learning with value-driven retrieval | Design Tree; cheap checks before device runs; variants raced |
| Measurement | Latency on the chip via profiling; % of peak | Real hardware, or a cycle-accurate simulator for Gemmini | Anti-hacking, compile and correctness gates; profiler latency | Per-cluster device traces; repeated acceptance test against the head |
| Kernel | Baseline | Final | Speedup | Peak | Final / |
|---|---|---|---|---|---|
| cycles | cycles | ( ) | cycles | peak | |
| qkv | 250,514 | 66,817 | 3.75 | 54,600 | 1.22 |
| attn_out | 404,633 | 36,227 | 11.17 | 29,000 | 1.25 |
| ffn | 3,703,473 | 193,246 | 19.16 | 174,000 | 1.11 |
| Geomean | 9.293 |
AI Accelerator Inference
AI accelerator inference studies hardware and software for running neural networks faster and more efficiently across data-center, edge, and mobile devices.
See all 107 papers