arXiv:2602.14293cs.LGcs.AI2026-02被引 13

用记忆增强强化学习,让代码生成模型持续优化跨代GPU的CUDA性能。

KernelBlaster: Continual Cross-Task CUDA Optimization via Memory-Augmented In-Context Reinforcement Learning

  • 构建可积累经验的持久化知识库,指导智能体系统探索优化策略。
  • 在多代GPU上实现1.43x~2.50x的加速,显著优于基线方法。
  • 适合需要长期迭代优化的GPU编程研究者与工程团队使用。

跨多代GPU架构优化CUDA代码极具挑战性,因达到峰值性能需探索日益复杂的硬件特定优化空间。传统编译器受限于固定启发式规则,而微调大语言模型成本高昂。基于LLM的代理在代码优化中缺乏对过往探索知识的聚合能力,导致采样偏差与次优解。本文提出KernelBlaster,一种内存增强的上下文强化学习(MAIC-RL)框架,提升基于LLM的GPU代码代理的优化搜索能力。通过将经验累积至可检索的持久化CUDA知识库,使代理能从历史中学习并系统化决策。我们设计了一种新的基于性能分析、文本梯度引导的代理流程,实现多代GPU架构下的高性能代码生成。KernelBlaster引导代理超越简单重写,系统探索高潜力优化策略。相比PyTorch基线,本方法在KernelBench Level 1、2、3上分别取得1.43x、2.50x、1.50x的几何平均加速。我们开源了KernelBlaster框架,附带测试工具、验证组件及可复现评估流水线。

原文摘要 · Abstract (English)

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.

CUDA优化强化学习代码生成智能代理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。