transformers/pr_46771/en/kernel_doc/loading_kernels.md

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 package.

pip install -U "transformers[kernels]"

Set use_kernels=True in from_pretrained() to load the most performant kernels available on the Hub for your device. This replaces supported PyTorch operations with the kernel implementation.

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.

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 builds such as kernels-community/aiter-rope. You don't need to set the device type yourself.

Browse available kernels in the kernels-community organization.

Attention kernels

Load attention kernels from the Hub with the attn_implementation argument.

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.

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.

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.

FlashAttention-4 support is in beta. APIs and behavior may change.

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 and inference modes based on PyTorch's model.training state.

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() function. Training mode also supports an additional torch.compile mode.

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

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.

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

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.

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. All patterns must share the same parent module. * matches any single path segment.

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

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.

export USE_HUB_KERNELS=0  # or OFF or NO

Processor kernels

Image and video processors can run some of their operations with Hub kernels from kernels-community/cv-utils. Pass use_kernels=True when loading the processor. The flag is not saved with the processor configuration.

Operation	What the kernel does	Processors
Resize and normalize	Resizes a batch of images of different sizes in one launch, then rescales and normalizes in the same pass.	Processors that use the default TorchvisionBackend preprocessing, such as ViT, CLIP and SigLIP
Resize, normalize and patchify	Writes each resized and normalized pixel directly at its place in the flattened patch sequence, for images and video frames.	Qwen2-VL and Qwen3-VL image and video processors, and the processors built on them (Cohere Compass, ERNIE 4.5 VL, GLM-Image, HunYuanVL, MiniMax-M3-VL, PaddleOCR-VL)
Connected components	Measures the area of every connected region of a mask, to fill small holes and remove small islands.	SAM 2 and SAM 3 image processors (max_hole_area, max_sprinkle_area)

from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", use_kernels=True)
inputs = processor(text=prompt, images=images, return_tensors="pt", device="cuda")

AutoProcessor passes the flag to its image and video processors. Kernels run on the device of the inputs, so pass device="cuda". The processor keeps its default implementation when the inputs are on CPU, when the kernel fails to load, or when the kernel does not support the arguments.

The resize kernels round to uint8 after each pass, like the torchvision path, but in a different order. On photos the mean difference stays below a quarter of a uint8 level. Bicubic resampling can differ by several levels on a few pixels near sharp edges.

Adding a processing kernel

A processing kernel has two levels in integrations/hub_processing_kernels.py.

register_processing_kernel maps an operation to a kernel of _HUB_KERNEL_MAPPING. The adapter receives the loaded kernel module and the arguments of the operation. It translates them to the kernel call and returns None for arguments the kernel cannot handle.

from transformers.integrations.hub_processing_kernels import register_processing_kernel

@register_processing_kernel("my_op", kernel_name="my-processing-kernel")
def my_op_kernel(kernel, images):
    if images[0].device.type != "cuda":
        return None
    return kernel.my_op(images)

use_processing_kernel connects a processor to it. It decorates _preprocess with a function that receives the processor and the named arguments of _preprocess. That function owns everything the kernel needs, such as the target sizes and the output BatchFeature, so the processor code does not change. When use_kernels=True and the function returns a result, _preprocess returns it. Otherwise _preprocess runs as usual.

from transformers import BatchFeature, TorchvisionBackend
from transformers.integrations.hub_processing_kernels import run_processing_kernel, use_processing_kernel

def my_op_with_kernel(processor, images, return_tensors, **kwargs):
    pixel_values = run_processing_kernel("my_op", images)
    if pixel_values is None:
        return None
    return BatchFeature(data={"pixel_values": pixel_values}, tensor_type=return_tensors)

class MyImageProcessor(TorchvisionBackend):
    @use_processing_kernel(my_op_with_kernel)
    def _preprocess(self, images, return_tensors, **kwargs):
        ...

Subclasses that keep _preprocess keep the kernel. The function reads what differs between models from the processor, such as its _resized_size method or its patch_block_size attribute. A model that replaces _preprocess with a different pipeline marks it with @no_inherit_decorator.

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.

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 repository
• Enhance Your Models in 5 Minutes with the Hugging Face Kernel Hub blog post
• Discover kernels in the kernels-community org