🚀 KernelLLM

KernelLLM是一款基于Llama 3.1 Instruct的大语言模型，专门针对使用Triton编写GPU内核的任务进行训练。它能够将PyTorch模块转换为Triton内核，让GPU编程更加便捷高效，满足了对高性能GPU内核不断增长的需求。

🚀 快速开始

KernelLLM是基于Llama 3.1 Instruct的大语言模型，专门用于使用Triton编写GPU内核。它能将PyTorch模块转换为Triton内核，在KernelBench-Triton上进行了评估。

✨ 主要特性

• 高性能表现：在KernelBench-Triton Level 1上，80亿参数的KernelLLM模型在单次推理性能上超越了GPT-4o和DeepSeek V3等模型。多次推理时，其性能优于DeepSeek R1，且参数比竞争对手少两个数量级。
• 自动化内核生成：旨在通过自动化生成高效的Triton实现，使内核开发更加容易和高效，满足对高性能GPU内核不断增长的需求。
• 广泛的评估：在约25000个PyTorch模块及其等效Triton内核实现的配对示例以及额外的合成生成样本上进行训练，并在KernelBench-Triton上进行评估。
• 多种使用方式：提供简单的接口，可通过脚本、交互式REPL和高级选项等多种方式与模型交互。

📦 安装指南

要使用KernelLLM，需安装所需的依赖项：

pip install transformers accelerate torch triton

💻 使用示例

基础用法

from kernelllm import KernelLLM

# Initialize the model
model = KernelLLM()

# Define your PyTorch module
pytorch_code = '''
import torch
import torch.nn as nn

class 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 = 128
input_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 code
optimized_code = model.generate_triton(pytorch_code, max_new_tokens=512)
print(optimized_code)

交互式REPL

你还可以使用内置的REPL界面：

python kernelllm.py

这将启动一个交互式会话，你可以输入PyTorch代码并获得Triton优化的实现。

高级用法

KernelLLM提供了几种自定义生成过程的方法：

from kernelllm import KernelLLM

model = KernelLLM()

# Stream output in real-time
model.stream_raw("Your prompt here", max_new_tokens=2048)

# Generate raw text without the Triton-specific prompt template
raw_output = model.generate_raw("Your prompt here", temperature=1.0, max_new_tokens=2048)

📚 详细文档

模型性能

KernelLLM的80亿参数模型在KernelBench Level 1的内核生成任务上与更大的模型相比具有竞争力或更优的性能，展示了其在特定训练方法上的有效性。所有测量均在Nvidia H100 GPU上进行。

模型细节

预期用途

• 预期用例：KernelLLM旨在用于英语、相关编程语言、Python和Triton的商业和研究用途。
• 超出范围的用途：以任何违反适用法律法规（包括贸易合规法律）的方式使用；使用英语以外的语言；以可接受使用政策和KernelLLM及其变体许可协议禁止的任何其他方式使用。

硬件和软件

• 训练因素：使用自定义训练库。
• 碳足迹：总体而言，训练KernelLLM需要在H100 - 80GB类型的硬件上进行250小时的计算，不包括基础模型的训练。估计的tCO2eq排放量的100%由Meta的可持续发展计划抵消。

伦理考虑和局限性

KernelLLM及其变体是一项新技术，使用时存在风险。到目前为止的测试仅在英语环境中进行，无法涵盖所有场景。因此，与所有大语言模型一样，KernelLLM的潜在输出无法提前预测，在某些情况下可能会对用户提示产生不准确或令人反感的响应。因此，在部署KernelLLM的任何应用程序之前，开发者应针对其特定应用对模型进行安全测试和调整。

请参阅负责任使用指南。

引用

@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={Corresponding authors: Aram Markosyan, Mark Saroufim},
    url={https://huggingface.co/facebook/KernelLLM},
}

🔧 技术细节

KernelLLM基于Llama 3.1 Instruct进行微调，在约25000个PyTorch模块及其等效Triton内核实现的配对示例以及额外的合成生成样本上进行训练。训练过程结合了TheStack过滤后的代码和通过torch.compile()及额外提示技术生成的合成示例，过滤和编译后的数据集为KernelBook。

模型使用监督指令微调在创建的数据集上对Llama3.1 - 8B - Instruct进行微调，并在KernelBench-Triton上测量其生成正确Triton内核和相应调用代码的能力。训练和评估期间，torch代码使用包含格式示例的提示模板作为指令。模型训练了10个epoch，批量大小为32，并使用标准的SFT配方，超参数通过在训练数据的保留子集上的困惑度进行选择。训练在16个GPU上大约花费了12小时（192 GPU小时），报告的是最佳检查点的验证结果。

📄 许可证

详情见LICENSE.pdf。