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,295 @@
|
||||
<!--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.
|
||||
|
||||
-->
|
||||
|
||||
# Loading kernels
|
||||
|
||||
A kernel works as a drop-in replacement for standard PyTorch operations. It swaps the `forward` method with the optimized kernel implementation without breaking model code.
|
||||
|
||||
This guide shows how to load kernels to accelerate inference.
|
||||
|
||||
Install the kernels package. We recommend the latest version which provides the best performance and bug fixes.
|
||||
|
||||
> [!NOTE]
|
||||
> kernels >=0.11.0 is the minimum required version for working with Transformers.
|
||||
|
||||
```bash
|
||||
pip install -U kernels
|
||||
```
|
||||
|
||||
Set `use_kernels=True` in [`~PreTrainedModel.from_pretrained`] to load the most performant kernels available on the Hub for your device. This replaces supported PyTorch operations with the kernel implementation.
|
||||
|
||||
```py
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
use_kernels=True,
|
||||
device_map="cuda"
|
||||
)
|
||||
```
|
||||
|
||||
The default kernels differ by device type. The table below lists the Hub repository that supplies each operation's default kernel. When no default kernel is registered, the operation falls back to standard PyTorch.
|
||||
|
||||
| Operation | NVIDIA (CUDA) | AMD (ROCm) | Intel (XPU) |
|
||||
|---|---|---|---|
|
||||
| RMSNorm | `kernels-community/liger-kernels` | `kernels-community/liger-kernels` | `kernels-community/rmsnorm` |
|
||||
| MoE MLP | `kernels-community/megablocks` | `kernels-community/megablocks` | `kernels-community/megablocks` |
|
||||
| MLP (SwiGLU, GeGLU) | `kernels-community/liger-kernels` | — | — |
|
||||
| Linear | `kernels-community/liger-kernels` | — | — |
|
||||
| Activations (GELU variants, SiLU) | `kernels-community/activation` | — | — |
|
||||
| Rotary embeddings | `kernels-community/rotary` | `kernels-community/aiter-rope` | `kernels-community/rotary` |
|
||||
| Causal LM loss | `kernels-community/liger-kernels` | — | — |
|
||||
| Deformable attention | `kernels-community/deformable-detr` | — | — |
|
||||
|
||||
> [!NOTE]
|
||||
> AMD GPUs report their device type as `cuda` in PyTorch. Transformers detects ROCm at runtime and routes supported operations to the AMD kernels above, including [AITER](https://github.com/ROCm/aiter) builds such as `kernels-community/aiter-rope`. You don't need to set the device type yourself.
|
||||
|
||||
Browse available kernels in the [kernels-community](https://huggingface.co/kernels-community) organization.
|
||||
|
||||
## Attention kernels
|
||||
|
||||
Load attention kernels from the Hub with the `attn_implementation` argument.
|
||||
|
||||
```py
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
attn_implementation="kernels-community/flash-attn2",
|
||||
device_map="cuda"
|
||||
)
|
||||
```
|
||||
|
||||
Note that for attention kernels, anything that is not part of the `kernels-community` repository (which is trusted - we may add more trusted repositories in the future) will require an additional `allow_all_kernels=True` kwarg to be used (similar to the `trust_remote_code=True` kwarg for non-HF models). This is because loading a kernel can lead to arbitrary code execution on the host machine, and we cannot verify every repo, so you need to explicitly allow it.
|
||||
|
||||
```py
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
attn_implementation="random-repo/random-attention",
|
||||
allow_all_kernels=True,
|
||||
device_map="cuda"
|
||||
)
|
||||
```
|
||||
|
||||
Specific kernels, like attention, accept several formats.
|
||||
|
||||
- `@v2.1.0` pins to a specific tag or branch.
|
||||
- `@>=2.0,<3.0` sets semantic versioning constraints.
|
||||
|
||||
```py
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
# pin to a specific version
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
attn_implementation="kernels-community/flash-attn2@v2.1.0",
|
||||
device_map="cuda"
|
||||
)
|
||||
# use semantic versioning constraints
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
attn_implementation="kernels-community/flash-attn2@>=2.0,<3.0",
|
||||
device_map="cuda"
|
||||
)
|
||||
```
|
||||
|
||||
## Mode-awareness
|
||||
|
||||
Kernels automatically adapt to [training](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.train) and [inference](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.eval) modes based on PyTorch's `model.training` state.
|
||||
|
||||
```py
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
use_kernels=True,
|
||||
device_map="cuda"
|
||||
)
|
||||
|
||||
# Switch to inference mode - uses inference-optimized kernels
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
output = model.generate(input_ids, max_new_tokens=50)
|
||||
|
||||
# Switch to training mode - uses training-optimized kernels with gradient support
|
||||
model.train()
|
||||
loss = model(input_ids, labels=labels).loss
|
||||
loss.backward()
|
||||
```
|
||||
|
||||
Explicitly enable training and inference modes with the `mode` argument in the [`~transformers.kernelize`] function. Training mode also supports an additional torch.compile mode.
|
||||
|
||||
```py
|
||||
from kernels import Mode
|
||||
from transformers import kernelize
|
||||
|
||||
# inference optimized kernels
|
||||
kernelize(model, mode=Mode.INFERENCE)
|
||||
|
||||
# training optimized kernels
|
||||
kernelize(model, mode=Mode.TRAINING)
|
||||
|
||||
# training and torch-compile friendly kernels
|
||||
kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE)
|
||||
```
|
||||
|
||||
## KernelConfig
|
||||
|
||||
[`KernelConfig`] customizes which kernels are used in a model.
|
||||
|
||||
The `:` separator names a specific kernel entry inside the repository and maps it to a layer.
|
||||
|
||||
```py
|
||||
from transformers import AutoModelForCausalLM, KernelConfig
|
||||
|
||||
kernel_config = KernelConfig(
|
||||
kernel_mapping={
|
||||
"RMSNorm": "kernels-community/liger_kernels:LigerRMSNorm",
|
||||
}
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
attn_implementation="kernels-community/flash-attn2:FlashAttention2",
|
||||
use_kernels=True,
|
||||
kernel_config=kernel_config,
|
||||
device_map="cuda"
|
||||
)
|
||||
```
|
||||
|
||||
Specify different kernel implementations for each device type.
|
||||
|
||||
```py
|
||||
from transformers import KernelConfig
|
||||
|
||||
kernel_config = KernelConfig(
|
||||
kernel_mapping={
|
||||
"RMSNorm": {
|
||||
"cuda": "kernels-community/liger_kernels:LigerRMSNorm",
|
||||
"rocm": "kernels-community/rocm-kernels:RocmRMSNorm",
|
||||
"metal": "kernels-community/metal-kernels:MetalRMSNorm",
|
||||
"xpu": "kernels-community/xpu-kernels:XpuRMSNorm"
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
## Module fusion
|
||||
|
||||
Fuse adjacent modules into a single kernel by passing a tuple of `(class_name, path_pattern)` pairs as the key in [`KernelConfig`]. All patterns must share the same parent module. `*` 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/fused-rmsnorm-mlp:RMSNormMLP",
|
||||
}
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
use_kernels=True,
|
||||
kernel_config=kernel_config,
|
||||
device_map="cuda",
|
||||
)
|
||||
```
|
||||
|
||||
Fusion requires the kernel repo to provide a companion `KernelNameLayout` class alongside the `KernelName` class. See the [Writing kernels](./writing_kernels) guide for how to implement one.
|
||||
|
||||
## Local kernels
|
||||
|
||||
Load kernels from local file paths with `use_local_kernel=True` in [`KernelConfig`]. This loads from a local filesystem path instead of a Hub repository.
|
||||
|
||||
Local kernels use `/abs/path:layer_name` instead of the Hub format `org/repo:layer_name`.
|
||||
|
||||
```py
|
||||
from transformers import KernelConfig, AutoModelForCausalLM
|
||||
|
||||
kernel_mapping = {
|
||||
"RMSNorm": "/path/to/liger_kernels:LigerRMSNorm",
|
||||
}
|
||||
kernel_config = KernelConfig(kernel_mapping, use_local_kernel=True)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
dtype="auto",
|
||||
device_map="auto",
|
||||
use_kernels=True,
|
||||
kernel_config=kernel_config
|
||||
)
|
||||
```
|
||||
|
||||
## Disabling kernels
|
||||
|
||||
Disable kernels for specific layers with an empty kernel mapping in [`KernelConfig`].
|
||||
|
||||
```py
|
||||
from transformers import AutoModelForCausalLM, KernelConfig
|
||||
|
||||
kernel_config = KernelConfig(
|
||||
kernel_mapping={
|
||||
"RMSNorm": "",
|
||||
}
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Qwen/Qwen3-0.6B",
|
||||
use_kernels=True,
|
||||
kernel_config=kernel_config,
|
||||
device_map="cuda"
|
||||
)
|
||||
```
|
||||
|
||||
Set the environment variable to disable kernels globally.
|
||||
|
||||
```bash
|
||||
export USE_HUB_KERNELS=0 # or OFF or NO
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
Kernel integration depends on hardware, drivers, and package versions working together. The following sections cover common failures.
|
||||
|
||||
### Installation issues
|
||||
|
||||
Import errors indicate the kernels library isn't installed.
|
||||
|
||||
```bash
|
||||
pip install -U kernels
|
||||
```
|
||||
|
||||
### Kernel loading failures
|
||||
|
||||
If specific kernels fail to load, try the following.
|
||||
|
||||
- Check your hardware compatibility with the kernel requirements.
|
||||
- Verify your CUDA/ROCm/Metal drivers are up to date.
|
||||
- Consult the kernel repository documentation for known issues.
|
||||
|
||||
### Device compatibility
|
||||
|
||||
Not all kernels support all devices. The library falls back to standard PyTorch operations if a kernel is unavailable for your hardware. Check kernel repository documentation for device-specific support.
|
||||
|
||||
## Resources
|
||||
|
||||
- [Kernels](https://github.com/huggingface/kernels) repository
|
||||
- [Enhance Your Models in 5 Minutes with the Hugging Face Kernel Hub](https://huggingface.co/blog/hello-hf-kernels) blog post
|
||||
- Discover kernels in the [kernels-community](https://huggingface.co/kernels-community) org
|
||||
@@ -0,0 +1,48 @@
|
||||
<!--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.
|
||||
|
||||
-->
|
||||
|
||||
# Kernels
|
||||
|
||||
PyTorch operations are general-purpose. Hardware vendors and the community create specialized implementations that run faster on specific platforms. Installing these optimized kernels is a challenge because it requires matching compiler versions, CUDA toolkits, and platform-specific builds.
|
||||
|
||||
| platform | supported devices |
|
||||
| :--- | :--- |
|
||||
| NVIDIA GPUs (CUDA) | Modern architectures with compute capability 7.0+ (Volta, Turing, Ampere, Hopper, Blackwell) |
|
||||
| AMD GPUs (ROCm) | Compatible with ROCm-supported devices |
|
||||
| Apple Silicon (Metal) | M-series chips (M1, M2, M3, M4 and newer) |
|
||||
| Intel GPUs (XPU) | Intel Data Center GPU Max Series and compatible devices |
|
||||
|
||||
[Kernels](https://huggingface.co/docs/kernels/index) solves this by distributing precompiled binaries through the [Hub](https://huggingface.co/kernels-community). It detects your platform at runtime and loads the right binary automatically.
|
||||
|
||||
When `use_kernels=True`, Transformers identifies layers with available optimized kernel implementations. It downloads and [caches](../installation#cache-directory) kernels from the Hub only when needed to reduce startup time. Kernels accelerate compute-intensive operations such as attention, normalization, and fused operations.
|
||||
|
||||
Not all operations have kernel implementations. The library falls back to standard PyTorch when no kernel is available.
|
||||
|
||||
## Determinism
|
||||
|
||||
Some kernels produce slightly different results than PyTorch due to operation reordering or accumulation strategies. These differences are functionally equivalent but affect reproducibility.
|
||||
|
||||
For deterministic behavior, try the following.
|
||||
|
||||
- Check kernel repository documentation for determinism guarantees. For example, the SDPA kernel in [gpt-oss-metal-kernels](https://huggingface.co/kernels-community/gpt-oss-metal-kernels#4-scaled-dot-product-attention-sdpa) matches the PyTorch implementation 97% of the time.
|
||||
- Disable specific kernels that affect your use case.
|
||||
- Set random seeds and PyTorch deterministic flags.
|
||||
|
||||
## Resources
|
||||
|
||||
- [Loading kernels](./loading_kernels) guide to get started
|
||||
- [Kernels](https://github.com/huggingface/kernels) GitHub repository
|
||||
- [Enhance Your Models in 5 Minutes with the Hugging Face Kernel Hub](https://huggingface.co/blog/hello-hf-kernels) blog post
|
||||
@@ -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