AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis
作者: Jinye Du, Quan Yuan, Zuyao Zhang, Yanzhi Yi, Jiahui Hu, Wangyi Chen, Yiyang Zhu, Qishui Zheng, Wenxiang Zou, Xiangyu Chang, Zuohe Zheng, Zichun Ye, Chao Liu, Shanni Li, Renwei Zhang, Yiping Deng, Xinwei Hu, Xuefeng Jin, Jie Zhao
分类: cs.AI, cs.LG
发布日期: 2025-12-29
💡 一句话要点
提出AKG内核代理以解决跨平台内核合成问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 内核生成 多代理系统 自动化优化 领域特定语言 跨平台支持
📋 核心要点
- 现有方法无法满足现代AI模型对高性能计算内核的需求，手动优化成为开发瓶颈。
- 提出AKG内核代理，通过多代理系统自动化内核生成、迁移和性能调优，支持多种DSL。
- 在KernelBench上评估时，AKG内核代理在GPU和NPU后端上实现了1.46倍的加速效果。
📝 摘要（中文）
现代AI模型对高性能计算内核的需求日益增加。大型语言模型、跨模态架构和推荐系统的复杂性，以及稀疏性和量化等技术，带来了显著的计算挑战。频繁的硬件更新和多样的芯片架构使得每个平台都需要定制的内核实现，而手动优化无法跟上这些需求，成为AI系统开发的瓶颈。本文提出了AKG内核代理（AI驱动的内核生成器），一个自动化内核生成、迁移和性能调优的多代理系统，支持多种领域特定语言（DSL），如Triton、TileLang、CPP和CUDA-C，能够针对不同硬件后端，同时保持正确性和可移植性。通过在KernelBench上使用Triton DSL评估，AKG内核代理在GPU和NPU后端上实现了比PyTorch Eager基线实现平均加速1.46倍，展示了其在现代AI工作负载内核开发中的有效性。
🔬 方法详解
问题定义：现代AI模型对计算内核的性能要求不断提高，现有的手动优化方法无法快速适应不同硬件平台的需求，导致开发效率低下。
核心思路：AKG内核代理通过多代理系统实现内核的自动生成、迁移和性能调优，旨在提高内核开发的效率和适应性，支持多种领域特定语言（DSL）。
技术框架：系统由多个模块组成，包括内核生成模块、迁移模块和性能调优模块，能够快速集成新的DSL和硬件目标。整体流程从接收用户需求开始，经过内核生成、优化到最终部署。
关键创新：AKG内核代理的核心创新在于其多代理架构和对多种DSL的支持，使其能够在不同硬件平台上保持高效性和可移植性，显著提升了内核开发的自动化水平。
关键设计：系统设计中考虑了多种硬件架构的特性，采用了模块化设计，允许快速集成新DSL，确保生成的内核在不同平台上都能高效运行。
🖼️ 关键图片
📊 实验亮点
在KernelBench的评估中，AKG内核代理在使用Triton DSL时，在GPU和NPU后端上实现了平均1.46倍的加速，相较于PyTorch Eager基线实现，展示了显著的性能提升，验证了其在现代AI工作负载中的有效性。
🎯 应用场景
AKG内核代理的潜在应用场景包括深度学习框架的内核优化、边缘计算设备的高效内核生成以及多模态AI系统的性能提升。其自动化特性将大幅降低开发成本，提高AI系统的适应性和性能，具有广泛的实际价值和未来影响。
📄 摘要（原文）
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.