chore: import upstream snapshot with attribution
Secret Leaks / trufflehog (push) Failing after 1s
Build documentation / build (push) Failing after 1s
Build documentation / build_other_lang (push) Failing after 0s
CodeQL Security Analysis / CodeQL Analysis (push) Failing after 0s
PR CI / pr-ci (push) Failing after 1s
Slow tests on important models (on Push - A10) / Get all modified files (push) Failing after 1s
Slow tests on important models (on Push - A10) / Model CI (push) Has been skipped
Self-hosted runner (benchmark) / Benchmark (aws-g5-4xlarge-cache) (push) Has been cancelled
New model PR merged notification / Notify new model (push) Has been cancelled
Update Transformers metadata / build_and_package (push) Has been cancelled
Secret Leaks / trufflehog (push) Failing after 1s
Build documentation / build (push) Failing after 1s
Build documentation / build_other_lang (push) Failing after 0s
CodeQL Security Analysis / CodeQL Analysis (push) Failing after 0s
PR CI / pr-ci (push) Failing after 1s
Slow tests on important models (on Push - A10) / Get all modified files (push) Failing after 1s
Slow tests on important models (on Push - A10) / Model CI (push) Has been skipped
Self-hosted runner (benchmark) / Benchmark (aws-g5-4xlarge-cache) (push) Has been cancelled
New model PR merged notification / Notify new model (push) Has been cancelled
Update Transformers metadata / build_and_package (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,168 @@
|
||||
<!--Copyright 2026 The HuggingFace Team. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||
the License. You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations under the License.
|
||||
|
||||
⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
|
||||
rendered properly in your Markdown viewer.
|
||||
|
||||
-->
|
||||
|
||||
# Writing kernels
|
||||
|
||||
This guide explains how to write kernels that go beyond a stateless `forward` replacement. It covers two capabilities the extended `KernelConfig` API supports:
|
||||
|
||||
1. Parameter transformation: the kernel expects weights in a different layout than the original model (for example, renamed or merged parameters).
|
||||
2. Module fusion: the kernel replaces multiple adjacent modules with a single fused implementation.
|
||||
|
||||
For basic kernels (stateless `forward` replacements with no parameter changes), see the [kernels](https://github.com/huggingface/kernels) library documentation.
|
||||
|
||||
## Two-class pattern
|
||||
|
||||
Any kernel that carries its own parameters follows a two-class pattern.
|
||||
|
||||
- `KernelName`: contains only the `forward` pass. The `kernels` library uses this class to kernelize the model because it does not allow stateful kernel classes.
|
||||
- `KernelNameLayout`: an `nn.Module` that holds the parameters and monkey-patches the original module before the checkpoint is loaded. At runtime, `kernelize` replaces its `forward` with the `forward` from `KernelName`'. You do not need to define `forward`. Transformers injects one automatically with the same signature as `KernelName.forward`.
|
||||
|
||||
> [!IMPORTANT]
|
||||
|
||||
The naming convention is strict. The layout class must be named `{KernelName}Layout` and defined in the same module as `KernelName`.
|
||||
|
||||
## Parameter transformation
|
||||
|
||||
Use this pattern when the kernel expects weights under different names or in a different shape than the original model checkpoint.
|
||||
|
||||
The `KernelNameLayout` class has the same `__init__` signature as the module it replaces and declares a `conversion_mapping` class attribute that tells Transformers how to remap checkpoint keys to the new parameter names (see [Dynamic weight loading](../weightconverter) for more details).
|
||||
|
||||
```python
|
||||
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)
|
||||
|
||||
|
||||
class layers:
|
||||
CustomRMSNorm = CustomRMSNorm
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> The `layers` class is required by the `kernels` library to expose the kernel entry point.
|
||||
|
||||
Load this kernel by passing the repo and class name to [`KernelConfig`]. The key is the original module class name from the model. The value points to the `KernelName` class (not the `Layout`) in the repo.
|
||||
|
||||
```python
|
||||
from transformers import AutoModelForCausalLM, KernelConfig
|
||||
|
||||
kernel_config = KernelConfig({"RMSNorm": "owner/my-kernel:CustomRMSNorm"})
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
use_kernels=True,
|
||||
kernel_config=kernel_config,
|
||||
device_map="cuda",
|
||||
)
|
||||
```
|
||||
|
||||
When the model loads, Transformers:
|
||||
1. Loads `CustomRMSNorm` from the repo and looks for `CustomRMSNormLayout` in the same module.
|
||||
2. Monkey-patches every `RMSNorm` in the model with `CustomRMSNormLayout`.
|
||||
3. Remaps checkpoint weights using `conversion_mapping` so they load into the new parameter names.
|
||||
4. Calls `kernelize`, which replaces `CustomRMSNormLayout.forward` with `CustomRMSNorm.forward`.
|
||||
|
||||
## Module fusion
|
||||
|
||||
Use this pattern when a kernel replaces multiple adjacent modules with a single fused implementation. Because the fused module combines parameters from several original modules, the `KernelNameLayout.__init__` receives the instantiated child modules rather than their constructor arguments.
|
||||
|
||||
```python
|
||||
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)
|
||||
|
||||
|
||||
class layers:
|
||||
RMSNormMLP = RMSNormMLP
|
||||
```
|
||||
|
||||
To fuse modules, pass a tuple of `(class_name, path_pattern)` pairs as the key in `KernelConfig` instead of a plain string. All patterns must share the same parent module (Transformers fuses the children in that parent). The `*` wildcard matches any single path segment.
|
||||
|
||||
```python
|
||||
from transformers import AutoModelForCausalLM, KernelConfig
|
||||
|
||||
kernel_config = KernelConfig(
|
||||
{
|
||||
(
|
||||
("RMSNorm", "model.layers.*.post_attention_layernorm"),
|
||||
("MLP", "model.layers.*.mlp"),
|
||||
): "owner/my-kernel:RMSNormMLP",
|
||||
}
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
use_kernels=True,
|
||||
kernel_config=kernel_config,
|
||||
device_map="cuda",
|
||||
)
|
||||
```
|
||||
|
||||
When the model loads, Transformers:
|
||||
1. Loads `RMSNormMLP` from the repo and finds `RMSNormMLPLayout` in the same module.
|
||||
2. Matches every decoder layer at `model.layers.*` and builds a fused parent class whose `__init__` calls `RMSNormMLPLayout(post_attention_layernorm, mlp)`.
|
||||
3. Replaces the remaining child (`mlp`) with `nn.Identity()` to preserve the parent module's interface.
|
||||
4. Remaps checkpoint weights using `conversion_mapping`.
|
||||
5. Calls `kernelize`, which replaces `RMSNormMLPLayout.forward` with `RMSNormMLP.forward`.
|
||||
|
||||
> [!TIP]
|
||||
> The order of pairs in the fusion tuple determines the argument order passed to `KernelNameLayout.__init__`.
|
||||
Reference in New Issue
Block a user