Menulis kernel

Panduan ini menjelaskan cara menulis kernel yang melampaui penggantian forward tanpa status. Ini mencakup dua kemampuan yang didukung oleh API KernelConfig yang diperluas:

1. Transformasi parameter: kernel mengharapkan bobot dalam tata letak yang berbeda dari model asli (misalnya, parameter yang namanya diubah atau digabungkan).
2. Fusi modul: kernel menggantikan beberapa modul bersebelahan dengan satu implementasi yang telah difusi.

Untuk kernel dasar (pengganti forward stateless tanpa perubahan parameter), lihat dokumentasi pustaka kernels.

Pola dua kelas

Setiap kernel yang membawa parameternya sendiri mengikuti pola dua kelas.

• KernelName: hanya berisi jalur forward. Perpustakaan kernels menggunakan kelas ini untuk meng-kernelisasi model karena tidak mengizinkan kelas kernel yang berstatus (stateful).
• KernelNameLayout: sebuah nn.Module yang menyimpan parameter dan melakukan monkey-patch pada modul asli sebelum checkpoint dimuat. Saat runtime, kernelize mengganti forward miliknya dengan forward dari KernelName. Anda tidak perlu mendefinisikan forward. Transformers menyuntikkannya secara otomatis dengan tanda tangan yang sama seperti KernelName.forward.

Konvensi penamaan bersifat ketat. Kelas layout harus diberi nama {KernelName}Layout dan didefinisikan di modul yang sama dengan KernelName.

Transformasi parameter

Gunakan pola ini ketika kernel mengharapkan bobot dengan nama atau bentuk yang berbeda dari checkpoint model asli.

Kelas KernelNameLayout memiliki tanda tangan __init__ yang sama dengan modul yang digantikannya dan mendeklarasikan atribut kelas conversion_mapping yang memberi tahu Transformers cara memetakan ulang kunci checkpoint ke nama parameter baru (lihat Pem bobotan dinamis untuk detail lebih lanjut).

# 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

Perpustakaan kernels menemukan kelas kernel melalui modul layers (baik layers.py maupun layers/__init__.py). Definisikan kelas kernel dan layout secara langsung di layers.py, dan paparkan modul layers dari __init__.py paket. Transformers menemukan CustomRMSNormLayout di modul yang sama dengan CustomRMSNorm.

Lindungi impor yang spesifik untuk Transformers (WeightConverter, Concatenate, WeightRenaming) di layers.py agar kernel juga dapat diimpor di lingkungan tempat Transformers tidak terpasang.

Muat kernel ini dengan meneruskan nama repo dan kelas ke KernelConfig. Kuncinya adalah nama kelas modul asli dari model. Nilainya menunjuk ke kelas KernelName (bukan Layout) di repo.

Nilainya berupa string "owner/repo:ClassName" atau tuple (repo_string, metadata). Dalam bentuk tuple, berikan version atau revision (minimal satu), dan secara opsional trust_remote_code. Gunakan version untuk versi kernel mayor atau revision untuk tag, branch, atau commit. Atur trust_remote_code=True untuk memuat dari luar organisasi kernels-community yang tepercaya. String polos secara default menggunakan version=1 dan 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",
)

Saat model dimuat, Transformers:

1. Memuat CustomRMSNorm dari repo dan mencari CustomRMSNormLayout di modul yang sama.
2. Monkey-patches setiap RMSNorm dalam model dengan CustomRMSNormLayout.
3. Memetakan ulang bobot checkpoint menggunakan conversion_mapping agar dapat dimuat ke nama parameter baru.
4. Memanggil kernelize, yang menggantikan CustomRMSNormLayout.forward dengan CustomRMSNorm.forward.

Fusi modul

Gunakan pola ini ketika sebuah kernel menggantikan beberapa modul bersebelahan dengan satu implementasi yang terpadu. Karena modul yang terpadu menggabungkan parameter dari beberapa modul asli, KernelNameLayout.__init__ menerima modul anak yang sudah diinstansiasi, bukan argumen konstruktor mereka.

# 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

Untuk menggabungkan modul, teruskan tuple pasangan (class_name, path_pattern) sebagai kunci di KernelConfig alih-alih string biasa. Semua pola harus berbagi modul induk yang sama (Transformers menggabungkan anak-anak di dalam induk tersebut). Wildcard * cocok dengan satu segmen path apa pun.

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",
)

Saat model dimuat, Transformers:

1. Memuat RMSNormMLP dari repo dan menemukan RMSNormMLPLayout di modul yang sama.
2. Mencocokkan setiap lapisan decoder di model.layers.* dan membangun kelas induk yang digabungkan (fused) yang __init__-nya memanggil RMSNormMLPLayout(post_attention_layernorm, mlp).
3. Menggantikan anak yang tersisa (mlp) dengan nn.Identity() untuk mempertahankan antarmuka modul induk.
4. Memetakan ulang bobot checkpoint menggunakan conversion_mapping.
5. Memanggil kernelize, yang menggantikan RMSNormMLPLayout.forward dengan RMSNormMLP.forward.

Urutan pasangan dalam tuple fusi menentukan urutan argumen yang diteruskan ke KernelNameLayout.__init__.