AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis
🤖 AI Summary
The growing demand for high-performance computing (HPC) kernels in AI—driven by large language models, multimodal architectures, and recommendation systems—exacerbates computational challenges, especially when combined with sparsity, quantization, and other optimizations. Concurrent hardware acceleration and architectural fragmentation further hinder manual kernel development, limiting efficiency and portability. Method: We propose the first multi-agent collaborative framework for AI kernel development, enabling DSL-agnostic kernel generation, cross-platform porting, and automated tuning. Its modular design supports rapid adaptation to new DSLs and hardware targets, integrating LLM-based code generation, multi-agent orchestration, DSL parsing, and compiler-level optimizations. Backend support includes Triton, TileLang, CUDA-C, and C++. Contribution/Results: Evaluated on KernelBench, our framework achieves 1.46× speedup over PyTorch Eager mode on average, while significantly improving developer productivity and cross-platform consistency.
Application Category
📝 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$ imes$ over PyTorch Eager baselines implementations, demonstrating its effectiveness in accelerating kernel development for modern AI workloads.
Problem
Research questions and friction points this paper is trying to address.
Automates kernel generation for diverse AI models and hardware
Addresses manual optimization bottlenecks in AI system development
Supports multiple DSLs for cross-platform kernel synthesis
Innovation
Methods, ideas, or system contributions that make the work stand out.
Multi-agent system automates kernel generation and tuning
Supports multiple DSLs for cross-platform hardware compatibility
Modular design enables rapid integration of new targets
🔎 Similar Papers
