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

2026-06-15

AI Summary

This work addresses the challenge that large language models often violate strict execution constraints when generating CUDA kernels, leading to functional errors. The authors propose CuSeT, a low-cost instruction fine-tuning approach that is the first to uncover the dual manifestation of CUDA sensitivity at both token-level and region-level granularities. Building on this insight, CuSeT introduces a synergistic optimization strategy that integrates adaptive token-level masking with region-aware sample reweighting within a supervised fine-tuning framework, effectively balancing the exploitation of high-confidence tokens and the preservation of low-confidence yet critical regions. Experiments demonstrate that CuSeT consistently improves the functional correctness of generated kernels across multiple model families and scales, outperforming both standard and enhanced supervised fine-tuning baselines while achieving state-of-the-art performance at lower inference cost.

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.

Problem

• CUDA kernel generation
• Large Language Models
• execution constraints
• CUDA sensitivity
• functional correctness

Innovation

• CUDA-sensitive
• instruction tuning
• token-level masking
• region-aware reweighting
• GPU kernel generation