让大模型生成更可靠的GPU代码,通过关注关键代码段提升正确率。
From Tokens to Regions: CUDA-Sensitive Instruction Tuning for GPU Kernel Generation

- 基于令牌置信度模式,区分高可信和低可信代码区域。
- 在多个模型上显著提升代码正确率,优于传统微调和先进方法。
- 低成本高效方案,适合需要快速部署的GPU代码生成场景。
高性能CUDA内核对可扩展人工智能系统至关重要,但大语言模型在生成正确内核时仍面临严格且隐式的执行约束挑战。现有基于LLM的方法要么依赖昂贵的代理或强化学习流程,要么采用监督微调(SFT)目标,未能显式建模CUDA敏感性——即代码令牌或区域与执行约束的紧密耦合。本文从令牌置信度模式角度研究CUDA敏感性,发现其同时存在于令牌和区域层面:大多数敏感令牌被高置信预测,而少数低置信子集构成对应执行关键结构的区域。这表明有效内核生成需同时利用高置信敏感令牌并保留低置信敏感区域。基于此,我们提出一种轻量级后训练方法——CUDA敏感指令微调(CuSeT),在简单SFT框架内实现‘从令牌到区域’的优化,结合自适应令牌掩码与区域感知样本重加权。实验表明,CuSeT在多种模型家族和规模下持续提升功能性正确率,超越标准SFT及先进SFT变体,在推理成本远低于前沿模型的前提下达到竞争力表现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。