HOW TO TRAIN YOUR LITTLE DRAGON：教导小型语言模型构建和改进 GPU 优化脚本

arXiv:2610.20980v1 公告类型：新 摘要：编译器变换之间会相互作用，因此 GPU 内核优化实质上是在搜索能够保持计算不变并提升性能的变换组合。我们提出 HOW TO TRAIN YOUR LITTLE DRAGON，这是一种训练方法，教会小型语言模型构建完整的优化脚本，并改进已有脚本。我们的模型不对候选方案进行排序，也不生成内核代码，而是生成优化脚本。我们设计了一种紧凑且无损的表示形式，其确定性解码能够重建精确的编译器轨迹，同时保留其中的变换参数及依赖关系。我们使用编译器框架 TVM MetaSchedule，在 54,116 个不同工作负载上构建训练数据，并构造了两个训练集中的 108,594 个样例：其中 60,236 个样例用于教模型根据未优化的内核生成优化脚本，48,358 个样例用于教模型根据内核和一个经过一定优化的脚本生成更好的脚本。监督微调利用这些从搜索中获得的样例，教模型构建和改进优化脚本。随后，我们加入强化学习，并采用一种塑形奖励，将正确性与实测性能相结合，仅对数值正确的内核给予基于延迟的额外奖励。我们将这一训练方法应用于一个拥有 80 亿参数的模型。在监督微调和强化学习之后，该模型在构建任务和改进任务中的每次尝试正确率分别达到 91.27% 和 92.01%；与此同时，在 102 个构建任务和 170 个改进提示中，强化学习生成的内核分别有 88 个和 145 个比监督训练生成的内核更快。尽管仅拥有 80 亿参数，我们的 LITTLE DRAGON 仍取得了高于六个前沿 LLM 的每次尝试正确率和更优的内核性能；这六个模型包括 GPT-6 Astra、Opus 5、Kimi-K3 和 GLM-5.3，它们的参数规模均从数千亿到数万亿不等，而针对特定任务的训练对这些模型而言要么无法进行，要么成本高昂。