Ask, Don't Equip: Question-Driven Coding Agents for Deep Kernel Optimization

AI summary

Overview

Research area: Machine learning systems — LLM coding agents applied to GPU kernel optimization (specifically, writing fast CUDA attention kernels).

Technical level: Advanced. The paper assumes familiarity with GPU kernel programming (warp specialization, TMA, tensor cores, PTX/SASS) and with LLM agent design; the argument about how to steer agents is nonetheless readable by a general ML audience.

Scope (one sentence): The paper argues that coding agents optimizing kernels fail for lack of direction rather than tools, and introduces KernelSeek, a question-driven agent that reaches 1.13× FlashAttention-3 on a controlled head-to-head and produces dependency-free FA3- and FA4-class kernels at 1.16× and 1.46× geometric mean.

What This Paper Is About

LLM agents that try to make GPU kernels faster currently sit between two extremes: scripted pipelines that hard-code what to try at each step, and open-ended agents that let the model freely mutate kernels and pay for that freedom in thousands of wasted trials. The authors observe that both extremes already receive plentiful information — profilers, documentation, reference libraries, lesson stores — but neither receives a hypothesis about where the remaining performance gap is. Their goal is to show that a short, well-placed question supplies that direction without constraining what the agent may do next, and to demonstrate it on the hard case of writing dependency-free FlashAttention kernels for Hopper and Blackwell GPUs.

Key Contributions

1. Diagnosis of autonomous-agent failure on deep kernel targets. The authors build Arsenal, a deliberately serious tool-rich autonomous agent in the style of prior work (AVO, KDA, CUCo), and document quantitatively why it stalls: it plateaued at 0.50× FA3 after 133 model-driven steps, with 82% of the text it read being its own kernel, 36 of 37 rewrite branches seeded from the same step, and every diagnostic located at the kernel layer.

2. KernelSeek, an operator-agnostic, question-driven kernel agent built from three parts: a minimal evaluator harness (checks outputs against the seed and an fp32 reference on every scenario, times against the seed, enforces constraints mechanically by inspecting binary symbols), a verifiable goal installed as a Stop hook that returns the agent to work when any clause is unmet, and questions asked at five decision points. No operator-specific knowledge and no faster implementation to imitate are required.

3. A reusable question bank and skill. Starting from 48 entries distilled from traces (21 of which named the upstream kernel, 13 attention structure — both hazards removed in the release), the bank grew to 101 core entries in ten groups ordered by where the gap might be, plus 5 entries that apply only when another implementation is available and 18 per-architecture checklist entries.

4. A controlled comparison plus an artifact release. A head-to-head on the equipped agent's own evaluator (26 configurations, same H200 node, overlapping days in June 2026), an analysis of 162 human messages and 11,532 agent tool calls, and release of fa3_hopper.cu and fa4_blackwell.cu under Apache 2.0 with harness, question bank, and full traces of both approaches.

Main Findings

• Head-to-head reversal on identical ground. On the same 26 FA3 configurations, evaluator, machine, and week, the equipped agent's champion reached 0.48× FA3 at step 35 (reported as 0.50× in the abstract and introduction) with a worst configuration of 0.24× and 0 of 26 configurations at or above parity. KernelSeek started at 0.009×, reached 0.96× after 21 hours, 1.02× after 66 hours, and 1.13× after 68.0 hours (41.3 active, after excluding a 26.7-hour idle gap), with 22 of 26 configurations at or above parity and a worst case of 0.87×.

• The produced kernels. fa3_hopper.cu (3.6K lines) is correct on 182/182 Hopper scenarios at 1.16× FA3 and 1.10× the fastest available implementation including cuDNN. fa4_blackwell.cu (10.2K lines) is correct on all 154 runnable Blackwell scenarios at 1.46× FA4 and 1.07× the fastest available, including 1.10× FA4 on production shapes from SOL-ExecBench.

• The baseline weakens as it strengthens. Against cuDNN — the fastest baseline on 54 of 182 Hopper scenarios and 92 of 124 Blackwell ones — the geomeans shrink to 1.10× and 1.07×. The authors also note that the 1.5–2× edge-set numbers partly reflect upstream dispatch overhead, and that excluding the largest fp8-KV claim wins moves the Blackwell geomeans to 1.37, 1.32, and 1.08×.

• Which questions moved the agent. Every large jump followed a question that either moved the agent to a different layer of the system or rejected a claimed limit — never one that supplied missing information. Examples from the traces: a per-configuration table exposed head dimensions 96/192/256 sitting on the scalar fallback at 0.01× (removing the fallback took the geomean from 0.44 to 0.96 in 4 hours); a host-side profile found a per-call tensor fill of 2–5 μs outside the kernel (split-KV 0.87 to 0.97, geomean 1.02 to 1.13); a benchmark-validity question exposed sink and left-pad outputs off by 495× and 103×; challenging a register "floor" led to a collective rewrite taking hd256 from 0.84 to 0.95–0.96×.

• Questions did not reduce tool use, they aimed it. KernelSeek spent 26% of the text it read on its own kernels versus 82% for the equipped agent, ran the profiler 429 times on raw output (versus one parsed NCU run per step), performed 36 documentation reads (versus 92 lookups), spawned 48 subagents, and averaged about 71 tool calls per human message.

• Autonomous agents seek closure; the harness and human refused it. Across the Claude Code sessions the agent posed 35 structured questions, 24 of which offered an option to accept the current state or stop, and asserted a limit or impossibility 126 times; the Stop hook returned it to work 59 times.

• Human steering was interrogative but not only questions. Of 162 messages: directives 58 (36%), questions 44 (27%), approvals 30 (19%), corrections 19 (12%), goals 11 (7%); 63 (39%) were interrogative in phrasing and 12 of 19 corrections were phrased as questions. Question categories: measurement check (13), bound differencing (10), benchmark validity (7), gap triage (7), agent process (6), challenge a limit (5), strategy options (5), hypothesis check (3).

• Evolution amplifies direction but does not supply it. Wrapping the same agent in an AVO-style lineage on Blackwell dense prefill, the first generation's question-shaped directive took the kernel from 0.29× to 0.68×; three more generations added 0.10× between them; the search stalled at 0.78× with cycle count already below upstream's. The final 22 points came from a question that moved the agent to profile the whole call rather than the kernel, revealing synchronous host copies. A separate routing sweep moved the Blackwell geomean from 0.88 to 1.15× without changing a kernel body.

• Every remaining loss has a measured cause. Blackwell scenarios below parity spend identical tensor work over a third more cycles from twice as many barrier waits per block, pay a cold start per CTA in batched decode, or run head dimension 96 on a padded 128-wide path. On Hopper the losses are fp8 dense prefill (0.67×), softcapped prefill, head dimension 96, and long sliding windows (0.34–0.79×), which have not been profiled to the same depth.

Methodology in Plain English

The authors first build the agent they don't recommend. Arsenal runs a deterministic Python loop: generate one candidate, build it, check for forbidden libraries by inspecting the binary, analyze the SASS instruction mix, verify correctness on 26 configurations, time against FA3 across four GPUs, profile the worst case with NCU, and on promotion A/B-test inside a live vLLM instance. A fresh model session is invoked at each step, with a parser that names the bottleneck, sixteen hand-written optimization "levers" keyed to that label, retrieved documentation, hints mined from library code, and a store of lessons. The median step prompt was 22.6K characters, 37% of it pre-fetched reference material.

KernelSeek keeps only the evaluator from that stack, strips the orchestration, and hands a general-purpose coding agent an ordinary developer toolkit (compiler, profiler, disassembler, sanitizer, debugger) with nothing pre-digested. Its two additions are a goal written as a predicate over the evaluator's output — enforced by a Stop hook, quantifying over every scenario rather than an aggregate, and judged mechanically rather than by a model — and short questions injected at five decision points: campaign start, plateau, a claimed limit or completion, a suspiciously large win, and a correctness break the agent explains away. Between those points the agent works alone.

The evaluation is a case study chosen because a strong upstream exists to measure against: FlashAttention forward kernels as single CUDA files with no CUTLASS, CuTe, cuBLAS, or cuDNN, spanning 182 scenarios per architecture (bf16, fp16, fp8; paged and contiguous KV; MHA, GQA, MQA; causal, sliding-window, softcapped, and attention-sink masks; ragged batches; prefill, chunked prefill, split-KV decode). The upstream kernels serve as both correctness oracle and yardstick, with cuDNN and FlashInfer added as comparison columns fed the same buffers. Timing is back-to-back on one idle GPU with an L2 flush, clocks unlocked so only ratios are reported, and four independent Hopper batches agree within 0.3% on every aggregate.

Why This Matters

Impact on research. The paper reframes a design question that the kernel-agent literature has treated as a tooling problem. If direction — not information, structure, or freedom — is the binding constraint, then benchmarks that reward agents for accumulating context may be measuring the wrong thing, and the field's default move of adding profiler parsers and retrieval stores may be actively counterproductive. The released traces and question bank give other groups a way to test that claim rather than take it on faith, and the kernels themselves set a measured frontier for dependency-free attention implementations.

Real-world applications:

• LLM serving and training stacks that need attention kernels adapted to new GPU generations without a multi-month expert rewrite.
• Kernel work on unusual or production-specific shapes (ragged batches, paged caches, split decode) where a general library kernel leaves performance on the table.
• Serving-engine porting, where the largest single win in the paper came from the host launch path and routing logic rather than the kernel body.
• Automated performance engineering of existing CUDA code more broadly, since nothing in the method is specific to attention or to the presence of a reference implementation.

Industry relevance. The results are reported in the currency industry cares about — speedup against the shipped upstream kernel, correctness across a production feature surface, and cost. The equipped agent burned $660 of API usage over 253 sessions and 62 hours of model activity to reach half of FA3, while KernelSeek crossed parity in about 41 active hours with 40 human turns. The authors also report the honest ceiling: against cuDNN the advantage narrows to 1.10× and 1.07×, and known gaps (fp8 dense prefill, head dimensions 192+ on Blackwell) are documented rather than hidden.

Future Directions

• Does the effect hold beyond a single case study? Every arm here is a single run, and the question bank's 72 entries drawn from public sources are explicitly labeled untested. Testing the bank on operators other than attention, and on tasks where no upstream kernel exists as a bound, is the open question the authors themselves flag.
• Can the questioning be automated? The bank's first version leaked the case study into 21 of 48 entries by naming the upstream kernel. Whether an agent can be trusted to select and phrase its own questions — without a human at the five decision points — is unresolved, and the paper's own negative result (a model judge accepting a premature completion) argues for caution.
• Closing the documented residual gaps. fp8 dense prefill at 0.67×, softcapped prefill, head dimension 96, long sliding windows on Hopper, and the 28 of 182 Blackwell scenarios (fp8 inputs and dense head dimensions of 192 and above) not yet built are concrete, measured targets left for the next agent.
• Better accounting for the comparison. Token usage was not recoverable for the Cursor sessions, the two arms ran different front-ends, and KernelSeek's 0.009× starting point understates its head start after two prior days on FA3's 40-scenario grid. A cleaner protocol — matched front-ends, recovered costs, multiple runs per arm — would test how large the gap really is.

Target Audience

Researchers and engineers working on LLM coding agents, automated performance engineering, and GPU kernel generation will find the core argument and the head-to-head directly useful. Kernel and compiler specialists will care about the specific mechanisms in the released Hopper and Blackwell kernels and their measured residual gaps. ML systems practitioners weighing whether to build an agentic kernel pipeline will benefit most from the diagnosis of why the tool-rich baseline stalled, and from the honest reporting of where the advantage disappears once cuDNN enters the comparison.

Authors’ abstract

LLM agents for GPU kernel optimization sit at two ends of a structure-freedom tradeoff. Scripted agents fix what the model does at each step, via a learned policy, an optimization menu, or fixed roles, and are sample-efficient where the recipe applies and stall where it does not. Open-ended agents let a self-directed model mutate a population of kernels and pay for the freedom in trials: thousands of evaluations, repeated edits, and wrong candidates. Both respond to hard targets by equipping the agent with more: parsed profiler output, disassembly analysis, retrieval over documentation and libraries, and stores of lessons. We argue that both lack direction, a hypothesis about where the remaining gap is, and that a short question asked at the right moment supplies it without narrowing what the agent may do next. We present KernelSeek, a question-driven agent that makes an existing CUDA kernel faster with no operator-specific knowledge and no faster implementation to imitate: a general-purpose coding agent, a minimal harness that enforces a verifiable goal, and questions at decision points, first from a human, then from a bank the agent applies to itself. On the same task, evaluator, hardware, and week, a tool-rich autonomous agent built in the prevailing style plateaued at 0.50x FlashAttention-3 after 133 model-driven steps; KernelSeek reached 1.13x. It then produced dependency-free FlashAttention forward kernels for Hopper and Blackwell, single CUDA files with no CUTLASS, CuTe, cuBLAS, or cuDNN, correct on all 182 scenarios per architecture and 1.16x FlashAttention-3 and 1.46x FlashAttention-4 (geometric mean). Across 162 human messages and 11.5K agent tool calls, every large gain followed a question that moved the agent to a different layer of the system or rejected a claimed limit, never one that supplied information. We release the kernels, harness, question bank, and all traces.