Towards Cold-Start Drafting and Continual Refining: A Value-Driven Memory Approach with Application to NPU Kernel Synthesis
作者: Yujie Zheng, Zhuo Li, Shengtao Zhang, Hanjing Wang, Junjie Sheng, Jiaqian Wang, Junchi Yan, Weinan Zhang, Ying Wen, Bo Tang, Muning Wen
分类: cs.LG, cs.AI, cs.CL
发布日期: 2026-03-11
💡 一句话要点
EvoKernel:面向NPU内核合成的价值驱动记忆方法,实现冷启动和持续优化
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: NPU内核合成 冷启动学习 强化学习 价值驱动记忆 领域专用架构
📋 核心要点
- 现有方法在数据丰富的平台(如CUDA)上表现出色,但在数据稀缺的NPU编程等生态系统中性能急剧下降,面临冷启动挑战。
- EvoKernel提出了一种自演进的Agent框架,通过价值驱动的记忆检索机制,学习特定阶段的Q值,从而优化内核合成过程。
- 实验结果表明,EvoKernel显著提高了NPU内核合成的正确率和速度,证明了价值引导的经验积累的有效性。
📝 摘要(中文)
本文提出EvoKernel,一个自演进的Agent框架,旨在解决大语言模型在数据稀缺的编程领域,特别是新兴领域专用架构(如NPU)上的内核合成问题。EvoKernel将内核合成过程建模为基于记忆的强化学习任务,通过新颖的价值驱动检索机制,学习特定阶段的Q值,从而优先考虑对当前目标(引导可行草案或迭代优化延迟)有贡献的经验。此外,通过启用跨任务记忆共享,Agent能够将从简单算子中获得的见解推广到复杂算子。在NPU版本的KernelBench上的评估表明,EvoKernel将前沿模型的正确率从11.0%提高到83.0%,并通过迭代优化实现了相对于初始草案3.60倍的中值加速。这表明,价值引导的经验积累使通用模型能够掌握小众硬件生态系统上的内核合成任务。
🔬 方法详解
问题定义:论文旨在解决在大语言模型应用于新兴领域专用架构(如NPU)的内核合成时,由于训练数据稀缺而导致的冷启动问题。现有模型在数据丰富的平台(如CUDA)上表现良好,但在数据稀缺的NPU编程等生态系统中性能显著下降,无法直接应用。
核心思路:论文的核心思路是将内核合成过程建模为一个基于记忆的强化学习任务,并设计一个自演进的Agent框架EvoKernel。EvoKernel通过价值驱动的检索机制,学习不同阶段的Q值,从而优先考虑对当前目标(例如,生成可行的初始草案或迭代优化延迟)有贡献的经验。这种方法允许模型从少量数据中学习,并逐步改进内核合成的质量。
技术框架:EvoKernel框架包含以下主要模块:1) 记忆模块,用于存储和检索经验;2) 价值评估模块,用于学习不同阶段的Q值;3) 策略优化模块,用于根据Q值选择合适的动作,例如,生成新的内核代码或修改现有代码。整个流程从一个初始的内核草案开始,通过不断地迭代和优化,最终生成高性能的NPU内核。
关键创新:EvoKernel的关键创新在于其价值驱动的检索机制。传统的记忆回放方法通常随机采样经验,而EvoKernel则根据经验对当前目标的贡献程度(即Q值)来选择经验。这种方法可以更有效地利用有限的训练数据,并加速模型的学习过程。此外,EvoKernel还支持跨任务的记忆共享,允许模型将从简单算子中获得的知识迁移到复杂算子上。
关键设计:EvoKernel使用深度Q网络(DQN)来学习Q值。Q网络的输入包括当前的状态(例如,当前的内核代码和目标延迟)和可能的动作(例如,修改内核代码的各种方式)。Q网络的输出是每个动作的Q值,表示执行该动作后获得的预期回报。EvoKernel使用一种特殊的损失函数来训练Q网络,该损失函数考虑了经验的价值和重要性。此外,EvoKernel还使用了一种探索策略,以鼓励模型探索新的动作和状态。
📊 实验亮点
EvoKernel在NPU版本的KernelBench上进行了评估,结果表明,EvoKernel将前沿模型的正确率从11.0%提高到83.0%,并通过迭代优化实现了相对于初始草案3.60倍的中值加速。这些结果表明,EvoKernel能够有效地解决NPU内核合成中的冷启动问题,并显著提高内核的性能。
🎯 应用场景
EvoKernel的应用场景广泛,包括各种领域专用架构(DSA)上的内核合成、编译器优化、以及自动代码生成等。该研究成果有助于降低在新型硬件平台上开发高性能软件的门槛,加速人工智能算法在各种设备上的部署,并推动领域专用计算的发展。
📄 摘要(原文)
Deploying Large Language Models to data-scarce programming domains poses significant challenges, particularly for kernel synthesis on emerging Domain-Specific Architectures where a "Data Wall" limits available training data. While models excel on data-rich platforms like CUDA, they suffer catastrophic performance drops on data-scarce ecosystems such as NPU programming. To overcome this cold-start barrier without expensive fine-tuning, we introduce EvoKernel, a self-evolving agentic framework that automates the lifecycle of kernel synthesis from initial drafting to continual refining. EvoKernel addresses this by formulating the synthesis process as a memory-based reinforcement learning task. Through a novel value-driven retrieval mechanism, it learns stage-specific Q-values that prioritize experiences based on their contribution to the current objective, whether bootstrapping a feasible draft or iteratively refining latency. Furthermore, by enabling cross-task memory sharing, the agent generalizes insights from simple to complex operators. By building an NPU variant of KernelBench and evaluating on it, EvoKernel improves frontier models' correctness from 11.0% to 83.0% and achieves a median speedup of 3.60x over initial drafts through iterative refinement. This demonstrates that value-guided experience accumulation allows general-purpose models to master the kernel synthesis task on niche hardware ecosystems. Our official page is available at https://evokernel.zhuo.li.