カーネルの作成

このガイドでは、ステートレスなforwardの置き換えを超えたカーネルの書き方を説明します。拡張されたKernelConfig APIがサポートする2つの機能について解説します。

パラメータ変換: カーネルは、元のモデルとは異なるレイアウトの重みを期待します(例: パラメータの名前変更やマージ)。
モジュール融合: カーネルが、複数の隣接するモジュールを単一の融合済み実装に置き換えます。

基本的なカーネル（パラメータの変更を伴わないステートレスなforwardの代替実装）については、kernelsライブラリのドキュメントを参照してください。

2クラスのパターン

独自のパラメータを持つカーネルはすべて、2クラスのパターンに従います。

KernelName: forwardパスのみを含むクラスです。kernelsライブラリは、ステートフルなカーネルクラスを許可しないため、このクラスを使用してモデルをカーネル化します。
KernelNameLayout: パラメータを保持し、チェックポイントの読み込み前に元のモジュールをモンキーパッチする nn.Module です。実行時、kernelize はその forward を KernelName の forward に置き換えます。forward を定義する必要はありません。Transformers が KernelName.forward と同じシグネチャを持つものを自動的に注入します。

命名規則は厳格です。レイアウトクラスは{KernelName}Layoutと命名し、KernelNameと同じモジュール内で定義する必要があります。

パラメータ変換

カーネルが、元のモデルチェックポイントとは異なる名前や異なる形状の重みを期待する場合に、このパターンを使用します。

KernelNameLayout クラスは、置き換えるモジュールと同じ __init__ シグネチャを持ち、チェックポイントのキーを新しいパラメータ名にどのように再マッピングするかを Transformers に伝える conversion_mapping クラス属性を宣言します（詳細は動的な重みの読み込みを参照）。

# layers.py
try:
    from transformers import Concatenate, WeightConverter
    from transformers.conversion_mapping import WeightRenaming
except ImportError:
    Concatenate = WeightConverter = WeightRenaming = None

import torch
import torch.nn as nn


class CustomRMSNormLayout(nn.Module):
    conversion_mapping = [...]  # rules that remap checkpoint keys to the new parameter names

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps


class CustomRMSNorm(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.scale * hidden_states.to(input_dtype)

# __init__.py
from . import layers

kernels ライブラリは、layers モジュール（layers.py または layers/__init__.py）を通じてカーネルクラスを検出します。カーネルクラスとレイアウトクラスを layers.py に直接定義し、パッケージの __init__.py から layers モジュールを公開してください。Transformers は、CustomRMSNorm と同じモジュール内に CustomRMSNormLayout を見つけます。

layers.py 内の Transformers 固有のインポート（WeightConverter、Concatenate、WeightRenaming）をガードし、Transformers がインストールされていない環境でもカーネルをインポートできるようにします。

KernelConfig にリポジトリ名とクラス名を渡すことで、このカーネルを読み込みます。キーはモデルの元のモジュールクラス名です。値は、リポジトリ内の KernelName クラス（Layout ではなく）を指します。

値は"owner/repo:ClassName"の文字列、または(repo_string, metadata)のタプルです。タプルの形式では、versionまたはrevision（少なくとも一方）を渡し、必要に応じてtrust_remote_codeを指定します。メジャーなカーネルバージョンにはversionを、タグ、ブランチ、またはコミットにはrevisionを使用します。信頼できるkernels-community組織の外から読み込む場合は、trust_remote_code=Trueを設定します。単なる文字列の場合、デフォルトでversion=1およびtrust_remode_code=Falseが設定されます。

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig({"RMSNorm": ("owner/my-kernel:CustomRMSNorm", {"trust_remote_code": True, "version": 1})})
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

モデルが読み込まれると、Transformers は:

CustomRMSNormをリポジトリから読み込み、同じモジュール内でCustomRMSNormLayoutを探します。
モデル内のすべてのRMSNormをCustomRMSNormLayoutでモンキーパッチします。
conversion_mapping を使用してチェックポイントの重みを再マッピングし、新しいパラメータ名に読み込めるようにします。
kernelizeを呼び出し、CustomRMSNormLayout.forwardをCustomRMSNorm.forwardに置き換えます。

モジュールの融合

カーネルが複数の隣接するモジュールを単一の融合済み実装に置き換える場合に、このパターンを使用します。融合済みモジュールは複数の元のモジュールのパラメータを組み合わせるため、KernelNameLayout.__init__には、コンストラクタの引数ではなく、インスタンス化された子モジュールが渡されます。

# layers.py
try:
    from transformers import Concatenate, WeightConverter
    from transformers.conversion_mapping import WeightRenaming
except ImportError:
    Concatenate = WeightConverter = WeightRenaming = None

import torch
import torch.nn as nn


class RMSNormMLPLayout(nn.Module):
    conversion_mapping = [...]  # rules that remap checkpoint keys to the fused parameter names

    def __init__(self, norm, mlp):
        super().__init__()
        self.variance_epsilon = norm.variance_epsilon
        self.scale = nn.Parameter(torch.empty_like(norm.weight))
        self.gate_up_proj = nn.Linear(
            mlp.gate_proj.in_features,
            mlp.gate_proj.out_features + mlp.up_proj.out_features,
            bias=mlp.gate_proj.bias is not None,
            device=mlp.gate_proj.weight.device,
            dtype=mlp.gate_proj.weight.dtype,
        )
        self.down_proj = nn.Linear(
            mlp.down_proj.in_features,
            mlp.down_proj.out_features,
            bias=mlp.down_proj.bias is not None,
            device=mlp.down_proj.weight.device,
            dtype=mlp.down_proj.weight.dtype,
        )
        self.act_fn = mlp.act_fn


class RMSNormMLP(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        hidden_states = self.scale * hidden_states.to(input_dtype)
        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
        return self.down_proj(self.act_fn(gate) * up)

# __init__.py
from . import layers

モジュールを融合するには、プレーンな文字列の代わりに、KernelConfig のキーとして (class_name, path_pattern) のペアのタプルを渡します。すべてのパターンは同じ親モジュールを共有する必要があります（Transformers はその親の子を融合します）。* ワイルドカードは、任意の単一パスセグメントに一致します。

from transformers import AutoModelForCausalLM, KernelConfig

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

モデルが読み込まれると、Transformers は:

RMSNormMLPをリポジトリから読み込み、同じモジュール内でRMSNormMLPLayoutを検索します。
model.layers.* に一致するすべてのデコーダー層にマッチし、__init__ が RMSNormMLPLayout(post_attention_layernorm, mlp) を呼び出す融合された親クラスを構築します。
残りの子モジュール(mlp)をnn.Identity()に置き換え、親モジュールのインターフェースを維持します。
conversion_mappingを使用して、チェックポイントの重みを再マッピングします。
kernelizeを呼び出し、RMSNormMLPLayout.forwardをRMSNormMLP.forwardに置き換えます。

フュージョンタプル内のペアの順序は、KernelNameLayout.__init__ に渡される引数の順序を決定します。

Update on GitHub