license: other
base_model:

• meta-llama/Llama-3.1-8B-Instruct
  datasets:
• ScalingIntelligence/KernelBench
• GPUMODE/KernelBook
  library_name: transformers

KernelLLM

在KernelBench-Triton Level 1测试中，我们的80亿参数模型在单次推理性能上超越了GPT-4o和DeepSeek V3等模型。通过多次推理，KernelLLM的表现甚至优于DeepSeek R1。这一切都来自一个参数量比竞争对手少两个数量级的模型。

用KernelLLM让内核开发更触手可及

我们推出KernelLLM——一个基于Llama 3.1 Instruct的大语言模型，专为使用Triton编写GPU内核任务训练。该模型能将PyTorch模块转换为Triton内核，并在KernelBench-Triton基准测试中验证（参见此处）。KernelLLM旨在通过降低内核开发门槛来推动GPU编程民主化。

我们的愿景是通过自动化生成高效Triton实现方案，满足日益增长的高性能GPU内核需求。随着工作负载规模扩大和加速器架构多样化，定制化内核解决方案的需求急剧上升。虽然已有相关研究和实践，但多数局限于测试时优化或针对KernelBench问题本身的方案调优，这限制了模型在分布外场景的泛化能力。据我们所知，KernelLLM是首个基于外部（torch+triton）数据对进行微调的LLM，我们希望通过开源模型加速智能内核编写系统的研发进程。

Triton内核生成工作流：使用KernelLLM将PyTorch代码（绿色）转换为候选Triton内核，输入输出组件以粗体标注。通过随机形状输入的单元测试验证生成结果，增加候选生成数量（pass@k）实现多轮评估，最终选择最优内核实现（绿色输出）。

模型训练使用了约25,000组PyTorch模块与对应Triton内核实现的数据对，以及合成生成的样本。数据来源包括TheStack的过滤代码[Kocetkov等，2022]和通过torch.compile()与提示技术生成的合成样本，最终形成KernelBook数据集。

我们在自建数据集上对Llama3.1-8B-Instruct进行监督式指令微调，评估其在KernelBench-Triton（专为Triton内核生成设计的新基准）上生成正确内核及调用代码的能力。训练采用包含格式示例的提示模板，历时10个epoch（批量32），超参数通过验证集困惑度选择。在16块GPU上耗时约12小时（合计192 GPU小时），报告最佳检查点的验证结果。

模型性能

在Level 1基准测试中，我们的80亿参数模型与参数量更大的模型相比展现出竞争优势，证明了专项训练的有效性。测试参数：temperature=1.0，top_p=0.97。

尽管体积小巧，该模型性能仍可比肩前沿LLM。我们在KernelBench开源基准（包含250个分难度级别的PyTorch模块）上评估模型生成正确且高效Triton内核的能力，所有测试在NVIDIA H100 GPU完成。

KernelLLM在pass@k测试中呈现拟对数线性增长趋势。

更多信息请访问Project Popcorn。

安装

pip install transformers accelerate torch triton

使用方式

基础用法

from kernelllm import KernelLLM  
model = KernelLLM()  
optimized_code = model.generate_triton(pytorch_code, max_new_tokens=512)  

交互式REPL

python kernelllm.py  

高级选项

model.stream_raw("输入提示", max_new_tokens=2048)  
raw_output = model.generate_raw("输入提示", temperature=1.0)  

当前局限与未来工作

• 可能产生API引用错误和语法问题
• 生成代码结构类似编译器输出，常无法实现有效内核
• 常见错误涉及变量命名、张量形状、类型处理和数值精度

模型详情

开发者: Meta
架构: 自回归语言模型，优化transformer架构
训练时间: 2025年3月
状态: 基于离线数据的静态模型
碳排放: 总计250小时H100-80GB算力，碳足迹由Meta可持续发展计划抵消

使用政策

适用范围: 英语环境下的商业与研究用途，涉及Python/Triton编程
禁止用途: 违反法律法规的行为，非英语场景使用，违反使用政策的行为

伦理声明

建议开发者部署前进行应用场景特定的安全测试，详见责任使用指南。

引用

@software{kernelllm2025,
    title={KernelLLM},
    author={Fisches, Zacharias and Paliskara, Sahan and Guo, Simon and Zhang, Alex and Spisak, Joe and Cummins, Chris and Leather, Hugh and Isaacson, Joe and Markosyan, Aram and Saroufim, Mark},
    year={2025},
    month={5},
    note={通讯作者: Aram Markosyan, Mark Saroufim},
    url={https://huggingface.co/facebook/KernelLLM},
}