Achieving Extreme Efficiency through Specialized GPU Kernel Generation

Achieving Extreme Efficiency through Specialized GPU Kernel Generation | Databricks Blog Skip to main content

Summary

New kernel drafts are cheap and easy to produce in parallel. Trust is not. The system only makes progress as fast as we can check those drafts.

A number that looks miraculously fast is often a measurement bug: leftover work from a previous run, a comparison where the two sides were not doing the same thing, or a kernel that only looks good under a hidden assumption.

Context is a tradeoff, not a pile to maximize. More text gives the model more to work with, but it costs more, and extra notes make it easier for the next attempt to drift. Too little and the loop cannot move.

An agent that can explore freely writes better kernels. A strict outer system has to define the feedback and decide what is allowed to ship. A good design needs both.

Generated individual Qwen 3.5 122B kernels were 1.8–5.2× faster than the best implementations available in vLLM.

Traditionally, production inference systems rely on generic kernels to handle diverse models and workloads. This is suboptimal because GPU operation shapes are determined by a combination of static model parameters and dynamic request-time factors; for instance, while a model defines one of the dimensions for a matrix multiplication, the other dimension fluctuates based on the specific token count of each request. There is growing interest in agentic GPU kernel generation, and recent efforts have shown promise. We explored a core question: if kernel generation can be automated, why should models of vastly different sizes (from 1 billion to 1 trillion parameters) rely on the same kernel? By specializing kernels to the specific shapes encountered at runtime, we can achieve extreme efficiency. In this blog, we share our successes and insights from using agents to generate GPU kernels. We built Proteus, a system designed to achieve extreme specialization, which requires a harness tailored for rigorous optimization, validation, and context management.

Figure 1: Simplified view of the Proteus harness

Conventional coding harnesses often fail here because agents tend to reward-hack: following the letter of the law rather than the spirit. If you give an agent a benchmark, it may optimize the benchmark and not the intended operation. To address this, Proteus proposes kernels, verifies them against a controlled reference implementation, times the successful ones, and iteratively improves upon the best results. While the process is straightforward, its success depends entirely on solving two foundational challenges. Figure 1 shows the simplified architecture of our design. Using our Proteus harness, we generated Qwen 3.5 122B kernels that were 1.8–5.2× faster than the best available in vLLM. Validation We originally treated kernel search as the hard part: how to explore a large space of programs without getting stuck in a plateau without improving? In practice the first question was more basic. Are we measuring what we think we are measuring? A model optimizes the score you give it. It does not need an exotic exploit: it may simply be that the evaluation is making an assumption. One example was kernels for rotary position embeddings (RoPE), a common step in attention layers. A candidate could reuse compiled code left over from an earlier attempt and look cheaper than a fair rebuild from scratch. Another could record a batch of GPU launches in a graph (e.g., CUDA graph) and replay them as one unit, while the baseline we compared against still launched each piece separately, so the two sides were not doing the same work. Another was strong on the input sizes we had put in the visible test set and weak on sizes it had not been shown. So we spent early design work on the checker, not the prompt. Time both sides the same way, including with more than one timer (e.g., CUDA event timer, wall clock time and CUPTI timer) when we need a cross-check. Clear leftover compiled state that should not persist, and keep the order of setup and teardown consistent so one side cannot skip work the other still pays for. Time the winners again before using them as the starting point for the next round. Keep some tests the candidate cannot see, so it cannot fit only the exam. To prevent evaluation "cheating" with artificially inflated performance, we implement automated consistency checks to flag theoretically impossible speedups (e.g., >100x) that exceed physical GPU bandwidth and compute limits. This protects against the same reward-hacking pitfalls seen in past industry cases, where agents optimized for the harness metrics rather than genuine performance gains. Without these constraints, generating more kernels mostly produced more noise. Emphasis on the checker also changes the bottleneck of agentic kernel generation. In just program-search (i.e., iterative optimization where the system searches over programs by repeatedly generating variants) work, good candidates are rare, so writing them dominates the cost. We can produce many drafts in parallel, but we cannot skip validation. We have to craft the validation carefully, and checking has to run on real GPUs, in isolation, and more than once. The system moves as fast as it can trust a kernel, not as fast as it can write one.

Figure 2: Token usage among different phases by our initial harness.

Figure 3: Token usage among different phases by our improved harness.

Context Management Another challenge is determining what the kernel generation model is allowed to see. It's a trade-off. Give the model a larger prompt and it has more information: the current best kernel, recent failures, profiler hints, notes from earlier runs. That can help. It also costs more, because we pay for every token the model reads. And as the prompt grows, it is easier for the next attempt to drift. Useful signals are mixed with stale advice, conflicting tips, and details that apply to a different input size or a different operation. The model does not always know which sentences to trust, so it follows the loudest ones, or all of them a little. Give it too little and the opposite happens. Every attempt starts from zero. The same dead ends come back. Nothing carries over from the last run, or from a related operation, and the loop does not advance. We wanted a knowledge layer to help with that: remember what worked, reuse it later, and do it without a person in the loop. That layer has a second tradeoff, between...