English

AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis

Artificial Intelligence 2025-12-30 v1 Machine Learning

Abstract

Modern AI models demand high-performance computation kernels. The growing complexity of LLMs, multimodal architectures, and recommendation systems, combined with techniques like sparsity and quantization, creates significant computational challenges. Moreover, frequent hardware updates and diverse chip architectures further complicate this landscape, requiring tailored kernel implementations for each platform. However, manual optimization cannot keep pace with these demands, creating a critical bottleneck in AI system development. Recent advances in LLM code generation capabilities have opened new possibilities for automating kernel development. In this work, we propose AKG kernel agent (AI-driven Kernel Generator), a multi-agent system that automates kernel generation, migration, and performance tuning. AKG kernel agent is designed to support multiple domain-specific languages (DSLs), including Triton, TileLang, CPP, and CUDA-C, enabling it to target different hardware backends while maintaining correctness and portability. The system's modular design allows rapid integration of new DSLs and hardware targets. When evaluated on KernelBench using Triton DSL across GPU and NPU backends, AKG kernel agent achieves an average speedup of 1.46×\times over PyTorch Eager baselines implementations, demonstrating its effectiveness in accelerating kernel development for modern AI workloads.

Keywords

Cite

@article{arxiv.2512.23424,
  title  = {AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis},
  author = {Jinye Du and Quan Yuan and Zuyao Zhang and Yanzhi Yi and Jiahui Hu and Wangyi Chen and Yiyang Zhu and Qishui Zheng and Wenxiang Zou and Xiangyu Chang and Zuohe Zheng and Zichun Ye and Chao Liu and Shanni Li and Renwei Zhang and Yiping Deng and Xinwei Hu and Xuefeng Jin and Jie Zhao},
  journal= {arXiv preprint arXiv:2512.23424},
  year   = {2025}
}
R2 v1 2026-07-01T08:44:15.136Z