KernelMosaic：面向 CUDA 内核生成的开源数据集与智能体框架
KernelMosaic: An Open Dataset and Agentic Framework for CUDA Kernel Generation
AI 导读
KernelMosaic 是一个开源数据集与智能体框架，将多样化 CUDA 任务与经验证的解决方案和轨迹配对，用于训练大语言模型生成 CUDA 内核。在小型开源模型上做监督微调后，KernelBench L1/L2/L3 的正确率分别从 4%、0%、0% 提升至 88%、82%、16%，且各层级快于 PyTorch eager 的内核比例均上升。训练实验表明，基于精选轨迹的监督显著优于直接响应。
正文 · 原文
Abstract:Training large language models to generate efficient CUDA kernels requires supervision spanning diverse computational requirements and iterative code optimization. Constructing this supervision is challenging because of heterogeneous specifications, and useful implementations emerge through iterative complicated testing and refinement. We introduce KernelMosaic, an open dataset pairing diverse CUDA tasks with verified solutions and trajectories. Human-guided review and diversity-aware selection curate tasks expressed as natural-language descriptions, reference programs, mathematical specifications, and their intersections. To build the solutions and trajectories, our agentic framework guides a strong LLM through implementation, correctness repair, and performance optimization with reusable CUDA guidance and execution feedback. Supervised fine-tuning of a small open-source model on KernelMosaic raises correctness from 4%, 0%, and 0% to 88%, 82%, and 16% on KernelBench L1, L2, and L3, respectively, while increasing the fraction of tasks with kernels faster than PyTorch eager at every level. Further training experiments show that supervision from the curated trajectories substantially outperforms direct responses and promotes sustained testing and revision. These results demonstrate that KernelMosaic enables a general-purpose model to achieve substantial gains in CUDA kernel generation.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2611.14604 [cs.LG] |
| (or arXiv:2611.14604v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2611.14604 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Ningyuan Xi [view email]
[v1]
Tue, 10 Nov 2026 17:45:06 UTC (679 KB)
来源：arXiv cs.AI · arxiv.org