KernelSeek improves GPU code by asking smarter questions
Ask, Don't Equip: Question-Driven Coding Agents for Deep Kernel Optimization
Machine Learning
Summary
Improving computer code that runs on powerful graphics processors is hard and often trial-and-error. The authors show that programming assistants can do much better if they ask the right questions at the right time instead of just trying lots of guesses or following fixed recipes. Their tool KernelSeek uses carefully chosen questions to guide code changes and ended up producing faster and simpler code than other automatic methods. This approach worked well for speeding up a complex GPU process called FlashAttention on several hardware versions.
What this means in practice
- •For gpu software developers: Create faster GPU programs by integrating question-driven coding agents to guide performance improvements without manual optimization expertise.
- •For hardware acceleration engineers: Generate efficient, independent CUDA kernels for new GPU architectures using a minimal harness that verifies correctness across scenarios.
Authors
Gabriele Oliaro, Jaeseong Lee, Zhaoyuan Su, Sreeram Vennam, Zhihao Jia, Samyam Rajbhandari
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.