用多智能体框架让GPU核优化更高效可解释
KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization
- 设计双层记忆多智能体系统，整合专家优化技能
- 在KernelBench上实现最高5.44倍加速，成功率100%
- 适合需要高效推理和可解释性的AI系统开发者
提升GPU核效率对推动AI系统发展至关重要。现有基于大语言模型的核优化方法依赖模型内部隐式学习的启发式策略，导致试错效率低且优化过程难以解释。本文提出KernelSkill，一个具备双层记忆架构的多智能体框架，通过协调具有长期可复用专家技能记忆和短期记忆的智能体，避免重复回溯。在KernelBench Level 1-3上，KernelSkill分别实现相对于Torch Eager的平均5.44x、2.82x、1.92x加速，成功率100%，显著优于现有基线方法。代码已开源。
原文摘要 · Abstract (English)
Improving GPU kernel efficiency is crucial for advancing AI systems. Recent work has explored leveraging large language models (LLMs) for GPU kernel generation and optimization. However, existing LLM-based kernel optimization pipelines typically rely on opaque, implicitly learned heuristics within the LLMs to determine optimization strategies. This leads to inefficient trial-and-error and weakly interpretable optimizations. Our key insight is to replace implicit heuristics with expert optimization skills that are knowledge-driven and aware of task trajectories. Specifically, we present KernelSkill, a multi-agent framework with a dual-level memory architecture. KernelSkill operates by coordinating agents with long-term memory of reusable expert skills and short-term memory to prevent repetitive backtracking. On KernelBench Levels 1-3, KernelSkill achieves a 100% success rate and average speedups of 5.44x, 2.82x, and 1.92x over Torch Eager on Levels 1, 2, and 3, respectively, outperforming prior baselines. Code is available at https://github.com/0satan0/KernelMem/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。