# 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](https://github.com/huggingface/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`.

> [!IMPORTANT]
> 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](../weightconverter) để biết thêm chi tiết).

```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]
> 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](/docs/transformers/main/vi/main_classes/kernels#transformers.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`.

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

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.

```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
```

Để 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.

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

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`.

> [!TIP]
> 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__`.