用多智能体LLM自动优化GPU核函数,提升大模型推理性能。
Astra: A Multi-Agent System for GPU Kernel Performance Optimization
- 多个LLM智能体协作生成、测试、优化CUDA代码。
- 在SGLang代码上实现平均1.32倍加速,无需额外训练。
- 适合需要高性能计算的开发者和研究者使用。
GPU核函数优化是高性能计算与机器学习交叉领域的核心挑战。高效核函数对加速大语言模型(LLM)训练与服务至关重要,但通常需大量手动调优。编译器系统虽减轻部分负担,仍需大量人工设计。近期研究尝试用LLM生成GPU核函数,但多聚焦于将PyTorch模块转为CUDA代码。本文提出Astra,首个基于LLM的多智能体系统,用于GPU核函数优化。不同于以往方法,Astra从广泛部署的SGLang框架中提取现有CUDA实现作为起点,而非以PyTorch模块为输入。在Astra中,专用LLM智能体通过迭代生成、测试、性能分析与规划协作,产出正确且高效的核函数。在SGLang的核函数上,Astra采用零样本提示(zero-shot prompting)与OpenAI o4-mini,平均实现1.32倍加速。详细案例研究显示,LLM可自主应用循环变换、优化内存访问模式、利用CUDA内建函数及快速数学运算,带来显著性能提升。本工作凸显多智能体LLM系统在GPU核优化中的潜力。代码已开源:https://github.com/Anjiang-Wei/Astra。
原文摘要 · Abstract (English)
GPU kernel optimization has long been a central challenge at the intersection of high-performance computing and machine learning. Efficient kernels are crucial for accelerating large language model (LLM) training and serving, yet attaining high performance typically requires extensive manual tuning. Compiler-based systems reduce some of this burden, but still demand substantial manual design and engineering effort. Recently, researchers have explored using LLMs for GPU kernel generation, though prior work has largely focused on translating high-level PyTorch modules into CUDA code. In this work, we introduce Astra, the first LLM-based multi-agent system for GPU kernel optimization. Unlike previous approaches, Astra starts from existing CUDA implementations extracted from SGLang, a widely deployed framework for serving LLMs, rather than treating PyTorch modules as the specification. Within Astra, specialized LLM agents collaborate through iterative code generation, testing, profiling, and planning to produce kernels that are both correct and high-performance. On kernels from SGLang, Astra achieves an average speedup of 1.32x using zero-shot prompting with OpenAI o4-mini. A detailed case study further demonstrates that LLMs can autonomously apply loop transformations, optimize memory access patterns, exploit CUDA intrinsics, and leverage fast math operations to yield substantial performance gains. Our work highlights multi-agent LLM systems as a promising new paradigm for GPU kernel optimization. Our code is publicly available at https://github.com/Anjiang-Wei/Astra.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。