用AI自动生成高性能AMD GPU核函数,提升代码效率与可及性
Geak: Introducing Triton Kernel AI Agent & Evaluation Benchmarks
- 基于反思式推理循环的AI代理生成Triton核函数
- 正确率最高达63%,运行速度提升2.59倍
- 适合希望自动化优化GPU代码的研究者与工程师
AI生成GPU内核的需求迅速增长,尤其在产业界和学术界对可扩展、硬件优化解决方案的迫切需求推动下。随着深度学习工作负载日益复杂多样,自动化底层内核开发成为满足性能与生产效率要求的关键。主要云服务商、半导体公司和研究机构正大力投入基于AI的GPU代码生成,目标是在AMD MI300X等硬件上实现接近专家水平的性能,同时减少人工优化成本。作为兼顾性能与编码便捷性的Python类域特定语言,Triton已成为此类生成任务的热门目标。本文提出GEAK(生成高效人工智能导向的GPU内核)框架,利用前沿大模型生成适用于AMD GPU(包括MI300X和MI250)的高性能Triton代码,并采用推理时计算规模扩展与类Reflexion反馈机制的推理循环。在两个评估基准上,GEAK显著优于直接提示前沿大模型及基于Reflexion的生成流水线,正确率最高达63%,执行速度最高提升2.59倍。结果表明,这类智能体式代码生成方法具有加速多样化硬件平台应用推广、普及专家级内核性能的巨大潜力。
原文摘要 · Abstract (English)
The demand for AI-generated GPU kernels is rapidly growing, influenced by the need for scalable, hardware-optimized solutions in both industry and academia. As deep learning workloads grow in complexity and diversity, it is imperative to automate low-level kernel development to meet performance and productivity demands. Major cloud providers, semiconductor companies, and research institutions are now investing heavily in AI-driven code generation for GPUs, aiming to reduce manual optimization efforts while achieving near-expert performance on hardware like AMD MI300X. The Triton language, a Python-based DSL for GPU programming, has emerged as a popular target for such AI-generated kernels due to its balance of performance and ease-of-coding. In this work, we present an evaluation suite for Triton-based GPU kernels and GEAK (Generating Efficient AI-centric GPU Kernels)-a framework that leverages cutting-edge LLMs to generate performant Triton code specifically for AMD GPUs, including the AMD MI300X and MI250. GEAK leverages inference-time compute scaling to produce Triton-based GPU kernels using a reasoning loop adapted from Reflexion-style feedback mechanisms. On two evaluation benchmarks, GEAK significantly outperformed the baselines of directly prompting frontier LLMs as well as Reflexion-based generation pipelines by achieving correctness up to $63$% and execution speed up of up to $2.59$X. These results highlight the promise of GEAK-like agentic code generation for accelerating the adoption of diverse hardware platforms and democratizing access to expert-level kernel performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。