Teaching LLMs to Write System Kernels for AI Accelerators: Post-Training, Reasoning, and Agentic Optimization
TL;DR
This tutorial presents emerging methods that leverage large language models (LLMs) to automate kernel generation for GPUs and custom accelerators such as AWS Trainium, and demonstrates how recent advances in reinforcement learning and inference-time compute scaling can be adapted to produce correct and performant kernels across diverse hardware backends.
Abstract
Writing high-performance system kernels for AI accelerators requires deep hardware expertise and low-level programming skill—a significant barrier for many researchers and engineers. This tutorial presents emerging methods that leverage large language models (LLMs) to automate kernel generation for GPUs and custom accelerators such as AWS Trainium. We cover the full pipeline: supervised fine-tuning on kernel corpora, reinforcement learning-based post-training (PPO, GRPO) with hardware-aware rewards, inference-time scaling for code reasoning, and multi-turn agentic systems for iterative kernel refinement. We also discuss key challenges unique to this domain, including reward hacking in compilation metrics, designing effective evaluation benchmarks for generated kernels, and strategies for data curation across heterogeneous hardware targets. The tutorial bridges the gap between LLM post-training research and systems programming, offering both foundational understanding and practical experience. We demonstrate how recent advances in reinforcement learning and inference-time compute scaling can be adapted to produce correct and performant kernels across diverse hardware backends. Attendees will gain hands-on experience generating Triton and NKI kernels using agentic LLM systems. Tutorial materials are available at https://neuron-science.github.io/llm_kernel_writing.