Buckets:

hf-doc-build/doc / transformers /main /en /kernel_doc /loading_kernels.md
|
download
raw
12.3 kB
# 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 Transformers with the supported version of the [kernels](https://github.com/huggingface/kernels) package.
```bash
pip install -U "transformers[kernels]"
```
Set `use_kernels=True` in [from_pretrained()](/docs/transformers/main/en/main_classes/model#transformers.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="auto"
)
```
The default mapping depends on the model, device, and execution mode. The table below lists common Hub repositories for
the devices shown. When no kernel is registered for a particular combination, 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` | — | — |
The table is not exhaustive. Models can register additional layers and functions, and some mappings are available only
for specific execution modes.
> [!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="auto"
)
```
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="auto"
)
```
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="auto"
)
# use semantic versioning constraints
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
attn_implementation="kernels-community/flash-attn2@>=2.0,<3.0",
device_map="auto"
)
```
### FlashAttention fallback
Requesting `attn_implementation="flash_attention_2"`, `"flash_attention_3"`, or `"flash_attention_4"` falls back to the matching Hub kernel when the compiled `flash-attn` package isn't installed or your device isn't CUDA.
> [!NOTE]
> FlashAttention-4 support is in beta. APIs and behavior may change.
```py
from transformers import AutoModelForCausalLM
# uses the compiled flash-attn package if present, otherwise the kernels-community/flash-attn2 Hub kernel
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
attn_implementation="flash_attention_2",
device_map="auto",
)
```
## 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="auto"
)
# 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 [kernelize()](/docs/transformers/main/en/main_classes/kernels#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](/docs/transformers/main/en/main_classes/kernels#transformers.KernelConfig) customizes which kernels are used in a model.
The `kernel_mapping` keys are names registered by the model. They can refer to a layer, such as `"RMSNorm"`, or a
registered function, such as `"rotary_pos_emb"`. The `:` separator names a specific kernel entry inside the repository
and maps it to that layer or function.
```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",
use_kernels=True,
kernel_config=kernel_config,
device_map="auto"
)
```
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/liger-kernels:LigerRMSNorm",
"xpu": "kernels-community/rmsnorm:RMSNorm"
}
}
)
```
### Kernel metadata
Add a metadata dict to control which kernel build is loaded.
| Option | Description | Default |
|---|---|---|
| `version` | Major version of the kernel repository. `2` loads the latest build on the repository's `v2` branch. | `1` |
| `revision` | Exact tag, branch, or commit to load instead of a version. | --- |
| `trust_remote_code` | Allows a repository outside the trusted `kernels-community` organization. Loading a kernel runs code from that repository on your machine. | `False` |
```py
from transformers import KernelConfig
kernel_config = KernelConfig(
kernel_mapping={
"RMSNorm": ("kernels-community/liger-kernels:LigerRMSNorm", {"version": 3}),
}
)
```
Older branches of a kernel repository may not have builds for your PyTorch and CUDA versions. If a kernel fails to load, try a newer `version` before concluding your hardware is unsupported.
### Inherited mappings
A `KernelConfig` inherits the default Transformers kernel mapping for Hub kernels, and entries in `kernel_mapping`
override the default for the corresponding layers and functions. Set `inherit_mapping=False` to use only the entries in `kernel_mapping`. Everything you leave out falls back to standard PyTorch. This is useful when benchmarking a
specific kernel or testing a custom implementation without applying the other defaults.
The configuration below maps the `rotary_pos_emb` function to a RoPE kernel and leaves every other operation on
PyTorch. The function name must be registered by the model.
```py
from transformers import AutoModelForCausalLM, KernelConfig
kernel_config = KernelConfig(
kernel_mapping={
"rotary_pos_emb": (
"kernels-community/rotary:apply_rotary_transformers",
{"version": 2},
),
},
inherit_mapping=False,
)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
use_kernels=True,
kernel_config=kernel_config,
device_map="auto",
)
```
## Module fusion
Fuse adjacent modules into a single kernel by passing a tuple of `(class_name, path_pattern)` pairs as the key in [KernelConfig](/docs/transformers/main/en/main_classes/kernels#transformers.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="auto",
)
```
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](/docs/transformers/main/en/main_classes/kernels#transformers.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
)
```
## Disable kernels globally
Set the `USE_HUB_KERNELS` environment variable to disable Hub kernels everywhere without changing your code.
```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 mean the kernels package is missing or its version falls outside the range Transformers supports. Reinstall through the extra to get a compatible version.
```bash
pip install -U "transformers[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

Xet Storage Details

Size:
12.3 kB
·
Xet hash:
e9496a6e9d9de7c7d1f07c6b1db994a1945ab6eeee3bdcf7ccead1c2a65539ea

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