From Tokens to Regions: CUDA-Sensitive Instruction Tuning for GPU Kernel Generation
作者: Wentao Chen, Jiace Zhu, Xing Zhe Chai, Zeng Qu, Qiaoling Xiao, Liucheng Duan, An Zou
分类: cs.LG, cs.AI
发布日期: 2026-06-15
💡 一句话要点
提出CuSeT以解决CUDA内核生成中的敏感性问题
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: CUDA内核生成 大型语言模型 敏感性分析 自适应掩蔽 样本重加权 深度学习优化 高性能计算
📋 核心要点
- 现有的LLM方法在生成CUDA内核时面临执行约束,导致生成的内核正确性不足。
- 论文提出CuSeT,通过结合自适应token级掩蔽与区域感知样本重加权,提升CUDA内核生成的准确性。
- 实验结果表明,CuSeT在多个模型上均显著提高了功能正确性,超越了标准SFT和其他变体。
📝 摘要(中文)
高性能的CUDA内核对于可扩展的AI系统至关重要,但大型语言模型(LLMs)在生成正确内核时仍面临严格的执行约束。现有基于LLM的方法要么依赖于昂贵的代理或强化学习(RL)流程,要么采用无法明确建模CUDA敏感性的监督微调(SFT)目标。本文从token置信度模式的角度研究CUDA敏感性,提出了CUDA-Sensitive Instruction Tuning(CuSeT),一种低成本的后训练方法,结合自适应token级掩蔽和区域感知样本重加权,显著提升了多个模型系列的功能正确性,超越了标准SFT和高级SFT变体,同时在推理成本上具有竞争力。
🔬 方法详解
问题定义:本文旨在解决现有LLM在生成CUDA内核时对执行约束的敏感性不足的问题。现有方法往往依赖于昂贵的强化学习或无法有效建模CUDA敏感性的监督微调,导致生成的内核在功能上不够准确。
核心思路:CuSeT的核心思路是从token到区域的转变,利用token置信度模式来识别CUDA敏感的token和区域,从而在生成过程中更好地利用高置信度的CUDA敏感token,同时保留低置信度的CUDA敏感区域。
技术框架:CuSeT的整体架构包括自适应token级掩蔽和区域感知样本重加权两个主要模块。自适应掩蔽用于选择高置信度的token,而区域重加权则帮助模型关注执行关键结构。
关键创新:CuSeT的创新在于结合了token级和区域级的敏感性分析,提出了一种新的训练策略,使得模型能够在生成CUDA内核时更有效地处理执行约束。这一方法与传统的SFT方法在设计理念上有本质区别。
关键设计:在参数设置上,CuSeT采用了动态的掩蔽策略和样本重加权机制,损失函数设计上强调了对CUDA敏感区域的关注,确保模型在训练过程中能够有效学习到执行约束的特征。
🖼️ 关键图片
📊 实验亮点
实验结果显示,CuSeT在多个模型系列上均显著提高了功能正确性,超越了标准SFT和其他高级变体,且在推理成本上具有竞争力,展示了其在CUDA内核生成中的有效性和优势。
🎯 应用场景
该研究的潜在应用领域包括高性能计算、深度学习框架优化和AI系统的内核生成。通过提升CUDA内核的生成准确性,CuSeT能够帮助开发者更高效地利用GPU资源,推动AI系统的可扩展性和性能提升,具有重要的实际价值和未来影响。
📄 摘要(原文)
High-performance CUDA kernels are essential for scalable AI systems, while Large Language Models (LLMs) still struggle to generate correct kernels due to strict and implicit execution constraints. Existing LLM-based approaches either rely on costly agentic or reinforcement-learning (RL) pipelines, or adopt supervised fine-tuning (SFT) objectives that fail to explicitly model CUDA sensitivity, namely code tokens or regions tightly coupled with execution constraints. In this work, we investigate CUDA sensitivity from the perspective of token confidence patterns, showing that CUDA sensitivity appears at both token and region levels, where most CUDA-sensitive tokens are predicted with high confidence, while a smaller low-confidence subset forms regions corresponding to execution-critical structures. These findings suggest that effective CUDA kernel generation should both leverage high-confidence CUDA-sensitive tokens and preserve low-confidence CUDA-sensitive regions. Building on these insights, we propose \textbf{\underline{CU}DA-\underline{Se}nsitive Instruction \underline{T}uning (CuSeT)}, a low-cost post-training method within a simple SFT framework. CuSeT follows the principle of ``from tokens to regions'' by combining \emph{adaptive token-level masking} with \emph{region-aware sample reweighting}. Experiments show that CuSeT consistently improves functional correctness across multiple model families and scales, outperforming standard SFT and advanced SFT variants, while achieving competitive performance against frontier CUDA kernel generation models with substantially lower inference cost.