# カーネルの作成

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

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

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

## 2クラスのパターン

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

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

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

## パラメータ変換

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

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

```python
# 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)
```

```python
# __init__.py
from . import layers
```

> [!NOTE]
> `kernels` ライブラリは、`layers` モジュール（`layers.py` または `layers/__init__.py`）を通じてカーネルクラスを検出します。カーネルクラスとレイアウトクラスを `layers.py` に直接定義し、パッケージの `__init__.py` から `layers` モジュールを公開してください。Transformers は、`CustomRMSNorm` と同じモジュール内に `CustomRMSNormLayout` を見つけます。
>
> `layers.py` 内の Transformers 固有のインポート（`WeightConverter`、`Concatenate`、`WeightRenaming`）をガードし、Transformers がインストールされていない環境でもカーネルをインポートできるようにします。

[KernelConfig](/docs/transformers/main/ja/main_classes/kernels#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`が設定されます。

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

## モジュールの融合

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

```python
# 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)
```

```python
# __init__.py
from . import layers
```

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

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

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