KernelBrain: Coarse-to-Fine, Budget-Aware Search for Agentic GPU Kernel Optimization
This work provides a practical solution for automating GPU kernel optimization, which is a difficult problem for developers seeking to improve performance of their applications.
This paper presents KernelBrain, an agent that optimizes GPU kernels by combining LLM-guided mutation, adaptive resource allocation, policy-gated evaluation, and profiler-informed diagnosis. It achieves 0.88x-6.72x speedup over PyTorch and up to 1.4x speedup over the state-of-the-art kernel agent, with up to 48% lower optimization time.
Automating GPU kernel optimization remains difficult in practice: generated variants can violate correctness constraints, runtime measurements are noisy, and search often stalls early. We present a practical optimization agent that combines LLM-guided mutation, adaptive resource allocation, policy-gated evaluation, and profiler-informed diagnosis. The system screens many candidates with low-cost evaluation and allocates higher-fidelity budget only to promising survivors to optimize and evolve GPU kernels. On important Triton kernel generation tasks, this design improves both kernel quality and search efficiency, reaching 0.88x-6.72x speedup over PyTorch and up to 1.4x speedup over the state-of-the-art kernel agent, with up to 48% lower optimization time.