Kernels

Purpose and Scope

Kernels in Transformers are optimized replacements for specific PyTorch operations. The kernel documentation frames them as targeted accelerators for matrix multiplications, attention, normalization, loss computation, activations, and similar hot paths. Instead of changing model architecture or user-facing model code, a kernel can swap the operation behind a module’s forward method. This matters when the bottleneck is not model quality but runtime efficiency: fewer GPU memory reads and writes, fewer individual operation launches, and fused implementations can improve throughput or reduce memory pressure during inference and training.

Sources: docs/source/en/kernels.md, docs/source/en/kernel_doc/loading_kernels.md

This page covers the public kernel workflow exposed by the Transformers documentation: loading default Hub kernels with use_kernels=True, selecting explicit kernels through KernelConfig, using Hub-hosted attention kernels through attn_implementation, enabling Liger Kernel from TrainingArguments, and writing custom kernels that need parameter conversion or module fusion. The source surface is intentionally split between a conceptual landing page, task guides for loading and authoring, and an API reference page that documents kernelize and KernelConfig.

Sources: docs/source/en/kernels.md, docs/source/en/kernel_doc/writing_kernels.md, docs/source/en/main_classes/kernels.md

Relevant Source Files

• docs/source/en/kernels.md - English landing page that explains what custom kernels are, shows Hub kernel loading with KernelConfig, documents the Liger training flag, and points readers to attention backends and torch.compile as next steps.
• docs/source/en/kernel_doc/loading_kernels.md - Task guide for installing the kernels package, enabling use_kernels=True, understanding device-specific default kernel repositories, loading attention kernels, and handling trust for non-community kernel repositories.
• docs/source/en/kernel_doc/writing_kernels.md - Authoring guide for kernels that need state, parameter transformation, checkpoint key remapping, or module fusion beyond a stateless forward replacement.
• docs/source/en/main_classes/kernels.md - API reference stub that publishes the documented kernel utilities, specifically kernelize and KernelConfig, through the Transformers documentation build.
• docs/source/zh/kernels.md - Chinese translation of the kernel landing page, preserving the same user-facing concepts, examples, Liger options, and next-step links for localized documentation.

Core Primitives

The central loading primitive is KernelConfig. The landing page shows it imported from transformers and constructed with a mapping from an original module class name, such as RMSNorm, to a Hub repository or repository-plus-class target. That config is passed to kernel_config in from_pretrained, while use_kernels=True opts the model into kernel replacement. Once loaded, the replacement is described as active for training, meaning the user continues to work with the normal model object rather than calling a separate execution API.

Sources: docs/source/en/kernels.md, docs/source/en/kernel_doc/loading_kernels.md

from transformers import AutoModelForCausalLM, KernelConfig
 
kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": "kernels-community/rmsnorm",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
 )

The automatic path is even shorter. The loading guide instructs readers to install the kernels package, with kernels >=0.11.0 as the minimum version for Transformers, and then pass use_kernels=True to PreTrainedModel.from_pretrained. Transformers selects the most performant registered kernels for the user’s device when defaults exist, and falls back to standard PyTorch when no default is registered. That fallback behavior is important because it keeps a model load usable even when only some operations have optimized implementations.

Sources: docs/source/en/kernel_doc/loading_kernels.md

pip install -U kernels

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="cuda"
)

Loading and Runtime Selection

Default kernel selection is device-aware. The loading guide lists default Hub repositories for NVIDIA CUDA, AMD ROCm, and Intel XPU across operations including RMSNorm, MoE MLP, SwiGLU or GeGLU MLPs, Linear, GELU and SiLU activations, rotary embeddings, causal language modeling loss, and deformable attention. AMD deserves special mention because PyTorch reports AMD GPU device type as cuda; the docs state that Transformers detects ROCm at runtime and routes supported operations to AMD-specific kernels such as kernels-community/aiter-rope without requiring the user to set a separate device type.

Sources: docs/source/en/kernel_doc/loading_kernels.md

Attention kernels use a separate loading path. Instead of kernel_config, the guide shows passing a Hub kernel identifier to attn_implementation in from_pretrained, for example kernels-community/flash-attn2. This keeps attention backend selection close to the rest of the model loading configuration while still treating external kernel code as a privileged execution path. For attention kernels outside trusted repositories such as kernels-community, the user must also pass allow_all_kernels=True, similar in spirit to trust_remote_code=True, because loading arbitrary kernel repositories can execute code on the host machine.

Sources: docs/source/en/kernel_doc/loading_kernels.md

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    device_map="cuda"
)

Liger Kernel for Training

The kernel landing page also documents Liger Kernel as a packaged training optimization. Liger fuses layers such as RMSNorm, RoPE, SwiGLU, CrossEntropy, and FusedLinearCrossEntropy into Triton kernels. The documented benefit is strongest for multi-GPU training: compatibility with FlashAttention, FSDP, and DeepSpeed, plus higher throughput and lower memory usage. In practice, this is a different workflow from picking a single Hub kernel mapping. Users install liger-kernel, then enable model patching through TrainingArguments rather than passing a KernelConfig to model loading.

Sources: docs/source/en/kernels.md, docs/source/zh/kernels.md

pip install liger-kernel

from transformers import TrainingArguments
 
training_args = TrainingArguments(
    ...,
    use_liger_kernel=True
)

For finer control, TrainingArguments accepts liger_kernel_config as a dictionary. The documentation notes that available options vary by model, but lists names such as rope, swiglu, cross_entropy, fused_linear_cross_entropy, and rms_norm. This is useful when a training run benefits from some fused kernels but needs to disable others for model support, debugging, numerical comparison, or compatibility with another optimization stack. The public configuration remains declarative: the trainer arguments express what should be patched, and the integration handles the replacement.

Sources: docs/source/en/kernels.md

from transformers import TrainingArguments
 
training_args = TrainingArguments(
    ...,
    use_liger_kernel=True,
    liger_kernel_config={
        "rope": True,
        "cross_entropy": True,
        "rms_norm": False,
        "swiglu": True,
    }
)

Writing Custom Kernels

The writing guide separates basic stateless kernels from kernels that require Transformers-specific layout handling. Basic kernels that only replace a forward implementation belong primarily to the external kernels library documentation. The Transformers guide focuses on two extended cases: parameter transformation, where the optimized kernel expects weights under different names or shapes than the original checkpoint, and module fusion, where one optimized implementation replaces multiple adjacent modules. These cases need coordination before checkpoint loading so that model parameters can land in the layout expected by the kernel.

Sources: docs/source/en/kernel_doc/writing_kernels.md

Stateful custom kernels use a strict two-class pattern. KernelName contains only the forward pass and is the class used by the kernels library during kernelization. KernelNameLayout is an nn.Module that holds parameters and monkey-patches the original module before checkpoint load. The naming rule is explicit: the layout class must be named {KernelName}Layout and defined in the same module as KernelName. At runtime, kernelize replaces the layout’s forward with the forward from KernelName, and Transformers injects a matching signature rather than requiring the layout class to define forward itself.

Sources: docs/source/en/kernel_doc/writing_kernels.md, docs/source/en/main_classes/kernels.md

Parameter transformation is declared with a conversion_mapping class attribute on the layout class. The guide’s RMSNorm example gives the layout the same __init__ signature as the module it replaces, defines new parameters such as scale, and stores values such as variance_epsilon. The kernel class then implements the normalized computation in forward. The repository and class name are passed through KernelConfig; the key is the original model module class name, and the value points to the kernel class, not the layout class.

Sources: docs/source/en/kernel_doc/writing_kernels.md

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",
)

API Components and Reference

The documented API surface is compact but powerful. KernelConfig is the user-facing configuration object for Hub kernels and custom kernel mappings. kernelize is the lower-level utility documented in the main classes reference and described by the writing guide as the runtime mechanism that replaces a layout module’s forward with the stateless kernel forward. use_kernels is the model-loading switch for automatic or configured non-attention kernels, kernel_config carries the mapping, attn_implementation selects Hub-hosted attention implementations, and allow_all_kernels is the explicit trust gate for attention kernels outside trusted repositories.

Sources: docs/source/en/main_classes/kernels.md, docs/source/en/kernel_doc/loading_kernels.md, docs/source/en/kernel_doc/writing_kernels.md

Practical Next Steps

Start with use_kernels=True if you want the safest high-level path and are comfortable with device-specific defaults. Move to KernelConfig when you need to select a particular Hub kernel, test a custom kernel, or map a model module such as RMSNorm to a known implementation. Use attn_implementation for attention-specific kernels, and be deliberate about allow_all_kernels=True because it changes the trust boundary. For training workloads, evaluate Liger separately through TrainingArguments, especially when using FlashAttention, FSDP, or DeepSpeed.

Sources: docs/source/en/kernels.md, docs/source/en/kernel_doc/loading_kernels.md

If you are implementing a kernel, decide first whether it is a stateless forward replacement or whether it needs parameters, checkpoint conversion, or module fusion. Stateless replacements can follow the external kernels library path, while stateful replacements should use the two-class KernelName and KernelNameLayout pattern documented here. Keep the layout naming convention exact, expose the kernel entry point as required by the kernels library, and test loading through KernelConfig against a real checkpoint before relying on it in training or production inference.

Sources: docs/source/en/kernel_doc/writing_kernels.md