Viết kernel

Hướng dẫn này giải thích cách viết các kernel vượt ra ngoài phạm vi thay thế forward không trạng thái. Nó bao gồm hai khả năng mà API KernelConfig mở rộng hỗ trợ:

1. Biến đổi tham số: kernel kỳ vọng trọng số ở một bố cục khác so với mô hình gốc (ví dụ, các tham số đã được đổi tên hoặc gộp lại).
2. Module fusion: kernel thay thế nhiều module liền kề bằng một triển khai đã hợp nhất duy nhất.

Đối với các kernel cơ bản (thay thế forward không có trạng thái và không có thay đổi tham số), hãy xem tài liệu thư viện kernels.

Mẫu hai lớp

Bất kỳ kernel nào mang theo các tham số của riêng nó đều tuân theo mô hình hai lớp.

• KernelName: chỉ chứa bước forward. Thư viện kernels sử dụng lớp này để kernel hóa mô hình vì nó không cho phép các lớp kernel có trạng thái.
• KernelNameLayout: một nn.Module giữ các tham số và monkey-patch module gốc trước khi checkpoint được nạp. Khi chạy, kernelize thay thế forward của nó bằng forward từ KernelName. Bạn không cần định nghĩa forward. Transformers tự động tiêm một hàm có cùng chữ ký với KernelName.forward.

Quy ước đặt tên là nghiêm ngặt. Lớp layout phải được đặt tên là {KernelName}Layout và được định nghĩa trong cùng module với KernelName.

Biến đổi tham số

Sử dụng mẫu này khi kernel kỳ vọng trọng số có tên khác hoặc có hình dạng khác so với checkpoint mô hình gốc.

Lớp KernelNameLayout có cùng chữ ký __init__ với module mà nó thay thế và khai báo thuộc tính lớp conversion_mapping để chỉ cho Transformers cách ánh xạ lại các khóa checkpoint sang tên tham số mới (xem Tải trọng động để biết thêm chi tiết).

# 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

Thư viện kernels phát hiện các lớp kernel thông qua module layers (chẳng hạn layers.py hoặc layers/__init__.py). Hãy định nghĩa cả lớp kernel và lớp layout trực tiếp trong layers.py, và phơi bày module layers từ __init__.py của gói. Transformers tìm thấy CustomRMSNormLayout trong cùng module với CustomRMSNorm.

Bảo vệ các import đặc thù của Transformers (WeightConverter, Concatenate, WeightRenaming) trong layers.py để kernel cũng có thể được import trong các môi trường không cài đặt Transformers.

Tải kernel này bằng cách truyền tên repo và tên class vào KernelConfig. Khóa là tên class module gốc từ model. Giá trị trỏ đến class KernelName (không phải Layout) trong repo.

Giá trị là một chuỗi "owner/repo:ClassName" hoặc một tuple (repo_string, metadata). Ở dạng tuple, hãy truyền version hoặc revision (ít nhất một), và tùy chọn trust_remote_code. Dùng version cho một phiên bản kernel chính hoặc revision cho một tag, branch, hoặc commit. Đặt trust_remote_code=True để tải từ bên ngoài tổ chức kernels-community được tin cậy. Một chuỗi đơn giản mặc định là version=1 và 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",
)

Khi mô hình được tải, Transformers:

1. Tải CustomRMSNorm từ kho và tìm CustomRMSNormLayout trong cùng một module.
2. Áp dụng monkey-patch cho mọi RMSNorm trong mô hình bằng CustomRMSNormLayout.
3. Áp dụng lại các trọng số checkpoint bằng conversion_mapping để chúng được nạp vào các tên tham số mới.
4. Gọi kernelize, thay thế CustomRMSNormLayout.forward bằng CustomRMSNorm.forward.

Nối module

Sử dụng mẫu này khi một kernel thay thế nhiều module liền kề bằng một triển khai đã hợp nhất (fused). Vì module đã hợp nhất kết hợp các tham số từ nhiều module gốc, nên KernelNameLayout.__init__ nhận các module con đã được khởi tạo thay vì các đối số khởi tạo của chúng.

# 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

Để hợp nhất các module, hãy truyền một tuple gồm các cặp (class_name, path_pattern) làm key trong KernelConfig thay vì một chuỗi thông thường. Tất cả các pattern phải cùng chia sẻ một module cha (Transformers sẽ hợp nhất các con trong module cha đó). Wildcard * khớp với bất kỳ đoạn path đơn nào.

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

Khi mô hình được tải, Transformers:

1. Tải RMSNormMLP từ kho và tìm RMSNormMLPLayout trong cùng một module.
2. Trùng khớp với mọi lớp decoder tại model.layers.* và xây dựng một lớp cha đã hợp nhất, trong đó __init__ gọi RMSNormMLPLayout(post_attention_layernorm, mlp).
3. Thay thế con còn lại (mlp) bằng nn.Identity() để giữ nguyên giao diện của module cha.
4. Áp dụng lại các trọng số checkpoint bằng conversion_mapping.
5. Gọi kernelize, thay thế RMSNormMLPLayout.forward bằng RMSNormMLP.forward.

Thứ tự của các cặp trong bộ hợp nhất (fusion tuple) xác định thứ tự của các đối số được truyền vào KernelNameLayout.__init__.