Code 2025
AI Kernel Generation: What's working, what's not, what's next – Natalie Serrino, Gimlet Labs
Overview
This talk explores the use of AI for generating and optimizing low-level computational kernels, which are crucial for the performance of complex AI workloads. The core idea is to leverage AI to automatically port and optimize these kernels for diverse hardware, addressing the scarcity of human experts in this specialized field. While AI shows promise in tasks like kernel fusion and translating optimizations, it's not a replacement for human expertise in highly complex algorithmic advancements.
Who should watch
- AI engineers and researchers working on performance optimization.
- Builders of agentic systems that involve complex, multi-stage pipelines.
- Product managers and engineers focused on shipping efficient AI products.
- Anyone interested in the intersection of AI and low-level systems programming.
- Individuals facing challenges in finding or affording specialized kernel optimization experts.
Key takeaways
- AI-driven kernel generation aims to automate the optimization of computational kernels for various hardware, addressing a bottleneck in AI performance.
- Optimizing kernels can lead to significant performance gains, such as a 3x throughput increase for a Llama model mentioned in an Nvidia blog.
- The process of kernel optimization is iterative and complex, involving compilation, correctness checks, and deep profiling, which AI can potentially assist with.
- Challenges in AI kernel synthesis include defining floating-point correctness, selecting appropriate input sizes for benchmarking, and reliably measuring performance without overhead.
- Preliminary results show AI agents achieving average speedups of around 25% on benchmarks, with a sweet spot in moderately complex problems.
- Successful AI optimizations include kernel fusion, where multiple operations are combined into a single function, and re-expressing problems using more optimized hardware-specific operations.
- AI struggles with highly optimized, human-expert-tuned operations like matrix multiplication and can sometimes produce results that are technically correct but not aligned with the intended benchmarking goals.
- A robust AI system for kernel optimization requires a human in the loop for supervision, a purpose-built harness for task management, and multiple sub-agents for idea generation and verification.
Notable quotes
*The thing that we've seen with all these talks so far is with agents, they're not just one chat model. They're complex pipelines of multiple models, multiple stages, tool calls, and the compute backing these is inherently should be heterogeneous.*
*So the idea here is to put AI as the as the kind of like where the human would go in that same loop right so the agentic flow here is to make sure it compiles and it executes and it's correct and then from there optimizing it.*
*This is not a silver bullet and I think that's really important to emphasize.*
Unofficial community note. Prefer the recording for nuance.