AI summaryⓘ
The authors look at how AI models create CUDA code, which needs to follow special rules to work correctly but is hard for models to get right. They find that certain parts of the code are predicted with high confidence, while other critical parts have lower confidence and correspond to important execution details. Using this insight, the authors develop CuSeT, a training method that focuses on both the confident tokens and the tricky regions by masking tokens and adjusting how training samples are weighted. Their experiments show CuSeT helps AI models generate better CUDA code with fewer errors, and it works well across different models and sizes without needing expensive extra training. This approach improves correctness while keeping inference costs low.
CUDA kernelsLarge Language Models (LLMs)Supervised fine-tuning (SFT)Token confidenceExecution constraintsAdaptive token maskingSample reweightingFunctional correctnessPost-trainingKernel generation
Authors
Wentao Chen, Jiace Zhu, Xing Zhe Chai, Zeng Qu, Qiaoling Xiao, Liucheng Duan, An Zou
Abstract
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.