arXiv cs.AI· Annabelle Sujun Tang, Wei Tang, Songyao Jiang, Tianhui Li, Shengzhe Xu, Jiin Woo, Kaan Ozkara, Youngsuk Park, Zhen Jia, Hadi Esmaeilzadeh·· 2 天前AI 评分34
如何训练你的小龙：教小语言模型构建与优化 GPU 优化脚本
How to Train Your Little Dragon: Teaching a Small Language Model to Construct and Refine GPU Optimization Scripts
AI 导读
研究者提出“HOW TO TRAIN YOUR LITTLE DRAGON”训练方法，让 8B 参数小模型直接生成 GPU 优化脚本而非 kernel 代码，并基于 TVM MetaSchedule 在 54,116 个负载上构建 108,594 条训练样本。
正文 · AI 翻译
摘要：编译器变换会相互影响，使得 GPU 内核优化成为对组合的搜索，这些组合在保持计算不变的同时提升性能。我们提出 HOW TO TRAIN YOUR LITTLE DRAGON，一种训练方法，教一个小型语言模型既能构建完整的优化脚本，也能改进现有脚本。我们的模型不排序候选方案，也不生成内核代码，而是生成优化脚本。我们设计了一种紧凑、无损的表示，其确定性解码可重建精确的编译器轨迹，保留其变换参数和依赖关系。我们使用编译器框架 TVM MetaSchedule 在 54,116 个不同工作负载上构建训练数据，并在两个训练集中构建了 108,594 个示例：60,236 个教模型从未优化的内核生成优化脚本，48,358 个教模型在内核和某个已部分优化的脚本基础上生成更好的脚本。监督微调教会模型从这些源自搜索的示例中构建和改进优化脚本。随后我们加入强化学习，使用结合正确性与实测性能的塑形奖励，仅对数值正确的内核给予基于延迟的奖励。我们将此训练方法应用于一个 80 亿参数的模型。经过监督微调和强化学习后，该模型在构建任务上单次尝试正确率达到 91.27%，在改进任务上达到 92.01%，而强化学习在 102 个构建任务中的 88 个和 170 个改进提示中的 145 个上产生了比监督训练更快的内核。仅有 8B 参数的 LITTLE DRAGON 实现了比六个前沿 LLM 更高的单次尝试正确率和更好的内核性能：其中包括 GPT-6 Astra、Opus 5、Kimi-K3 和 GLM-5.3，它们各自拥有数千亿到数万亿参数，且针对特定任务的训练要么不可用，要么成本高昂。
| 主题： | 机器学习（cs.LG）；人工智能（cs.AI） |
| 引用为： | arXiv:2610.20980 [cs.LG] |
| （或 arXiv:2610.20980v1 [cs.LG] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2610.20980 arXiv 颁发的 DOI 通过 DataCite（待注册） |
提交历史
来自：Annabelle Sujun Tang [查看电子邮件]
[v1]
2026年10月14日星期三 21:49:24 UTC（209 KB）
来源：arXiv cs.AI · arxiv.org