Skip to content
Book Open access

Teaching LLMs to Write System Kernels for AI Accelerators: Post-Training, Reasoning, and Agentic Optimization

Aug 2026 · Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2 · pp. 13356-13359 · 0 citations · 2 references

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.

Read PDF

Similar papers

Preprint Aug 2026

LEAP: Lean Environment-Feedback via Adaptive Pruning for Code RL in GPU Kernel Generation

This work introduces LEAP (Lean Environment-Feedback via Adaptive Pruning), a scalable and computationally efficient multi-turn RL framework optimized for low-level hardware accelerator alignment and proposes a Rank-Based Reward formulation, establishing a practical paradigm for low-level code RL.

Tankun Li, Zhi Chen, Yao-Hua Tang · 0 citations
Preprint Jul 2026

Triton for MTIA: Bridging the Programming Model Gaps for Custom AI Accelerators

This work presents the first production-scale application of Triton on a custom ML accelerator, MTIA-2i, developed by Meta, and develops a new compiler backend that targets it, introduces enhancements to TorchInductor code generation, and proposes minimal language extensions that expose MTIA-specific architectural feat...

Haishan Zhu, Domi Yan, Michael Levesque-Dion et al. · 0 citations
#artificial intelligence Preprint Sep 2026

MaxKernel: Agentic Kernel Generation for TPUs

This work presents MaxKernel, a multi-agent system that implements three distinct paradigms for TPU kernel development: a Human-in-the-Loop (HITL) agent for collaborative, step-by-step design; an Autonomous (Auto) agent that executes a fully automated, metric/trace-driven optimization loop; and a Graph-Based Autonomous...

Shang-Kun Wang, Nina Cai, Charles Hoong et al. · 0 citations
#natural language process... Preprint Sep 2026

AMDKernelVault: Large-Scale Datasets and Agentic Training for AMD GPU Kernel Optimization

We introduce AMDKernelVault, an open HIP and Triton kernel corpus and training framework for recent AMD CDNA GPUs. Existing LLM-based kernel agents are largely CUDA/NVIDIA-centric and often depend on repeated frontier-LLM calls for generation, reflection, and optimization. To address this gap, we develop HIPKernelGen a...

Ji Liu, S. Majumder, Yi-Qing Huang et al. · 0 citations
Review Aug 2026

AI Hardware Accelerators for Large Language Models: Architectures and the Memory Wall

This review surveys the full landscape of AI hardware accelerators for LLMs, including general-purpose GPUs, custom ASICs, reconfigurable FPGAs, processing-in-memory and near-memory architectures, and emerging neuromorphic and photonic approaches across cloud and edge deployment.

Siddharth Patel, Rohit Singh · 0 citations
Preprint Sep 2026

AccelForge: Comprehensive Modeling and Co-Design Framework for AI Accelerators

Tensor algebra workloads, of which deep neural networks are prominent examples, are energy-intensive workloads in modern datacenter and edge deployments, making accelerators necessary to achieve energy efficiency and high throughput. To quickly evaluate and iterate on accelerator designs, we need an accelerator modelin...

Tanner Andrulis, Michael Gilbert, Vivienne Sze et al. · 0 citations

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.