AI Papers Reader

Personalized digests of latest AI research

View on GitHub

Asking Better Questions Helps AI Agents Write Faster GPU Code

Attention, the core operation in large language models, runs fastest on hand-engineered GPU kernels, the low-level programs that execute directly on the chip. Each new NVIDIA generation has needed a fresh kernel to approach its potential, and state-of-the-art versions take expert teams months to build. A new paper from Snowflake AI Research and Carnegie Mellon University argues that AI coding agents struggle with this work not because they lack information, but because they lack direction.

The authors describe two common approaches to agent-driven kernel optimization. Scripted agents follow a fixed recipe, which makes them efficient where the recipe applies and limiting where it does not. Open-ended agents let a model freely mutate candidate programs, which can require thousands of evaluations and produces many incorrect candidates. Both, the authors note, typically respond to difficult targets by giving the agent more: profiler summaries, documentation, reference code, and stores of lessons.

The authors propose KERNELSEEK, which takes a different path. It pairs a general-purpose coding agent with ordinary developer tools, a minimal harness that checks each candidate against the original kernel and a high-precision reference, and a stop condition that returns the agent to work whenever its claim of completion fails verification. Its distinctive element is questions posed at decision points: when progress stalls, when the agent claims a hardware limit, and after suspiciously large gains. A good question names a hypothesis without supplying the answer and requires a measurement the agent can make itself. Recurring questions are collected into a bank that the agent eventually applies to itself.

The paper gives one example. An agent treated the performance of a split-KV decode kernel as a hardware ceiling. Asked whether that ceiling reflected the hardware or its own design, and where time went outside the kernel, the agent profiled the host side and found a per-call tensor fill costing 2–5 microseconds. Removing it raised split-KV performance from 0.87× to 0.97× and the overall geometric mean from 1.02× to 1.13×, measured against the upstream FlashAttention-3 (FA3) kernel.

On the same 26-configuration FA3 task, an autonomous baseline built in the prevailing style, driven by Claude Opus in 133 model-driven steps, plateaued at about half of FA3’s speed. The paper reports this as 0.50× in its text and 0.48× in a table. KERNELSEEK, which received 40 human messages, 17 of them questions, reached 1.13× FA3 in the same week on the same machine. The authors caution that each arm was run once and that the two used different front-ends.

The final kernels, released as fa3_hopper.cu and fa4_blackwell.cu, are single CUDA files that use no CUTLASS, cuBLAS, or cuDNN libraries. The Hopper kernel runs at 1.16× FA3 across 182 scenarios and was correct on all of them, though it trails the fastest available implementation, including cuDNN, by less (1.10×). The Blackwell kernel runs at 1.46× FA4 across 124 scenarios, or 1.07× the fastest available.

The limitations are substantial. The Hopper kernel is slower than upstream on fp8 dense prefill (0.67×). The Blackwell kernel does not yet support fp8 inputs or head dimensions of 192 and above. The work covers attention alone, and the question bank’s transfer to other operators is untested. The authors have not yet run the experiment that would test whether the bank can replace the human questioner.

The broader significance, as the authors frame it, is that for deep kernel work, direction may be a scarcer resource than information. The results are a single case study in attention kernels, so they suggest rather than establish that targeted questioning generalizes.