Meta「輕量級」KernelLLM顛覆GPU內核生成，8B參數碾壓GPT-4o

新智元報道

編輯：英智

【新智元導讀】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怎麼用？

先裝幾個依賴包：

用的時候，先導入庫，調用generate_triton函數，就能生成優化後的Triton代碼啦。

KernelLLM提供了一個簡單的接口，用於從PyTorch代碼生成Triton核。

要是不想寫腳本，還能直接運行python kernelllm.py，使用內置的REPL接口，打開交互式界面，實時看結果。

kernelllm.py提供了多種與模型交互的方法。

KernelLLM提供了幾種自定義生成過程的方法：

有時它會犯點小錯誤，比如API引用不對、語法出錯，有時候還不太能按指令生成理想的內核。

生成的代碼結構有點像編譯器自動吐出來的，有時在變量命名、張量形狀、類型處理和數值精度這些細節上也容易出問題。

參考資料：

https://x.com/reach_vb/status/1924478755898085552

https://huggingface.co/facebook/KernelLLM