Buckets:
| # 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.