license: other base_model:

• meta-llama/Llama-3.1-8B-Instruct datasets:
• ScalingIntelligence/KernelBench
• GPUMODE/KernelBook library_name: transformers

KernelLLM

KernelBench-Triton Level 1において、当社の80億パラメータモデルは、GPT-4oやDeepSeek V3などのモデルをシングルショット性能で上回りました。複数回の推論では、KernelLLMの性能がDeepSeek R1を凌駕します。これは競合他社のモデルに比べて2桁少ないパラメータ数で達成された成果です。

KernelLLMによるカーネル開発の民主化

KernelLLMは、Llama 3.1 Instructをベースとした大規模言語モデルで、Tritonを使用したGPUカーネル作成タスクに特化してトレーニングされました。KernelLLMはPyTorchモジュールをTritonカーネルに変換し、KernelBench-Tritonで評価されました（詳細はこちら）。KernelLLMは、カーネル開発をよりアクセスしやすく効率的にすることで、GPUプログラミングの民主化を目指しています。

KernelLLMのビジョンは、効率的なTriton実装の自動生成を通じて、高性能GPUカーネルに対する需要の高まりに対応することです。ワークロードが大規模化し、多様なアクセラレータアーキテクチャが登場する中、カスタマイズされたカーネルソリューションの必要性が大幅に高まっています。既存の研究（1、2）の多くはテスト時最適化や推論時スケーリングに限定されており、KernelBench問題自体のトレースソリューションを調整するものもありますが、これらは分布外汎化に関する結果の情報量を制限しています。私たちの知る限り、KernelLLMは外部の（torch、triton）ペアでファインチューンされた最初のLLMであり、モデルを公開することでインテリジェントなカーネル作成システムへの進展を加速できることを願っています。

Tritonカーネル生成のためのKernelLLMワークフロー：当社のアプローチでは、KernelLLMを使用してPyTorchコード（緑色）をTritonカーネル候補に変換します。入力と出力のコンポーネントは太字で示されています。生成されたコードは、既知の形状を持つランダムな入力でカーネルを実行するユニットテストに対して検証されます。このワークフローにより、カーネル候補の生成数を増やすことで複数の生成（pass@k）を評価できます。最適なカーネル実装が選択され、返されます（緑色の出力）。

このモデルは、約25,000組のPyTorchモジュールとそれに対応するTritonカーネル実装のペア、および追加の合成生成サンプルでトレーニングされました。当社のアプローチは、TheStack [Kocetkov et al. 2022]からフィルタリングされたコードと、torch.compile()および追加のプロンプト技術を通じて生成された合成例を組み合わせています。フィルタリングおよびコンパイルされたデータセットは[KernelBook]](https://huggingface.co/datasets/GPUMODE/KernelBook)です。

作成されたデータセットに対してLlama3.1-8B-Instructを教師あり指示チューニングでファインチューンし、Tritonカーネル生成を対象としたKernelBenchの新バージョンであるKernelBench-Triton [Ouyang et al. 2025]において、正しいTritonカーネルと対応する呼び出しコードを生成する能力を測定しました。トレーニングと評価の両方で、フォーマット例を含むプロンプトテンプレートとともにtorchコードが使用されました。モデルは10エポック、バッチサイズ32でトレーニングされ、トレーニングデータの保持サブセットでのパープレキシティによって選択されたハイパーパラメータを持つ標準的なSFTレシピが使用されました。トレーニングは16GPU（192GPU時間）で約12時間の実時間を要し、最良のチェックポイントの検証結果を報告しています。

モデル性能

当社の80億パラメータモデルは、カーネル生成タスクにおいて、はるかに大規模なモデルと比較して競争力のあるまたは優れた性能を達成し、KernelBench Level 1に対する専門的なトレーニングアプローチの有効性を示しています。KernelLLMの推論はtemperature=1.0、top_p=0.97で実行されました。

結果として得られたモデルは、その小さなサイズにもかかわらず、最先端のLLMと競争力があります。当社は、LLMが効率的なGPUカーネルを記述する能力を評価するためのオープンソースベンチマークであるKernelBenchでモデルを評価しました。これは、Conv2DやSwish（レベル1）などの単一のtorch演算子から、完全なモデルアーキテクチャ（レベル3）まで、難易度レベルごとに整理された250の選択されたPyTorchモジュールを含みます。ベンチマークは、正しさ（参照PyTorch出力との比較）と性能（ベースライン実装に対する高速化の測定）の両方を測定します。当社は、LLMのTritonカーネル生成能力を評価する新しいKernelBench-Tritonバリアントを実装し、KernelLLMの能力を評価する理想的なベンチマークとしました。すべての測定はNvidia H100 GPUで行われました。

KernelLLMは、pass@k分析中に準対数線形のスケーリング動作を示します。

詳細については、Project Popcornをご覧ください。

インストール

KernelLLMを使用するには、必要な依存関係をインストールしてください：

pip install transformers accelerate torch triton

使用方法

KernelLLMは、PyTorchコードからTritonカーネルを生成するためのシンプルなインターフェースを提供します。付属のkernelllm.pyスクリプトは、モデルと対話するための複数の方法を提供します。

基本的な使用法

from kernelllm import KernelLLM

# モデルの初期化
model = KernelLLM()

# PyTorchモジュールの定義
pytorch_code = '''
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    バイナリ分類タスクのためのヒンジ損失を計算するモデル。
    """
    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 []
'''

# 最適化されたTritonコードの生成
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()

# リアルタイムで出力をストリーミング
model.stream_raw("プロンプトをここに入力", max_new_tokens=2048)

# Triton固有のプロンプトテンプレートなしで生のテキストを生成
raw_output = model.generate_raw("プロンプトをここに入力", temperature=1.0, max_new_tokens=2048)

現在の制限と今後の課題

有望な結果を示しているにもかかわらず、KernelLLMにはいくつかの制限があります：

• モデルは依然として誤ったAPI参照や構文エラーを生成する可能性があり、指示に従う能力が制限されています。
• 生成されたコードは構造的にコンパイラ生成の出力に似ており、モデルはしばしば意味のあるカーネルを実装できません。
• エラー分析では、変数名、テンソル形状、型処理、数値精度に関する指示に従うことに関連する一般的な問題が示されています。

モデル詳細

モデル開発者: Meta.

入力: モデルはテキストのみを入力として受け取ります。

出力: モデルはテキストのみを生成します。

モデルアーキテクチャ: KernelLLMは、最適化されたトランスフォーマーアーキテクチャを使用する自己回帰型言語モデルです。

モデル日付: KernelLLMは2025年3月にトレーニングされました。

ステータス: これはオフラインデータセットでトレーニングされた静的モデルです。

ライセンス: 詳細はLICENSE.pdfをご覧ください。

意図された使用法

意図された使用ケース: KernelLLMは、英語、関連するプログラミング言語、Python、およびTritonにおける商業および研究用途を意図しています。

範囲外の使用法: 適用可能な法律または規制（貿易コンプライアンス法を含む）に違反する方法での使用。英語以外の言語での使用。KernelLLMとそのバリアントの使用ポリシーおよびライセンス契約で禁止されているその他の方法での使用。

ハードウェアとソフトウェア

トレーニング要因: カスタムトレーニングライブラリを使用しました。

カーボンフットプリント: 合計で、KernelLLMのトレーニングにはH100-80GBタイプのハードウェアで250時間の計算が必要でした（ベースモデルのトレーニングは含まれません）。推定tCO2eq排出量の100%はMetaのサステナビリティプログラムによって相殺されました。

倫理的考慮事項と制限

KernelLLMとそのバリアントは、使用にリスクを伴う新しい技術です。これまでに行われたテストは英語で行われ、すべてのシナリオをカバーしておらず、カバーすることもできませんでした。これらの理由から、すべてのLLMと同様に、KernelLLMの潜在的な出力を事前に予測することはできず、場合によってはユーザーのプロンプトに対して不正確または不快な応答を生成する可能性があります。したがって、KernelLLMのアプリケーションを展開する前に、開発者はモデルの特定のアプリケーションに合わせた安全性テストとチューニングを実施する必要があります。

責任ある使用ガイドはhttps://ai.meta.com/llama/responsible-use-guideで入手できます。

引用

@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},
}