用多智能体自动生成跨平台高性能计算内核,加速AI模型部署
AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis
- 构建多智能体系统,自动完成内核生成与迁移
- 在GPU和NPU上平均提速1.46倍,超越基础实现
- 支持Triton等语言,适合追求高效推理的工程师
现代AI模型需要高性能计算内核。随着大语言模型、多模态架构及推荐系统的复杂度上升,结合稀疏化和量化等技术,计算挑战日益严峻。同时,硬件频繁更新、芯片架构多样,要求为每种平台定制内核,人工优化难以跟上需求,成为AI系统开发的瓶颈。近期大模型代码生成能力的进步为自动化内核开发带来新可能。本文提出AKG kernel agent(AI驱动内核生成器),一个支持多领域特定语言(DSL)的多智能体框架,包括Triton、TileLang、CPP和CUDA-C,可适配不同硬件后端,保证正确性与可移植性。系统模块化设计便于快速集成新语言和硬件目标。在KernelBench上使用Triton DSL评估,该框架在GPU与NPU后端平均相较PyTorch Eager基线提升1.46倍性能,证明其对现代AI工作负载内核开发的有效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。