KernelMosaic：用于 CUDA 内核生成的开放数据集和代理框架

arXiv:2611.14604v1 公告类型：新 摘要：训练大型语言模型以生成高效的 CUDA 内核需要跨越不同计算要求和迭代代码优化的监督。由于规范异构，构建这种监督具有挑战性，并且通过迭代复杂的测试和细化会出现有用的实现。我们引入 KernelMosaic，这是一个开放数据集，将不同的 CUDA 任务与经过验证的解决方案和轨迹配对。人工指导的审查和多样性意识的选择将任务表达为自然语言描述、参考程序、数学规范及其交叉点。为了构建解决方案和轨迹，我们的代理框架通过可重用的 CUDA 指导和执行反馈，通过实施、正确性修复和性能优化来指导强大的法学硕士。对 KernelMosaic 上的小型开源模型进行监督微调，在 KernelBench L1、L2 和 L3 上的正确率分别从 4%、0% 和 0% 提高到 88%、82% 和 16%，同时在每个级别都增加了比 PyTorch eager 更快的内核任务比例。进一步的训练实验表明，来自策划轨迹的监督大大优于直接响应，并促进持续的测试和修订。这些结果表明 KernelMosaic 使通用模型能够在 CUDA 内核生成方面取得实质性进展。