新智元报道
编辑：英智
Meta 推出 KernelLLM，这个基于 Llama 3.1 微调的 8B 模型，竟能将 PyTorch 代码自动转换为高效 Triton GPU 内核。实测数据显示，它的单次推理性能超越 GPT-4o 和 DeepSeek V3，多次生成时得分飙升。
在 AI 领域，参数规模曾被视为「性能天花板」。
Meta 最新发布的 KernelLLM，却用 8B 参数的「小身板」，在 GPU 内核生成任务中把 200B 的 GPT-4o 按在地上摩擦。
这是一个基于 Llama 3.1 Instruct 进行微调的 8B 参数模型，旨在将 PyTorch 模块自动转换为高效的 Triton GPU 内核。
KernelLLM 简直是 GPU 内核开发神器，用更少的参数实现了更强的性能，且简单易用。
它只有 8B 参数，但是在 KernelBench-Triton Level 1，单次推理性能超过了 GPT-4o 和 DeepSeek V3。
通过多次推理，KernelLLM 性能优于 DeepSeek R1。
这一切都来自一个参数规模比竞争对手小两个数量级的模型。
@Denis Kanonik 吐槽「这又是用测试集训练的吗？」
KernelLLM 让内核开发更易上手
KernelLLM 是一款基于 Llama 3.1 Instruct 的 8B 模型，专门针对用 Triton 编写 GPU 内核的任务进行了训练。
它能让 GPU 编程变得更简单，实现高性能 GPU 内核生成的自动化。
KernelLLM 通过自动化生成高效的 Triton 实现，满足对高性能 GPU 内核日益增长的需求。
随着工作负载的增大和加速器架构的多样化，对定制化内核解决方案的需求显著增加。
现在市面上很多相关工具，要么只能在测试的时候优化，要么就只盯着 KernelBench 的问题调优，很难应对更广泛的场景。
KernelLLM 是首个在外部（PyTorch，Triton）代码对数据上进行微调的 LLM。
Triton 内核生成工作流程
把 PyTorch 代码输进去，KernelLLM 就会生成 Triton 内核候选代码。
然后用单元测试来验证这些代码，用随机输入跑一跑，看看输出对不对。要是生成好几个候选代码，还能比比哪个最好，挑出最优的。
KernelLLM 的 Triton 内核生成流程：用 KernelLLM 把 PyTorch 代码翻译成 Triton 内核的候选代码。生成的代码会通过单元测试验证，测试用已知形状的随机输入数据运行内核。这个流程支持生成多个候选代码（通过 pass@k评估），增加候选数量来提高质量，最后选出最好的 Triton 内核实现作为输出（绿色部分）
为了训练这个模型，团队可是下了大功夫，用了 25000 多对（PyTorch，Triton）代码示例，还有合成的样本。
这些数据一部分来自 TheStack 的过滤代码，一部分是通过 torch.compile () 和提示技术生成的。
数据集 KernelBook，参考链接：https://huggingface.co/datasets/GPUMODE/KernelBook。
训练时用的是 Llama3.1-8B-Instruct 模型，在自定义数据集上做了监督微调（SFT），测试它在 KernelBench-Triton 上生成正确 Triton 内核及调用代码的能力。
KernelBench-Triton 是基于 KernelBench[Ouyang et al. 2025]开发的变体，专注 Triton 内核生成。
训练和评估时，PyTorch 代码会配置一个包含格式示例的提示模板作为指令。
模型训练了 10 个 epoch，批大小为 32，采用标准 SFT 方法，超参数根据验证集的困惑度（perplexity）来选择。
训练用了 16 个 GPU，共耗时 12 小时（192 GPU 小时），报告了最佳检查点的验证结果。
性能评估
尽管模型规模较小，但其性能可与最先进的 LLM 相媲美。
KernelBench-Triton 测试中，8B 参数的 KernelLLM，单次推理得分 20.2，比 671B 参数的 DeepSeek V3（16 分）和 200B 参数的 GPT-4o（15 分）都高。
要是多生成几个候选代码，得分还能蹭蹭往上涨，生成 10 个的时候能到 51.8 分，20 个的时候能到 57.1 分。
KernelLLM 推理用 temperature=1.0 和 top_p=0.97 运行。
在 KernelBench 上测试了模型，这是一个开源基准测试，用于评估 LLM 编写的高效 GPU 内核的能力。
它包含 250 个精心挑选的 PyTorch 模块，按负载调整，从简单的单操作（如 Conv2D 或 Swish，Level 1）到完整的模型架构（Level 3）。
它在不同难度的任务里表现都很稳，不管是简单的单个操作符，还是复杂的模型架构，都能应对。
测试会同时降低代码的正确性（通过与参考 PyTorch 输出对比）和性能（通过与基准实现的加速比）。
团队开发了一个新的 KernelBench-Triton 变体，专门评估 LLM 生成 Triton 内核的能力，非常适合测试 KernelLLM。
所有测试都在 NVIDIA H100 GPU 上完成。
KernelLLM 在 pass@k中表现出近似对数线性的扩展行为
KernelLLM 怎么用？
先装几个依赖包：
pip install transformers accelerate torch triton用的时候，先导入库，调用 generate_triton 函数，就能生成优化后的 Triton 代码啦。
KernelLLM 提供了一个简单的接口，用于从 PyTorch 代码生成 Triton 核。
from kernelllm import KernelLLM# Initialize the modelmodel = KernelLLM ()# Define your PyTorch modulepytorch_code = '''import torchimport torch.nn as nnclass Model (nn.Module): """ A model that computes Hinge Loss for binary classification tasks. """ def __init__(self): super (Model, self).__init__() def forward (self, predictions, targets): return torch.mean (torch.clamp (1 - predictions * targets, min=0)) batch_size = 128input_shape = (1,) def get_inputs (): return [torch.randn (batch_size, *input_shape)， torch.randint (0, 2, (batch_size, 1)) .float () * 2 - 1]def get_init_inputs (): return []'''# Generate optimized Triton codeoptimized_code = model.generate_triton (pytorch_code, max_new_tokens=512) print (optimized_code)要是不想写脚本，还能直接运行 python kernelllm.py，使用内置的 REPL 接口，打开交互式界面，实时看结果。
kernelllm.py 提供了多种与模型交互的方法。
python kernelllm.pyKernelLLM 提供了几种自定义生成过程的方法：
from kernelllm import KernelLLMmodel = KernelLLM ()# Stream output in real-timemodel.stream_raw ("Your prompt here"， max_new_tokens=2048)# Generate raw text without the Triton-specific prompt templateraw_output = model.generate_raw ("Your prompt here"， temperature=1.0, max_new_tokens=2048)有时它会犯点小错误，比如 API 引用不对、语法出错，有时候还不太能按指令生成理想的内核。
生成的代码结构有点像编译器自动吐出来的，有时在变量命名、张量形状、类型处理和数值精度这些细节上也容易出问题。
参考资料：
