论文

从词元到区域:GPU 内核生成的 CUDA 敏感指令调优

From Tokens to Regions: CUDA-Sensitive Instruction Tuning for GPU Kernel Generation

模型训练监督微调与指令调优

摘要

高性能 CUDA 内核对于可扩展的 AI 系统至关重要,而 大语言模型 (LLM) 由于严格和隐式的执行约束,仍然难以生成正确的内核。现有的基于 LLM 的方法要么依赖于昂贵的代理或强化学习 (RL) 管道,要么采用无法显式建模 CUDA 敏感性的监督 微调 (SFT) 目标,即与执行约束紧密耦合的代码词元或区域。在这项工作中,我们从词元置信模式的角度研究了 CUDA 敏感性,表明 CUDA 敏感性出现在词元和区域级别,其中大多数 CUDA 敏感词元以高置信度进行预测,而较小的低置信度子集形成与执行关键结构相对应的区域。这些发现表明,有效的 CUDA 内核生成应该既利用高置信度 CUDA 敏感词元,又保留低置信度 CUDA 敏感区域。基于这些见解,我们提出了 \textbf{\underline{CU}DA-\underline{Se}nsitivestruct \underline{T}uning (CuSeT)},这是一种简单 SFT 框架内的低成本 后训练 方法。 CuSeT 通过将 \emph{自适应 词元级 掩蔽}与 \emph{区域感知样本重新加权}相结合,遵循“从标记到区域”的原则。实验表明,CuSeT 持续提高了多个模型系列和规模的功能正确性,优于标准 SFT 和高级 SFT 变体,同时以大幅降低的推理成本实现与前沿 CUDA 内核生成模型的竞争性能。