Kernelllm

🚀 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。