KernelBlaster: 通过记忆增强的上下文强化学习实现持续的跨任务CUDA优化
KernelBlaster: Continual Cross-Task CUDA Optimization via Memory-Augmented In-Context Reinforcement Learning
AI总结：
KernelBlaster通过记忆增强的上下文强化学习框架提升CUDA代码优化性能，实现跨多个GPU架构世代的高效优化。
摘要：
在多个GPU架构世代中优化CUDA代码具有挑战性，因为要实现峰值性能需要对一个日益复杂且硬件特定的优化空间进行广泛探索。传统编译器受到固定启发式方法的限制，而微调大型语言模型（LLMs）成本较高。然而，针对CUDA代码优化的智能工作流在整合先前探索知识方面能力有限，导致采样偏差和次优解决方案。我们提出了KernelBlaster，一种记忆增强的上下文强化学习（MAIC-RL）框架，旨在提高基于LLM的GPU编码代理的CUDA优化搜索能力。KernelBlaster使代理能够通过积累知识到可检索的持久CUDA知识库中，从而学习经验并在未来任务中做出系统性决策。我们提出了一种新颖的基于配置文件的、文本梯度的智能流程，用于CUDA生成和优化，以在多个GPU架构世代中实现高性能。KernelBlaster引导LLM代理系统性地探索高潜力的优化策略，而非简单的重写。与PyTorch基线相比，我们的方法在KernelBench级别1、2和3上分别实现了几何均值加速1.43倍、2.50倍和1.50倍。我们发布了KernelBlaster作为开源智能框架，附带测试工具、验证组件和可重复的评估流程。
英文摘要：
Optimizing CUDA code across multiple generations of GPU architectures is challenging, as achieving peak performance requires an extensive exploration of an increasingly complex, hardware-specific optimization space. Traditional compilers are constrained by fixed heuristics, whereas finetuning Large Language Models (LLMs) can be expensive. However, agentic workflows for CUDA code optimization have limited ability to aggregate knowledge from prior exploration, leading to biased sampling and suboptimal solutions. We propose KernelBlaster, a Memory-Augmented In-context Reinforcement Learning (MAIC-RL) framework designed to improve CUDA optimization search capabilities of LLM-based GPU coding agents. KernelBlaster enables agents to learn from experience and make systematically informed decisions on future tasks by accumulating knowledge into a retrievable Persistent CUDA Knowledge Base. We propose a novel profile-guided, textual-gradient-based agentic flow for CUDA generation and optimization to achieve high performance across generations of GPU architectures. KernelBlaster guides LLM agents to systematically explore high-potential optimization strategies beyond naive rewrites. Compared to the PyTorch baseline, our method achieves geometric mean speedups of 1.43x, 2.50x, and 1.50x on KernelBench Levels 1, 2, and 3, respectively. We release KernelBlaster as an open-source agentic framework, accompanied by a test harness, verification components, and a reproducible evaluation pipeline.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
