Kernel-Smith:进化式内核优化的统一方案 Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization
进化搜索+逐步训练框架实现GPU内核自动优化,超越Gemini-3.0-pro
前置知识
GPU Kernel(GPU内核)
GPU内核是在GPU上并行执行的函数,是高性能计算的基础单元。现代深度学习框架(如PyTorch)的底层算子(矩阵乘法、卷积、归一化等)都需要高效内核来充分利用GPU的数千个计算核心。内核优化包括选择合适的分块(tiling)策略、融合(fusion)操作、内存访问模式等,直接影响训练和推理的吞吐量。
本文的核心目标就是自动生成高性能GPU内核,理解内核的基本概念是理解整个工作的前提。
Triton编程语言
Triton是OpenAI开发的领域特定语言(DSL),用于编写GPU内核。相比直接编写CUDA,Triton提供了更高层的抽象,允许开发者在更友好的编程模型中表达分块、内存访问等优化策略,同时保持接近手写CUDA的性能。Triton编译器会将高级描述转换为高效的GPU指令。
本文在NVIDIA GPU上使用Triton作为目标后端,生成的内核代码就是Triton代码。
进化算法(Evolutionary Algorithm)
进化算法是一类受生物进化启发的优化方法,通过维护一个种群(population)的候选解,在每一代中通过选择(selection)、变异(mutation)和重组(recombination)操作产生新解。优秀个体被保留在档案(archive)中,种群在多代迭代中逐步改进。在代码优化场景中,每个候选解就是一个可执行程序,适应度由编译、正确性和性能指标决定。
Kernel-Smith采用进化搜索框架来探索内核优化空间,这是与传统单轮或多轮对话式优化的本质区别。
GRPO(Group Relative Policy Optimization)
GRPO是一种强化学习算法,通过在每个数据点上采样多个候选输出,然后用组内相对排名作为奖励信号来更新策略模型。相比PPO等需要独立价值网络的算法,GRPO直接利用同一组样本之间的相对表现差异,更高效且稳定。在本文中,每个数据点采样8个候选内核,以相对于父代码的加速比作为奖励。
GRPO是Kernel-Smith训练流程中RL阶段的核心算法,理解它才能理解模型是如何从进化轨迹中学习的。
KernelBench
KernelBench是评估LLM驱动内核生成的标准基准,任务是将PyTorch参考实现替换为更快的GPU内核。评估指标包括正确性(corr)、快速比例(fastp,即产生加速的内核百分比)和平均加速比(avg_amsr)。它按难度分为三个级别:Level 1(简单)、Level 2(中等)、Level 3(困难)。
Kernel-Smith的主要评估基准就是KernelBench,论文的核心结果都在此基准上报告。
MAP-Elites
MAP-Elites(Multi-dimensional Archive of Phenotypic Elites)是一种质量-多样性算法,不同于传统进化算法只追求最高适应度,它同时维护解的多样性。它在特征空间(如复杂度、性能评分)的网格中保存每个区域的最优解,确保搜索过程不会过早收敛到单一模式。在本文中,档案按内核复杂度和综合评分(编译+正确性+加速)组织。
Kernel-Smith的进化框架借鉴了MAP-Elites的思想来平衡探索与利用,这是搜索效率的关键。
研究动机
高性能GPU内核的生成对于大模型系统和科学计算工作负载至关重要,但现有的基于LLM的方法仍然难以超越一次性代码生成的局限。具体来说,当前方法面临两个耦合的挑战:第一,高效内核通常需要在多种实现选择(包括不同的融合模式、分块策略和重写方向)中进行搜索才能出现,但现有系统的多轮精炼或历史条件化agent循环虽然有助于局部调试,却容易将后续提案锚定到早期决策上,限制了探索多样性。第二,功能正确性和高性能是两种不同的能力。当性能分析噪声较大时,搜索可能保留次优内核或淘汰真正有前途的内核,而这些错误会在多代搜索中累积。例如,GPU环境中算子执行的wall-clock时间即使在固定硬件和驱动版本下也会出现不可忽略的波动,对于小规模输入张量,内核启动时间占比过高,导致波动尤为明显。
本文的目标是本文的具体目标是构建一个统一框架Kernel-Smith,将稳定的评估驱动进化agent与专门针对进化搜索设计的训练方案相结合。在agent端,系统维护一个可执行候选程序的种群,使用高性能和多样性程序的档案以及结构化的执行反馈(编译、正确性和加速)来迭代改进候选内核。在训练端,将长期进化轨迹转换为以单步为中心的监督微调和强化学习信号,通过保留正确性保持的高增益修订来训练模型成为进化循环中的强局部改进器。最终目标是在KernelBench基准上达到最先进的整体性能,并在真实推理引擎(如SGLang和LMDeploy)中产生可部署的上游贡献。
与已有工作不同的是,本文的独特切入角度体现在两个核心设计选择上。首先,在agent设计方面,不同于现有的多轮对话式精炼(如PRAGMA、CUDA Agent等),Kernel-Smith采用进化搜索范式,维护候选程序种群而非单条轨迹,从而更好地利用测试时计算资源。关键创新是将评估稳定性作为agent设计的中心,通过固定计算图、重复测量和离群值去除来抑制计时噪声,将执行时间波动控制在1%以内。其次,在训练方面,不同于优化一次性内核生成的常规SFT,本文将长期进化轨迹转化为以单步为中心的训练信号。具体来说,它只保留高增益的修订步骤作为训练样本,这种过滤策略作为一种轨迹压缩形式,让模型学习对最终加速贡献最大的原子改进,而非模仿每一个中间转换。这一视角弥合了'生成正确内核'与'迭代优化内核'之间的能力鸿沟。
核心方法
Kernel-Smith的整体方法可以分为两大支柱:进化agent框架和面向进化的训练方案。直觉上,内核优化是一个高度非凸的搜索空间,单次生成很难找到全局最优解。进化搜索通过维护候选程序种群,允许性能增益在多轮搜索中累积,天然适合这类任务。但进化搜索的可靠性高度依赖评估的稳定性——如果性能分析噪声大,搜索会做出错误的淘汰和保留决策。因此,Kernel-Smith首先构建了一个稳定的评估后端(包括Triton和MACA两个异构平台),通过预热执行、多次测量、离群值去除和CUDAGraph技术将计时波动控制在1%以内。在此基础上,进化agent在每一步都接收参考实现和从档案中采样的候选程序,生成新的内核候选,并获得结构化的执行反馈(编译状态、正确性结果、加速比、运行时测量、硬件元数据和错误日志)。训练方案则从进化轨迹中提取关键改进步骤,通过SFT冷启动和GRPO强化学习来训练模型成为进化循环中的'强局部改进器'。
Kernel-Smith的核心创新在于将模型训练与进化搜索过程深度耦合,而非独立地优化代码生成能力。与已有方法的本质区别体现在三个层面:第一,评估稳定性作为首要设计原则。现有系统(如Astra、CudaForge等)的进化或迭代搜索依赖不稳定的性能分析,导致搜索质量退化。Kernel-Smith通过预热、多次测量、CUDAGraph等技术将波动控制在1%,保证了进化搜索的可靠性。第二,'最佳步骤选择'训练策略。实验发现,使用进化过程的所有步骤训练会导致信息泄露的捷径学习(模型从输入prompt中看到更好的内核示例后直接记忆),只使用第一步又因分布偏移和任务过于简单而效果不佳。而选择进化过程中的'最佳步骤'(既有一定基线加速又需要进一步优化的步骤)作为训练样本,能在任务复杂度和学习效果之间取得最佳平衡。第三,以加速比作为RL奖励信号。不同于代码风格或语义匹配等软奖励,加速比是一个可验证的、客观的性能指标,直接优化这个指标让模型学会真正提升性能的编辑操作。
方法步骤详情
Kernel-Smith的完整工作流程包含以下步骤。第一步,数据合成:从GitHub爬取多样化的PyTorch仓库,通过静态分析管道提取nn.Module子类,递归解析文件内依赖并内联必要组件,使每个示例自包含。经过嵌入和图去重、LLM辅助测试生成和执行过滤后,得到覆盖20个功能家族的59k高质量模块。第二步,冷启动数据生成:使用DeepSeek-V3.2-Speciale作为教师模型在Kernel-Smith框架中运行,生成轨迹数据后用正确性和加速指标过滤,保留功能有效且性能改进的样本。第三步,聚类种子专家数据:对数据集进行嵌入并应用HDBSCAN聚类,识别代表性聚类中心进行人工清洗和专家标注,然后将这些专家操作符反馈到Kernel-Smith中进行额外的rollout轮次。第四步,SFT冷启动:将多轮进化轨迹分解为单轮训练样本,使用正确性导向过滤(初始翻译步骤,保留所有功能正确的输出)和性能导向过滤(迭代进化步骤,仅保留加速比$> 1.0$的样本),最终得到超过200k高质量单轮样本,使用64k上下文长度进行SFT训练。第五步,RL训练:在聚类种子专家数据上,每个问题使用Gemini-3.0-pro进行40轮进化精炼,选择'最佳步骤'构建训练集。使用GRPO算法,每个数据条目采样8个候选,以相对于父代码的加速比作为奖励信号。第六步,推理时的进化搜索:在KernelBench上运行40轮迭代进化,每轮生成候选内核,通过评估后端验证编译、正确性和加速,使用MAP-Elites档案组织搜索空间。
技术新颖性
Kernel-Smith的技术新颖性体现在多个方面。首先,它是首个将评估稳定性提升为一等公民的内核生成框架。虽然进化算法(如AlphaEvolve)已在代码优化中应用,但现有工作未充分考虑GPU性能分析的噪声问题。Kernel-Smith通过固定计算图、多次测量、离群值去除和CUDAGraph等系统级优化,将计时波动控制在1%以内,这一工程贡献对进化搜索的可靠性至关重要。其次,'逐步中心'的训练范式是新颖的。不同于模仿整个轨迹的SFT或对最终结果的端到端RL,本文从进化轨迹中识别关键的原子改进步骤,将长期优化问题分解为可学习的单步操作。这种轨迹压缩策略让模型学习最有价值的编辑模式,而不是冗余的中间转换。第三,异构平台的统一评估协议是另一个创新。通过后端解耦设计,将任务规范、执行编排和指标计算与设备特定的编译和运行时接口分离,同一评估协议可以在NVIDIA GPU(Triton)和MetaX GPU(MACA)上复用,甚至可以扩展到华为NPU等平台。最后,从真实代码库(而非标准基准)爬取训练数据,并通过静态分析管道规范化,构建了覆盖20个功能家族的59k高质量模块数据集,这种数据合成方法论本身也具有独立价值。
实验结果
Kernel-Smith在多个平台和基准上展示了全面的性能提升。在NVIDIA GPU的KernelBench基准上,Kernel-Smith-235B-RL达到了96.33%的平均正确率(corr),仅略低于Claude-4.6-opus的99.33%,但显著超过所有其他基线(如DeepSeek-v3.2-Speciale的94.67%、Gemini-3.0-pro的94.33%)。更重要的是,它以3.70的平均加速比(avg_amsr)建立了新的最优水平,在中等难度子集(Level 2)上达到7.77,远超Claude-4.6-opus的5.83。在困难任务(Level 3)上,模型维持了94%的正确率,超过Gemini-3.0-pro和所有开源基线。在MetaX平台的MACA内核生成任务上,Kernel-Smith-MACA-30B以13.27的avg_amsr超过了DeepSeek-v3.2-think(8.01)和Qwen3-235B-2507-think(12.30),而235B版本进一步提升到14.26。在真实应用中,SGLang的融合内核实现了4.78倍加速,LMDeploy的MoE路由内核实现1.36倍加速并带来1.85%-3.00%的端到端吞吐量提升,DeepSeek Engram内核实现了14.59倍加速。关键的进化轨迹对比显示,Kernel-Smith-235B-RL的最佳分数增长曲线在整个搜索过程中形成了其他模型的上包络线,表明该模型能更有效地利用额外的测试时计算资源。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| KernelBench Level 1 (简单) | avg_amsr | 2.30 | Gemini-3.0-pro: 2.46 | 与最强专有模型基本持平 |
| KernelBench Level 2 (中等) | avg_amsr | 7.77 | Claude-4.6-opus: 5.83 | 提升33.3% |
| KernelBench Level 3 (困难) | avg_amsr | 1.02 | Claude-4.6-opus: 2.02 | 略低于最强专有模型 |
| KernelBench 整体 | avg_amsr | 3.70 | Claude-4.6-opus: 3.33 | 提升11.1%,整体最优 |
| KernelBench 整体 | corr | 96.33 | Claude-4.6-opus: 99.33 | 仅次于最强专有模型,超过所有开源模型 |
| MACA内核生成整体 | avg_amsr | Kernel-Smith-MACA-30B: 13.27, 235B: 14.26 | DeepSeek-v3.2-think: 8.01, Qwen3-235B: 12.30 | 30B版本即超过235B基线,235B版本进一步领先 |
| SGLang内核(真实应用) | speedup | 4.78倍 | PyTorch eager baseline | 端到端延迟改善0.28%-1.75% |
| LMDeploy MoE路由内核 | speedup | 1.36倍 | PyTorch eager baseline | 端到端吞吐量提升1.85%-3.00% |
| DeepSeek Engram内核 | speedup | 14.59倍 | PyTorch eager baseline | 最大真实应用场景加速 |
局限与改进
尽管Kernel-Smith取得了显著成果,本文存在几个值得关注的局限性。首先,在困难任务(Level 3)上,模型的avg_amsr仅为1.02,远低于Claude-4.6-opus的2.02,说明对于高度复杂的内核优化,进化搜索和训练方案仍有较大改进空间。其次,correctness指标(96.33%)仍低于最强专有模型(99.33%),虽然差距不大,但在生产环境中任何正确性问题都可能导致严重后果。第三,论文报告的内核级加速(如4.78倍、14.59倍)在转化为端到端系统性能时会大幅衰减——SGLang的端到端延迟改善仅0.28%-1.75%,LMDeploy的吞吐量提升仅1.85%-3.00%。这反映了内核优化在整体系统性能中的占比有限。第四,训练数据的合成依赖DeepSeek-V3.2-Speciale和Gemini-3.0-pro等强教师模型,数据质量和多样性受限于这些模型的能力。第五,论文缺乏与其他进化式内核优化方法(如AlphaEvolve本身、KernelSkill、K-Search等)的直接对比,使得Kernel-Smith各组件的贡献难以精确量化。最后,评估只在NVIDIA H200和MetaX平台上进行,对于其他硬件平台(如AMD GPU、华为NPU)的效果尚不清楚。
独立分析的弱点
独立分析Kernel-Smith的弱点,可以识别出以下几个方面。第一,进化搜索的计算开销较大:40轮迭代进化,每轮需要编译、运行和评估候选内核,加上100次单元测试来报告平均性能,总计算成本显著高于单轮生成方法。对于资源有限的研究团队,这可能是一个实际障碍。改进方向包括开发更高效的早期剪枝策略(如基于快速编译检查的预筛选)和自适应终止条件(当搜索收敛时提前停止)。第二,模型在困难任务上的表现仍有较大提升空间(Level 3的avg_amsr仅1.02),可能是因为复杂内核的优化空间太大,进化搜索难以充分探索。改进方向包括引入更细粒度的变异操作(如针对特定硬件特性的编辑)和层次化搜索(先优化算法层面再优化实现层面)。第三,训练依赖强教师模型生成进化轨迹,存在'天花板效应'——学生模型难以超越教师的能力。改进方向包括自我对弈式训练(用模型自身生成的轨迹迭代改进)和课程学习(从简单到复杂逐步增加训练难度)。第四,异构平台支持需要为每个新后端构建评估服务,工程成本较高。改进方向包括自动化评估服务生成和更通用的跨平台抽象层。
未来方向
论文作者提出了几个重要的未来方向:扩展到更多后端(如华为NPU等),自动化生产引擎的端到端pull-request过程,以及开发更灵活的agent工作流(包含更丰富的工具和自适应搜索策略)。基于Kernel-Smith的成果,还可以延伸出以下研究方向:第一,将进化搜索框架应用于其他高性能代码生成任务,如CUDA kernel、自定义操作符、甚至整个计算图的优化。第二,探索多目标进化(同时优化延迟、内存占用、能耗等),利用MAP-Elites天然支持多维特征空间的优势。第三,开发在线学习机制,让模型在实际部署中持续从新的优化轨迹中学习,逐步适应特定硬件和工作负载。第四,将训练方案应用于更小的模型(如7B-30B),降低部署成本,使进化式内核优化对更广泛的用户可用。第五,结合形式验证方法,在保证功能正确性的前提下进行更激进的性能优化,解决当前方法在Level 3任务上正确率不高的问题。第六,研究如何将内核级优化与系统级优化(如调度、内存管理)联合优化,缩小内核加速与端到端系统性能之间的差距。
复现评估
从复现角度来看,Kernel-Smith的复现存在一些挑战。积极的方面是,论文明确描述了评估配置(温度0.6、top-p 0.95、最大32K tokens、40轮进化、100次单元测试),且基于OpenEvolve开源框架构建。训练流程也有详细的描述:SFT使用200k样本、64k上下文长度,RL使用GRPO算法、每条数据8个候选、加速比作为奖励。然而,存在以下复现障碍:第一,数据访问需要联系作者通过邮件获取,且训练数据依赖特定的教师模型(DeepSeek-V3.2-Speciale、Gemini-3.0-pro),这些模型的可用性和成本可能成为问题。第二,235B参数模型的训练需要大量GPU资源,论文未报告具体的训练硬件和时间。第三,评估需要特定的硬件配置(NVIDIA H200 GPU、MetaX GPU),普通研究者可能无法获取。第四,KernelBench基准的具体配置和预处理细节(如hack检测的具体规则)可能需要额外澄清。总体而言,复现的难度较高,主要受限于算力要求和数据可访问性,但评估流程和框架设计是相对清晰和可复用的。
论文图表