编写内核

本指南介绍如何编写超越无状态 forward 替代方案的 kernel。它涵盖了扩展的 KernelConfig API 所支持的两项功能：

参数转换：内核期望的权重布局与原始模型不同（例如，参数被重命名或合并）。
模块融合：该内核将多个相邻模块替换为单个融合实现。

对于基础内核（无状态 forward 替换且不涉及参数更改），请参阅 kernels 库文档。

两类模式

任何携带自身参数的内核都遵循双类模式。

KernelName：仅包含 forward 前向传播。kernels 库使用此类对模型进行 kernel 化，因为它不允许使用有状态的 kernel 类。
KernelNameLayout：一个 nn.Module，用于保存参数，并在加载检查点之前对原始模块进行 monkey-patch。在运行时，kernelize 会将其 forward 替换为来自 KernelName 的 forward。你无需定义 forward。Transformers 会自动注入一个与 KernelName.forward 具有相同签名的 forward。

命名约定是严格的。布局类必须命名为 {KernelName}Layout，并且定义在与 KernelName 相同的模块中。

参数转换

当内核期望的权重名称或形状与原始模型检查点不同时，请使用此模式。

KernelNameLayout 类具有与其所替换模块相同的 __init__ 签名，并声明了一个 conversion_mapping 类属性，用于告知 Transformers 如何将检查点键重新映射到新的参数名称（有关更多详细信息，请参阅动态权重加载）。

# 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）来发现 kernel 类。请在 layers.py 中直接定义 kernel 类和 layout 类，并从包的 __init__.py 中暴露 layers 模块。Transformers 会在与 CustomRMSNorm 相同的模块中找到 CustomRMSNormLayout。

在 layers.py 中，将 Transformers 特有的导入（WeightConverter、Concatenate、WeightRenaming）进行保护，以便在未安装 Transformers 的环境中也能导入该 kernel。

通过向 KernelConfig 传入仓库和类名来加载此内核。键是模型中的原始模块类名。值指向仓库中的 KernelName 类（而非 Layout）。

该值是一个 "owner/repo:ClassName" 字符串，或一个元组 (repo_string, metadata)。在元组形式中，请传入 version 或 revision（至少一个），并可选地传入 trust_remote_code。对于主要内核版本，请使用 version；对于标签、分支或提交，请使用 revision。设置 trust_remote_code=True 可从受信任的 kernels-community 组织之外加载。仅使用字符串时，默认值为 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。
使用 CustomRMSNormLayout 对模型中的每个 RMSNorm 进行 monkey-patch。
使用 conversion_mapping 重新映射检查点权重，使其能够加载到新的参数名称中。
调用 kernelize，将 CustomRMSNormLayout.forward 替换为 CustomRMSNorm.forward。

模块融合

当某个 kernel 将多个相邻模块替换为单个融合实现时，请使用此模式。由于融合模块会合并来自多个原始模块的参数，因此 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