カーネルの読み込み

カーネルは、標準的な PyTorch 操作のドロップイン代替として機能します。モデルコードを壊すことなく、forward メソッドを最適化されたカーネル実装に置き換えます。

このガイドでは、推論を高速化するためにカーネルを読み込む方法を示します。

サポートされているバージョンのkernelsパッケージとともにTransformersをインストールします。

pip install -U "transformers[kernels]"

use_kernels=True を from_pretrained() に設定すると、デバイスで利用可能な最も高性能なカーネルを Hub から読み込むことができます。これにより、サポートされている PyTorch の演算がカーネル実装に置き換えられます。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

デフォルトのマッピングは、モデル、デバイス、および実行モードによって異なります。以下の表には、表示されているデバイスに関する一般的な Hub リポジトリが記載されています。特定の組み合わせにカーネルが登録されていない場合、その操作は標準の PyTorch にフォールバックします。

操作	NVIDIA (CUDA)	AMD (ROCm)	Intel (XPU)
RMSNorm	`kernels-community/liger-kernels`	`kernels-community/liger-kernels`	`kernels-community/rmsnorm`
MoE MLP	`kernels-community/megablocks`	`kernels-community/megablocks`	`kernels-community/megablocks`
MLP (SwiGLU、GeGLU)	`kernels-community/liger-kernels`	—	—
Linear	`kernels-community/liger-kernels`	—	—
活性化関数(GELUのバリアント、SiLU)	`kernels-community/activation`	—	—
Rotary embeddings	`kernels-community/rotary`	`kernels-community/aiter-rope`	`kernels-community/rotary`
因果言語モデルの損失	`kernels-community/liger-kernels`	—	—
Deformable attention	`kernels-community/deformable-detr`	—	—

この表は網羅的ではありません。モデルは追加のレイヤーや関数を登録でき、一部のマッピングは特定の実行モードでのみ利用可能です。

AMD GPU は PyTorch においてデバイスタイプを cuda として報告します。Transformers は実行時に ROCm を検出し、AITER のビルド（kernels-community/aiter-rope など）を含む、上記の AMD カーネルにサポートされている操作をルーティングします。デバイスタイプを自分で設定する必要はありません。

kernels-community 組織で利用可能なカーネルを閲覧してください。

Attentionカーネル

attn_implementation 引数を使用して、Hub からアテンションカーネルを読み込みます。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    device_map="auto"
)

アテンションカーネルについては、kernels-community リポジトリ（信頼できるリポジトリであり、将来的により多くの信頼できるリポジトリを追加する可能性があります）の一部でないものは、追加の allow_all_kernels=True キーワード引数の使用が必要になります（HF 以外のモデルにおける trust_remote_code=True キーワード引数と同様です）。これは、カーネルの読み込みがホストマシン上で任意のコードの実行につながる可能性があるため、すべてのリポジトリを検証できないため、明示的に許可する必要があるためです。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="random-repo/random-attention",
    allow_all_kernels=True,
    device_map="auto"
)

attention などの特定のカーネルは、複数のフォーマットを受け入れます。

@v2.1.0 は、特定のタグまたはブランチに固定します。
@>=2.0,<3.0 はセマンティックバージョニングの制約を設定します。

from transformers import AutoModelForCausalLM

# pin to a specific version
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2@v2.1.0",
    device_map="auto"
)
# use semantic versioning constraints
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2@>=2.0,<3.0",
    device_map="auto"
)

FlashAttention のフォールバック

attn_implementation="flash_attention_2"、"flash_attention_3"、または "flash_attention_4" を指定した場合、コンパイル済みの flash-attn パッケージがインストールされていないか、デバイスが CUDA でない場合は、対応する Hub カーネルにフォールバックします。

FlashAttention-4 のサポートはベータ版です。API および動作は変更される可能性があります。

from transformers import AutoModelForCausalLM

# uses the compiled flash-attn package if present, otherwise the kernels-community/flash-attn2 Hub kernel
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="flash_attention_2",
    device_map="auto",
)

モード認識

Kernelsは、PyTorchのmodel.training状態に基づいて、学習モードと推論モードに自動的に適応します。

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

# Switch to inference mode - uses inference-optimized kernels
model.eval()
with torch.no_grad():
    output = model.generate(input_ids, max_new_tokens=50)

# Switch to training mode - uses training-optimized kernels with gradient support
model.train()
loss = model(input_ids, labels=labels).loss
loss.backward()

mode 引数を使用して、kernelize() 関数でトレーニングモードと推論モードを明示的に有効にします。トレーニングモードでは、追加の torch.compile モードもサポートされています。

from kernels import Mode
from transformers import kernelize

# inference optimized kernels
kernelize(model, mode=Mode.INFERENCE)

# training optimized kernels
kernelize(model, mode=Mode.TRAINING)

# training and torch-compile friendly kernels
kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE)

KernelConfig

KernelConfig は、モデル内で使用されるカーネルをカスタマイズします。

kernel_mapping のキーは、モデルによって登録された名前です。これは "RMSNorm" のようなレイヤーや、"rotary_pos_emb" のような登録済み関数を指すことができます。: 区切り文字は、リポジトリ内の特定のカーネルエントリを指定し、そのレイヤーまたは関数にマッピングします。

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": "kernels-community/liger-kernels:LigerRMSNorm",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto"
)

デバイス種別ごとに異なるカーネル実装を指定します。

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": {
            "cuda": "kernels-community/liger-kernels:LigerRMSNorm",
            "rocm": "kernels-community/liger-kernels:LigerRMSNorm",
            "xpu": "kernels-community/rmsnorm:RMSNorm"
        }
    }
)

カーネルメタデータ

metadata辞書を追加して、読み込まれるカーネルビルドを制御します。

オプション	説明	デフォルト
`version`	カーネルリポジトリのメジャーバージョン。`2` は、リポジトリの `v2` ブランチ上の最新ビルドを読み込みます。	`1`
`revision`	バージョンではなく、読み込むための正確なタグ、ブランチ、またはコミット。	---
`trust_remote_code`	信頼できる `kernels-community` 組織外のリポジトリを許可します。カーネルを読み込むと、そのリポジトリのコードがあなたのマシン上で実行されます。	`False`

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": ("kernels-community/liger-kernels:LigerRMSNorm", {"version": 3}),
    }
)

カーネルリポジトリの古いブランチには、お使いの PyTorch および CUDA のバージョンに対応するビルドが含まれていない場合があります。カーネルの読み込みに失敗した場合は、ハードウェアがサポートされていないと判断する前に、新しい version を試してください。

継承されたマッピング

KernelConfig は、Hub カーネルに対する Transformers のデフォルトカーネルマッピングを継承し、kernel_mapping のエントリは、対応するレイヤーおよび関数に対してデフォルトを上書きします。inherit_mapping=False を設定すると、kernel_mapping のエントリのみが使用されます。指定しなかった項目はすべて、標準の PyTorch にフォールバックします。これは、他のデフォルトを適用せずに特定のカーネルをベンチマークしたり、カスタム実装をテストしたりする場合に便利です。

以下の設定では、rotary_pos_emb 関数を RoPE カーネルにマッピングし、その他のすべての操作は PyTorch のままにします。関数名はモデルによって登録されている必要があります。

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "rotary_pos_emb": (
            "kernels-community/rotary:apply_rotary_transformers",
            {"version": 2},
        ),
    },
    inherit_mapping=False,
)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

モジュールの融合

(class_name, path_pattern) のペアのタプルを KernelConfig のキーとして渡すことで、隣接するモジュールを単一のカーネルに融合できます。すべてのパターンは同じ親モジュールを共有する必要があります。* は任意の単一パスセグメントに一致します。

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    {
        (
            ("RMSNorm", "model.layers.*.post_attention_layernorm"),
            ("MLP",     "model.layers.*.mlp"),
        ): "owner/fused-rmsnorm-mlp:RMSNormMLP",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Fusion には、KernelName クラスのほかに、KernelNameLayout クラスを伴うクラスをカーネルのリポジトリで提供する必要があります。実装方法については、カーネルの作成ガイドを参照してください。

ローカルカーネル

use_local_kernel=True を KernelConfig に指定して、ローカルファイルパスからカーネルを読み込みます。これにより、Hub リポジトリではなくローカルファイルシステム上のパスから読み込みが行われます。

ローカルカーネルでは、Hub形式のorg/repo:layer_nameの代わりに/abs/path:layer_nameを使用します。

from transformers import KernelConfig, AutoModelForCausalLM

kernel_mapping = {
    "RMSNorm": "/path/to/liger-kernels:LigerRMSNorm",
}
kernel_config = KernelConfig(kernel_mapping, use_local_kernel=True)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype="auto",
    device_map="auto",
    use_kernels=True,
    kernel_config=kernel_config
)

カーネルをグローバルで無効にする

コードを変更せずに、すべての場所で Hub カーネルを無効にするには、USE_HUB_KERNELS 環境変数を設定してください。

export USE_HUB_KERNELS=0  # or OFF or NO

トラブルシューティング

カーネル統合は、ハードウェア、ドライバー、パッケージのバージョンが互いに正常に動作していることに依存します。以下のセクションでは、一般的な失敗例について説明します。

インストールの問題

インポートエラーは、kernels パッケージが欠落しているか、そのバージョンが Transformers がサポートする範囲外であることを示しています。互換性のあるバージョンを取得するには、extra を通じて再インストールしてください。

pip install -U "transformers[kernels]"

カーネルの読み込み失敗

特定のカーネルの読み込みに失敗した場合は、以下の手順をお試しください。

カーネルの要件に対するハードウェアの互換性を確認してください。
CUDA/ROCm/Metal ドライバーが最新であることを確認してください。
既知の問題については、カーネルリポジトリのドキュメントを参照してください。

デバイス互換性

すべてのカーネルがすべてのデバイスに対応しているわけではありません。お使いのハードウェアでカーネルが利用できない場合、ライブラリは標準的な PyTorch 操作にフォールバックします。デバイス固有のサポート状況については、カーネルのリポジトリのドキュメントをご確認ください。

リソース

Kernels リポジトリ
Hugging Face Kernel Hubで5分以内にモデルを強化するブログ記事
kernels-community 組織でカーネルを発見

Update on GitHub