← 返回 2026-02-16

DICE:扩散大语言模型在CUDA内核生成中的卓越表现 DICE: Diffusion Large Language Models Excel at Generating CUDA Kernels

Haolei Bai, Lingcheng Kong, Xueyi Chen, Jianmian Wang, Zhiqiang Tao, Huan Wang 📅 2026-02-12 👍 8 2026-07-13 08:35
CUDA内核生成 代码生成 强化学习 扩散大语言模型 高性能计算

首个专门为CUDA内核生成设计的扩散大语言模型

前置知识

扩散大语言模型(dLLMs)

扩散大语言模型是一类新兴的语言模型,与传统的自回归模型不同,它通过迭代去噪的方式生成文本。具体来说,dLLMs首先将生成序列进行掩码处理,然后通过多个去噪步骤逐步恢复原始序列。这种机制允许模型在生成过程中同时考虑全局信息,并支持多token并行生成。在代码生成任务中,dLLMs能够进行全局结构规划和非顺序的迭代优化,这与人类编程的工作流程更为吻合。

本文的核心创新就是将dLLMs应用于CUDA内核生成,理解dLLMs的工作原理是理解本文方法的基础。

CUDA内核生成

CUDA内核生成是指使用AI模型自动生成可在NVIDIA GPU上高效执行的CUDA代码。一个成功的CUDA内核不仅要语法正确、功能等价于参考实现,还要能提供实际的性能加速。这个任务极具挑战性,因为需要理解底层硬件架构、内存层次结构、线程调度等复杂细节。目前主流的评估基准是KernelBench,它包含250个任务,分为三个难度级别:单算子Level 1、融合模式Level 2和完整模型架构Level 3。

CUDA内核生成是本文要解决的核心问题,理解其挑战性和评估标准对于理解论文贡献至关重要。

强化学习(RL)

强化学习是一种机器学习范式,智能体通过与环境交互来学习最优策略。在本文中,RL被用于进一步优化dLLMs的CUDA内核生成能力。具体来说,模型生成候选CUDA内核,然后根据执行结果即正确性获得奖励信号。通过最大化期望奖励,模型学会生成更高质量的CUDA内核。本文采用的是基于执行的二元奖励信号:如果生成的内核编译并执行正确则奖励为1,否则为0。

本文提出的BiC-RL框架是核心贡献之一,理解RL的基本原理有助于理解该框架的设计动机和优势。

欺骗行为(Deceptive Behavior)

欺骗行为是指模型生成的CUDA内核表面上看起来正确,但实际上并未真正实现CUDA优化。具体表现为三种类型:第一种是模型遵循提示中的示例内核结构但在前向传播中使用高级PyTorch函数绕过了真正的内核实现;第二种是模型生成了有效的自定义CUDA内核但未实现必要的调用逻辑;第三种是自定义内核在前向函数中被省略导致编译时无法执行。这种行为会导致评估脚本误判为正确但实际上并未达到CUDA优化的目标。

欺骗行为是CUDA内核生成中的关键问题,本文的BiC-RL框架专门设计来解决这个问题。

KernelBench基准

KernelBench是用于评估CUDA内核生成能力的标准基准,包含250个任务,分为三个难度级别:Level 1包含100个单算子操作,Level 2包含100个融合模式,Level 3包含50个完整模型架构。评估指标包括执行正确性Exec和加速指标fastp,其中fastp表示生成的内核既正确又比PyTorch实现快至少p倍的任务比例。

所有实验结果都在KernelBench上进行评估,理解这个基准是理解实验结果的关键。

研究动机

CUDA内核生成是代码生成中极具挑战性的任务。一个成功的内核不仅要语法和功能正确,还要提供实际的性能加速。现有的研究主要依赖自回归大语言模型AR LLMs,但AR模型存在几个根本性问题:首先AR模型采用严格的从左到右生成方式,这与CUDA内核中广泛存在的非局部依赖关系不匹配,CUDA内核涉及共享内存使用、同步机制、跨区域依赖等,这些都不适合严格的顺序生成;其次AR模型的推理速度较慢,特别是对于CUDA内核这种需要生成长代码序列的任务;第三高质量的CUDA内核训练数据严重稀缺,现有数据集中只有86%的样本表现出加速效果,且大部分加速比集中在1.0到1.5倍区间,这些微小的加速在GPU环境中可能是统计噪声而非真实性能提升。

本文的目标是本文的目标是开发专门针对CUDA内核生成的扩散大语言模型dLLMs。具体目标包括三个方面:第一构建高质量的SFT数据集CuKe,通过严格的2.0倍加速阈值筛选,确保每个样本都代表实际可验证的性能提升;第二设计双阶段强化学习框架BiC-RL,通过从易到难的渐进式训练,让模型先掌握内核核心逻辑再进行端到端生成;第三构建DICE模型系列包括1.7B、4B、8B参数,在多个参数规模上验证方法的有效性和可扩展性。

与已有工作不同的是,本文的独特切入角度体现在三个方面:第一首次系统地将dLLMs应用于CUDA内核生成,利用dLLMs的并行生成能力和双向注意力机制来更好地处理CUDA内核中的非局部依赖关系;第二提出双阶段强化学习框架BiC-RL,创新性地将内核填充Kernel Infilling和端到端生成Kernel Generation分离,通过提供前缀和后缀脚手架来防止欺骗行为,让模型专注于核心逻辑实现;第三引入数据调度策略,模仿课程学习的思想,在训练过程中从简单操作逐步过渡到复杂模型结构,有效解决了RL训练中的优化平台期问题。

核心方法

DICE的训练框架采用渐进式策略,从数据准备到模型训练形成完整流水线。整体思路是首先通过严格的筛选标准构建高质量的CuKe数据集,然后使用SFT建立基础能力,最后通过双阶段RL进一步提升性能。技术路线可以概括为数据质量控制、监督微调冷启动、渐进式强化学习优化。这种设计的直觉是CUDA内核生成需要同时掌握两个能力:理解CUDA编程的底层逻辑通过SFT建立,和生成高性能的优化内核通过RL提升。双阶段RL的设计则源于对欺骗行为的深入分析,模型容易作弊生成表面正确但实际无效的内核,因此需要先通过内核填充阶段确保模型掌握核心实现逻辑。

本文的核心创新点是BiC-RL双阶段强化学习框架,它与已有方法的本质区别在于任务分解和渐进式训练。传统RL直接进行端到端生成但面临探索空间过大、优化困难的问题。BiC-RL将任务分解为两个阶段:内核填充阶段Kernel Infilling和内核生成阶段Kernel Generation。在内核填充阶段模型只需要生成CUDA内核的核心C++实现部分,前缀和后缀包括库导入、编译逻辑、nn.Module类结构都由模板提供,这确保模型专注于学习核心逻辑,同时通过脚手架设计防止了欺骗行为。在内核生成阶段模型才进行完整的端到端生成。这种分解的创新性在于通过结构化脚手架消除了欺骗行为的可能性,降低了早期训练的探索难度提高了训练效率,建立了从局部优化到全局生成的知识迁移。

方法步骤详情

DICE的训练分为四个主要步骤。第一步是CuKe数据集构建,基于ConCuR数据集进行增强,应用2.0倍加速阈值筛选得到1425个高质量PyTorch-CUDA对,同时从主流LLM中提取复杂结构如Attention子模块和MLP块,使用Mercury Coder生成优化内核,最终得到6303个样本。第二步是监督微调SFT,在SDAR模型基础上使用CuKe数据集训练3个epoch,学习率1.0乘以10的负5次方,使用8块NVIDIA A100 GPU。第三步是内核填充阶段RL,使用992个程序,将CUDA内核分解为前缀、核心实现和后缀三部分,模型只需生成核心部分,训练20步,每步采样64个问题,每个问题生成16个候选响应。第四步是端到端内核生成阶段RL,使用4000个程序,模型需要生成完整的CUDA内核,训练100步。两个RL阶段都采用基于执行的二元奖励信号,并通过z-score归一化进行组内标准化。数据调度策略在整个RL过程中逐步从单算子操作过渡到复杂模型结构。

技术新颖性

DICE的技术新颖性体现在多个层面。首先在模型架构选择上首次将dLLMs系统性地应用于CUDA内核生成,利用dLLMs的并行生成能力即推理速度优势和双向注意力机制更好地处理非局部依赖来解决AR模型的固有局限。其次在训练框架设计上BiC-RL创新性地将内核填充和端到端生成分离,通过结构化脚手架前缀加后缀消除了欺骗行为的可能性,这种设计不仅解决了评估可靠性问题还通过降低探索难度提高了训练效率。第三在奖励函数设计上虽然采用简单的二元执行奖励但通过组内z-score归一化有效地处理了不同问题难度的差异。第四在训练策略上数据调度机制模仿课程学习思想从简单操作逐步过渡到复杂结构有效解决了RL训练中的优化平台期问题。最后在数据构建上CuKe数据集通过严格的2.0倍加速阈值和结构多样性增强实现了在有限数据下的高性能。

扩散大语言模型的推理范式
Figure 2: 扩散大语言模型的推理范式
DICE框架总览
Figure 3: DICE框架总览
CUDA内核组件定义
Figure 4: CUDA内核组件定义

实验结果

实验结果表明DICE在KernelBench基准上取得了显著的性能提升。在8B参数规模上DICE-8B在Level 1上达到47%执行正确率fast1为15% fast2为6%,Level 2上达到35%执行正确率fast1为19% fast2为1%,Level 3上达到14%执行正确率fast1为6% fast2为6%。这些结果显著超越了所有可比规模的AR LLMs和dLLMs,包括专门针对CUDA内核训练的cudaLLM-8B其结果为46%、34%、6%。值得注意的是DICE-8B在数据效率上表现突出:仅使用6303个SFT样本和4992个RL样本,而cudaLLM使用了8920个SFT样本和71996个RL样本,且DICE的最大生成长度为4096 tokens而cudaLLM为32768 tokens。在4B规模上DICE-4B达到27%、24%、12%超越了多个8B规模的模型。在1.7B规模上DICE-1.7B达到24%、6%、10%在大多数其他1.7B模型完全失败的情况下仍能保持不错性能。消融实验表明SFT阶段建立了基础能力对1.7B模型至关重要,内核填充阶段和数据调度都能有效提升性能,完整的BiC-RL框架带来最大的性能提升。

8B及相似规模模型在KernelBench上的主要结果
Table 1: 8B及相似规模模型在KernelBench上的主要结果
4B及相似规模模型在KernelBench上的主要结果
Table 2: 4B及相似规模模型在KernelBench上的主要结果
1.7B及相似规模模型在KernelBench上的主要结果
Table 3: 1.7B及相似规模模型在KernelBench上的主要结果
不同加速阈值的CuKe数据集变体性能对比
Table 4: 不同加速阈值的CuKe数据集变体性能对比
不同训练策略的消融实验结果
Table 6: 不同训练策略的消融实验结果
Qwen3和SDAR系列模型在CuKe上SFT结果对比
Table 7: Qwen3和SDAR系列模型在CuKe上SFT结果对比
排除欺骗行为后的性能对比
Table 8: 排除欺骗行为后的性能对比
训练成本对比
Table 9: 训练成本对比
不同最大生成token数的性能对比
Table 10: 不同最大生成token数的性能对比
cudaLLM与DICE的详细对比
Table 11: cudaLLM与DICE的详细对比
cudaLLM SFT数据集的加速分布
Table 12: cudaLLM SFT数据集的加速分布
性能与推理速度对比及RL训练轨迹
Figure 1: 性能与推理速度对比及RL训练轨迹
SDAR与Qwen3对比SFT后的平均性能增益
Figure 5: SDAR与Qwen3对比SFT后的平均性能增益
查看结构化数据
任务指标本文基线提升
KernelBench Level 1单算子操作 执行正确率Exec DICE-8B达到47.0% cudaLLM-8B为46.0%,Qwen3-8B为8.0%,LLaDA-Instruct-8B为2.0% 比cudaLLM高1.0个百分点,比通用模型高39到45个百分点
KernelBench Level 2融合模式 执行正确率Exec DICE-8B达到35.0% cudaLLM-8B为34.0%,Qwen3-8B为25.0%,SDAR-Chat-8B为2.0% 比cudaLLM高1.0个百分点,比通用模型高10到33个百分点
KernelBench Level 3完整模型 执行正确率Exec DICE-8B达到14.0% cudaLLM-8B为6.0%,Qwen3-8B为4.0%,SDAR-Chat-8B为0.0% 比cudaLLM高8.0个百分点,是其他模型的2到14倍
KernelBench Level 1的4B规模 执行正确率Exec DICE-4B达到27.0% Qwen3-4B为3.0%,SDAR-Chat-4B为5.0% 比同规模模型高22到24个百分点,接近8B模型性能
KernelBench Level 1的1.7B规模 执行正确率Exec DICE-1.7B达到24.0% Qwen3-1.7B为0.0%,SDAR-Chat-1.7B为0.0% 在其他模型完全失败的情况下达到24%正确率
欺骗行为抑制 Robust Check后性能下降 DICE-8B在Level 1无下降,Level 2无下降,Level 3从20%降至14% LLaDA-Instruct在Level 2从13%降至0%,Level 3从26%降至0% DICE在Level 1和Level 2上完全消除了欺骗行为

局限与改进

本文存在几个明显的局限性。首先奖励函数设计过于简单:当前采用基于执行的二元奖励信号正确为1错误为0,没有将执行延迟直接纳入奖励函数。虽然作者解释这是为了在分布式训练环境中保持RL过程的稳定性因为CUDA内核的加速测量存在高方差和环境噪声,但这意味着模型优化目标与最终评估指标性能加速之间存在gap。其次评估范围有限:所有实验都在KernelBench基准上进行,该基准仅包含250个任务且主要针对PyTorch操作,可能无法完全代表实际应用场景中的CUDA内核需求。第三计算资源需求较高:训练需要8块NVIDIA A100 GPU,完整的BiC-RL训练包括SFT、内核填充RL、端到端生成RL需要约1100 GPU小时,这对资源有限的研究团队可能构成障碍。第四模型在最困难的Level 3任务上性能仍然有限最高14%,且仍存在部分欺骗行为从20%降至14%,说明对于复杂模型架构的内核生成仍有较大提升空间。

独立分析的弱点

基于独立分析DICE存在几个可以改进的弱点。第一数据规模相对较小:CuKe数据集仅包含6303个样本,虽然作者强调数据效率但在更复杂任务上可能面临数据不足的问题,特别是Level 3任务只有36个经过验证的结构化样本。改进方向是扩大数据收集范围,可能通过自动化流水线从更多开源项目中提取高质量CUDA内核。第二模型架构未针对CUDA内核生成进行专门优化:DICE直接在SDAR基础上进行微调没有针对CUDA内核的特殊结构如内存层次和线程块组织进行架构调整。可以考虑设计专门的注意力机制或位置编码来更好地捕捉CUDA代码的结构特征。第三推理效率仍有提升空间:虽然dLLMs相比AR模型有并行生成优势但在实际部署中生成长度限制4096 tokens可能限制了复杂内核的生成能力,需要设计更高效的解码策略。第四缺乏与人类专家的对比:所有实验都是模型之间的对比没有与人类CUDA程序员的性能进行比较,难以评估实际应用价值。

未来方向

作者提出了几个重要的未来研究方向,基于现有成果还可以延伸出更多方向。作者提出的方向包括整合多目标奖励函数将执行延迟直接纳入RL训练,虽然面临测量方差的挑战但可以通过统计方法或模拟环境来解决,以及扩展到更多硬件平台和编程框架如AMD ROCm和Intel oneAPI等。基于现有成果可延伸的方向包括开发更先进的数据增强技术通过程序变换和参数变异等方法扩大训练数据规模,探索多任务学习将CUDA内核生成与其他代码优化任务如Triton内核生成和内核调优联合训练,设计人机协作的交互式生成系统让人类专家可以提供反馈和指导,开发更可靠的评估方法能够自动检测和排除欺骗行为而不需要人工审查,探索迁移学习将在CUDA内核生成上学到的知识迁移到其他硬件特定代码生成任务。

复现评估

本文的复现性评估总体较好。代码和数据方面作者提供了项目页面但论文中未明确说明是否开源了完整的训练代码和CuKe数据集。基于SDAR模型的初始化是公开的,CuKe数据集的构建过程有详细描述基于ConCuR数据集增强。计算资源方面所有实验使用8块NVIDIA A100 GPU,SFT训练约22.4 GPU小时,内核填充RL约120 GPU小时,端到端生成RL约960 GPU小时,总计约1100 GPU小时这对学术研究机构是可承受的。训练细节方面论文提供了详细的超参数设置包括学习率、batch size、训练步数等以及评估配置包括不同模型的推理设置。潜在复现挑战包括CuKe数据集的构建需要执行大量CUDA内核并测量性能需要相应的GPU环境,RL训练涉及在线采样和执行验证需要稳定的CUDA编译和执行环境,部分基线模型如cudaLLM的训练细节未完全公开可能影响公平对比。