Post-Training, Reasoning, and Agentic Optimization
224 Jungmungwangwang-ro, Seogwipo-si, Jeju-do, South Korea
Rajarshi Saha is an Applied Scientist at Amazon Annapurna Labs, where he focuses on developing algorithms for resource-efficient training and inference of large foundation models. His research broadly explores theoretical approaches to challenges in this space, with an emphasis on investigating optimality and developing theoretically-backed system-aware algorithms. Before joining AWS, he earned a PhD in Electrical Engineering from Stanford University. Prior to that, he completed his Bachelors' and Masters' from IIT Kharagpur, where he received the Prime Minister of India Gold Medal as the class valedictorian.
Zhen Jia is a Senior Applied Scientist at Amazon Annapurna Labs, leading the system-for-ML and ML-for-system directions. His primary research interests span systems, high-performance computing, and deep learning compilers. He has contributed over 70 peer-reviewed papers in top computer system conferences and journals, including SOSP, EuroSys, HPCA, PPoPP, PACT, TPDS, TACO, and MLSys.
Jiin Woo is a PhD student in ECE at Carnegie Mellon University, advised by Professors Yuejie Chi and Gauri Joshi. Previously, she earned her M.S. in Electrical Engineering and B.S. in Mathematical Sciences from KAIST. Her research focuses on the foundations of collaborative reinforcement learning for system efficiency and robustness, with an emphasis on federated RL and RL for system kernel optimization.
Youngsuk Park is a Senior Applied Scientist and Research Lead at Amazon Annapurna Labs, where he leads the Core Algorithm team. His research focuses on efficient training and inference of large-scale foundation models, including low-precision training, optimization methods, and hardware-aware algorithms for AI accelerators such as AWS Trainium. He has authored 40+ papers on the topics of optimization, RL, LLM, and ML systems at top venues including ICLR, ICML, and AISTATS. He has delivered multiple tutorials at AAAI, IJCAI, KDD, and Big Data. He received his PhD in Electrical Engineering from Stanford University.
Kaan Ozkara
Applied Scientist
Amazon Annapurna Labs
Shaowei Zhu
Applied Scientist
Amazon Annapurna Labs
Lingfan Yu
Applied Scientist
Amazon Annapurna Labs
Wei Tang
Applied Scientist
Amazon Annapurna Labs
Ziyang Xu
Senior Applied Scientist
Amazon Annapurna Labs
Emily Webber
Principal Solutions Architect
Amazon Annapurna Labs
Yida Wang
Principal Scientist
Amazon Annapurna Labs
Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD '26).
BibTeX
@inproceedings{10.1145/3770855.3816468,
author = {Park, Youngsuk and Saha, Rajarshi and Ozkara, Kaan and Zhu, Shaowei and Woo, Jiin and Yu, Lingfan and Tang, Wei and Xu, Ziyang and Jia, Zhen and Webber, Emily and Wang, Yida},
title = {Teaching LLMs to Write System Kernels for AI Accelerators: Post-Training, Reasoning, and Agentic Optimization},
year = {2026},
isbn = {9798400722592},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3770855.3816468},
doi = {10.1145/3770855.3816468},
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.},
booktitle = {Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2},
pages = {13356–13359},
numpages = {4},
keywords = {llm, system kernels, post-training, reinforcement learning, kernel generation, ai accelerators, inference-time scaling, agentic systems, cuda, triton, nki, code optimization},
location = {Republic of Korea},
series = {KDD '26}
}