Kirill Nagaitsev CS PhD Prospectus: Shaping the Search in LLM-Driven Computational Kernel Optimization
Webcast Link (Hybrid)
Optimizing computational kernels for target GPU and accelerator hardware is an important, ongoing challenge. These kernels are the performance-critical units of AI/ML models and HPC applications, where small inefficiencies compound into large costs at scale. Traditional approaches include writing custom kernels manually, which requires significant time and expertise, as well as using compilers, which apply transformations from fixed catalogs. Recent work shows that LLM-based methods can also optimize kernels effectively, often outperforming existing compilers. PIKE, the first component of my thesis, demonstrates that in LLM-based multi-agent systems for PyTorch inference optimization, achieving a cost-effective search requires careful allocation of agents and search resources.
I propose extending PIKE to harness the strengths of tensor program superoptimization, a search-based technique that explores the space of programs within a fixed operator vocabulary to find the fastest one equivalent to the original kernel. I describe the design space that emerges from these hybrid approaches, in which a classical search technique is driven by an LLM. I further propose treating the kernel optimization search space itself as adaptable, using the LLM to modify the operator vocabulary that the superoptimizer searches over. I outline completed work on PIKE, initial steps toward integrating superoptimization as a classical search method driven by an LLM, and future directions for shaping the search in LLM-driven kernel optimization.
Jensen Smith
Email