Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / transformers /main /en /kernel_doc /writing_kernels.md

HuggingFaceDocBuilder

2 days ago

preview code

download

raw

8.33 kB

# 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
# 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]
> The `kernels` library discovers kernel classes through a `layers` module (either `layers.py` or `layers/__init__.py`). Define both the kernel and layout classes directly in `layers.py`, and expose the `layers` module from the package's `__init__.py`. Transformers finds `CustomRMSNormLayout` in the same module as `CustomRMSNorm`.
>
> Guard the Transformers-specific imports (`WeightConverter`, `Concatenate`, `WeightRenaming`) in `layers.py` so the kernel can also be imported in environments where Transformers is not installed.

Load this kernel by passing the repo and class name to [KernelConfig](/docs/transformers/main/en/main_classes/kernels#transformers.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.

The value is either an `"owner/repo:ClassName"` string or a tuple `(repo_string, metadata)`. In the tuple form, pass `version` or `revision` (at least one), and optionally `trust_remote_code`. Use `version` for a major kernel version or `revision` for a tag, branch, or commit. Set `trust_remote_code=True` to load from outside the trusted `kernels-community` organization. A bare string defaults to `version=1` and `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",
)
```

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

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", {"trust_remote_code": True, "version": 1}),
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)
```

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

Xet Storage Details

Size:: 8.33 kB
Xet hash:: 564ca9ac1923a78afb149aeeb601f3005d024dd1b22983bff3b5fd7c96f94544

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# 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
	# 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]
	> The `kernels` library discovers kernel classes through a `layers` module (either `layers.py` or `layers/__init__.py`). Define both the kernel and layout classes directly in `layers.py`, and expose the `layers` module from the package's `__init__.py`. Transformers finds `CustomRMSNormLayout` in the same module as `CustomRMSNorm`.
	>
	> Guard the Transformers-specific imports (`WeightConverter`, `Concatenate`, `WeightRenaming`) in `layers.py` so the kernel can also be imported in environments where Transformers is not installed.

	Load this kernel by passing the repo and class name to [KernelConfig](/docs/transformers/main/en/main_classes/kernels#transformers.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.

	The value is either an `"owner/repo:ClassName"` string or a tuple `(repo_string, metadata)`. In the tuple form, pass `version` or `revision` (at least one), and optionally `trust_remote_code`. Use `version` for a major kernel version or `revision` for a tag, branch, or commit. Set `trust_remote_code=True` to load from outside the trusted `kernels-community` organization. A bare string defaults to `version=1` and `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",
	)
	```

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

	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", {"trust_remote_code": True, "version": 1}),
	}
	)
	model = AutoModelForCausalLM.from_pretrained(
	"Qwen/Qwen3-0.6B",
	use_kernels=True,
	kernel_config=kernel_config,
	device_map="auto",
	)
	```

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