Meta「輕量級(jí)」KernelLLM顛覆GPU內(nèi)核生成，8B參數(shù)碾壓GPT-4o
在AI領(lǐng)域，參數(shù)規(guī)模曾被視為「性能天花板」。
Meta最新發(fā)布的KernelLLM，卻用8B參數(shù)的「小身板」，在GPU內(nèi)核生成任務(wù)中把200B的GPT-4o按在地上摩擦。
這是一個(gè)基于Llama 3.1 Instruct進(jìn)行微調(diào)的8B參數(shù)模型，旨在將PyTorch模塊自動(dòng)轉(zhuǎn)換為高效的Triton GPU內(nèi)核。
圖片
KernelLLM簡(jiǎn)直是GPU內(nèi)核開(kāi)發(fā)神器，用更少的參數(shù)實(shí)現(xiàn)了更強(qiáng)的性能，且簡(jiǎn)單易用。
它只有8B參數(shù)，但是在KernelBench-Triton Level 1，單次推理性能超過(guò)了GPT-4o和DeepSeek V3。
通過(guò)多次推理，KernelLLM性能優(yōu)于DeepSeek R1。
圖片
這一切都來(lái)自一個(gè)參數(shù)規(guī)模比競(jìng)爭(zhēng)對(duì)手小兩個(gè)數(shù)量級(jí)的模型。
@Denis Kanonik吐槽「這又是用測(cè)試集訓(xùn)練的嗎？」
圖片
KernelLLM讓內(nèi)核開(kāi)發(fā)更易上手
KernelLLM是一款基于Llama 3.1 Instruct的8B模型，專門針對(duì)用Triton編寫GPU內(nèi)核的任務(wù)進(jìn)行了訓(xùn)練。
它能讓GPU編程變得更簡(jiǎn)單，實(shí)現(xiàn)高性能GPU內(nèi)核生成的自動(dòng)化。
KernelLLM通過(guò)自動(dòng)化生成高效的Triton實(shí)現(xiàn)，滿足對(duì)高性能GPU內(nèi)核日益增長(zhǎng)的需求。
隨著工作負(fù)載的增大和加速器架構(gòu)的多樣化，對(duì)定制化內(nèi)核解決方案的需求顯著增加。
現(xiàn)在市面上很多相關(guān)工具，要么只能在測(cè)試的時(shí)候優(yōu)化，要么就只盯著KernelBench的問(wèn)題調(diào)優(yōu)，很難應(yīng)對(duì)更廣泛的場(chǎng)景。
KernelLLM是首個(gè)在外部（PyTorch，Triton）代碼對(duì)數(shù)據(jù)上進(jìn)行微調(diào)的LLM。
Triton內(nèi)核生成工作流程
把PyTorch代碼輸進(jìn)去，KernelLLM就會(huì)生成Triton內(nèi)核候選代碼。
然后用單元測(cè)試來(lái)驗(yàn)證這些代碼，用隨機(jī)輸入跑一跑，看看輸出對(duì)不對(duì)。要是生成好幾個(gè)候選代碼，還能比比哪個(gè)最好，挑出最優(yōu)的。
圖片
KernelLLM的Triton內(nèi)核生成流程：用KernelLLM把PyTorch代碼翻譯成Triton內(nèi)核的候選代碼。生成的代碼會(huì)通過(guò)單元測(cè)試驗(yàn)證，測(cè)試用已知形狀的隨機(jī)輸入數(shù)據(jù)運(yùn)行內(nèi)核。這個(gè)流程支持生成多個(gè)候選代碼（通過(guò) pass@k評(píng)估），增加候選數(shù)量來(lái)提高質(zhì)量，最后選出最好的Triton內(nèi)核實(shí)現(xiàn)作為輸出（綠色部分）
為了訓(xùn)練這個(gè)模型，團(tuán)隊(duì)可是下了大功夫，用了25000多對(duì)（PyTorch，Triton）代碼示例，還有合成的樣本。
這些數(shù)據(jù)一部分來(lái)自TheStack的過(guò)濾代碼，一部分是通過(guò)torch.compile () 和提示技術(shù)生成的。
數(shù)據(jù)集KernelBook，參考鏈接：https://huggingface.co/datasets/GPUMODE/KernelBook。
訓(xùn)練時(shí)用的是Llama3.1-8B-Instruct模型，在自定義數(shù)據(jù)集上做了監(jiān)督微調(diào)（SFT），測(cè)試它在KernelBench-Triton上生成正確Triton內(nèi)核及調(diào)用代碼的能力。
KernelBench-Triton是基于KernelBench[Ouyang et al. 2025]開(kāi)發(fā)的變體，專注Triton內(nèi)核生成。
訓(xùn)練和評(píng)估時(shí)，PyTorch代碼會(huì)配置一個(gè)包含格式示例的提示模板作為指令。
模型訓(xùn)練了10個(gè)epoch，批大小為32，采用標(biāo)準(zhǔn)SFT方法，超參數(shù)根據(jù)驗(yàn)證集的困惑度（perplexity）來(lái)選擇。
訓(xùn)練用了16個(gè)GPU，共耗時(shí)12小時(shí)（192 GPU小時(shí)），報(bào)告了最佳檢查點(diǎn)的驗(yàn)證結(jié)果。
性能評(píng)估
盡管模型規(guī)模較小，但其性能可與最先進(jìn)的LLM相媲美。
圖片
KernelBench-Triton測(cè)試中，8B參數(shù)的KernelLLM，單次推理得分20.2，比671B參數(shù)的DeepSeek V3（16分）和200B參數(shù)的GPT-4o（15分）都高。
圖片
要是多生成幾個(gè)候選代碼，得分還能蹭蹭往上漲，生成10個(gè)的時(shí)候能到51.8分，20個(gè)的時(shí)候能到57.1分。
KernelLLM推理用temperature=1.0和top_p=0.97運(yùn)行。
在KernelBench上測(cè)試了模型，這是一個(gè)開(kāi)源基準(zhǔn)測(cè)試，用于評(píng)估LLM編寫的高效GPU內(nèi)核的能力。
它包含250個(gè)精心挑選的PyTorch模塊，按負(fù)載調(diào)整，從簡(jiǎn)單的單操作（如Conv2D或Swish，Level 1）到完整的模型架構(gòu)（Level 3）。
它在不同難度的任務(wù)里表現(xiàn)都很穩(wěn)，不管是簡(jiǎn)單的單個(gè)操作符，還是復(fù)雜的模型架構(gòu)，都能應(yīng)對(duì)。
測(cè)試會(huì)同時(shí)降低代碼的正確性（通過(guò)與參考PyTorch輸出對(duì)比）和性能（通過(guò)與基準(zhǔn)實(shí)現(xiàn)的加速比）。
團(tuán)隊(duì)開(kāi)發(fā)了一個(gè)新的KernelBench-Triton變體，專門評(píng)估LLM生成Triton內(nèi)核的能力，非常適合測(cè)試KernelLLM。
所有測(cè)試都在NVIDIA H100 GPU上完成。
圖片
KernelLLM在pass@k中表現(xiàn)出近似對(duì)數(shù)線性的擴(kuò)展行為
KernelLLM怎么用？
先裝幾個(gè)依賴包：
pip install transformers accelerate torch triton
pip install transformers accelerate torch triton用的時(shí)候，先導(dǎo)入庫(kù)，調(diào)用generate_triton函數(shù)，就能生成優(yōu)化后的Triton代碼啦。
KernelLLM提供了一個(gè)簡(jiǎn)單的接口，用于從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)
from kernelllm import KernelLLM
# Initialize the model
model = KernelLLM()
# Define your PyTorch module
pytorch_code =
'''
import torch
import 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 = 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)要是不想寫腳本，還能直接運(yùn)行python kernelllm.py，使用內(nèi)置的REPL接口，打開(kāi)交互式界面，實(shí)時(shí)看結(jié)果。
kernelllm.py提供了多種與模型交互的方法。
python kernelllm.py
python kernelllm.pyKernelLLM提供了幾種自定義生成過(guò)程的方法：
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)
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)有時(shí)它會(huì)犯點(diǎn)小錯(cuò)誤，比如API引用不對(duì)、語(yǔ)法出錯(cuò)，有時(shí)候還不太能按指令生成理想的內(nèi)核。
生成的代碼結(jié)構(gòu)有點(diǎn)像編譯器自動(dòng)吐出來(lái)的，有時(shí)在變量命名、張量形狀、類型處理和數(shù)值精度這些細(xì)節(jié)上也容易出問(wèn)題。
