加载内核

Kernel 可作为标准 PyTorch 操作的直接替代方案。它在不破坏模型代码的情况下，将 forward 方法替换为优化后的 kernel 实现。

本指南介绍如何加载内核以加速推理。

使用受支持的 kernels 包版本安装 Transformers。

pip install -U "transformers[kernels]"

在 from_pretrained() 中设置 use_kernels=True，即可加载 Hub 上适用于您设备的高性能内核。这将把受支持的 PyTorch 操作替换为对应的内核实现。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

默认映射取决于模型、设备和执行模式。下表列出了所示设备的常见 Hub 仓库。当某个特定组合未注册内核时，该操作将回退到标准 PyTorch。

操作	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`	—	—
激活函数（GELU 变体、SiLU）	`kernels-community/activation`	—	—
旋转嵌入	`kernels-community/rotary`	`kernels-community/aiter-rope`	`kernels-community/rotary`
因果语言模型损失	`kernels-community/liger-kernels`	—	—
可变形注意力	`kernels-community/deformable-detr`	—	—

该表格并非详尽无遗。模型可以注册额外的层和函数，并且某些映射仅在特定的执行模式下可用。

AMD GPU 在 PyTorch 中将其设备类型报告为 cuda。Transformers 会在运行时检测 ROCm，并将受支持的操作路由到上述 AMD 内核，包括 AITER 构建，例如 kernels-community/aiter-rope。您无需自行设置设备类型。

浏览 kernels-community 组织中的可用内核。

注意力内核

使用 attn_implementation 参数从 Hub 加载注意力内核。

from transformers import AutoModelForCausalLM

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

请注意，对于注意力内核，任何不属于 kernels-community 仓库（该仓库受信任——我们未来可能会添加更多受信任的仓库）的内容，都需要额外使用 allow_all_kernels=True 关键字参数（类似于非 HF 模型的 trust_remote_code=True 关键字参数）。这是因为加载内核可能导致在主机上执行任意代码，而我们无法验证每个仓库，因此需要您显式允许。

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

某些特定的内核（例如 attention）支持多种格式。

@v2.1.0 固定到特定的标签或分支。
@>=2.0,<3.0 用于设置语义化版本约束。

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 回退

请求 attn_implementation="flash_attention_2"、"flash_attention_3" 或 "flash_attention_4" 时，如果未安装编译好的 flash-attn 包，或者您的设备不是 CUDA，则会回退到 Hub 上对应的内核。

FlashAttention-4 支持目前处于测试阶段。API 和行为可能会发生变化。

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

模式感知

内核会根据 PyTorch 的 model.training 状态，自动适配训练和推理模式。

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()

在 kernelize() 函数中，使用 mode 参数显式启用训练和推理模式。训练模式还支持额外的 torch.compile 模式。

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 用于自定义模型中使用的内核。

kernel_mapping 的键是模型注册的名称。它们可以指向某个层，例如 "RMSNorm"，或某个已注册的函数，例如 "rotary_pos_emb"。: 分隔符用于指定仓库中特定的内核条目，并将其映射到该层或函数。

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

为每种设备类型指定不同的内核实现。

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

内核元数据

添加一个元数据字典，以控制加载哪个内核构建版本。

选项	描述	默认
`version`	内核仓库的主要版本。`2` 会加载该仓库 `v2` 分支上的最新构建。	`1`
`revision`	要加载的确切标签、分支或提交，而不是版本。	---
`trust_remote_code`	允许使用位于受信任的 `kernels-community` 组织之外的仓库。加载内核时，将在您的机器上执行该仓库中的代码。	`False`

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": ("kernels-community/liger-kernels:LigerRMSNorm", {"version": 3}),
    }
)

内核仓库的旧分支可能没有针对您的 PyTorch 和 CUDA 版本的构建。如果内核加载失败，请先尝试使用更新的 version，然后再断定您的硬件不受支持。

继承的映射

KernelConfig 会继承 Transformers 针对 Hub 内核的默认内核映射，而 kernel_mapping 中的条目会覆盖对应层和函数的默认设置。设置 inherit_mapping=False 以仅使用 kernel_mapping 中的条目。所有未指定的部分都会回退到标准 PyTorch。这在基准测试特定内核或测试自定义实现而不应用其他默认设置时非常有用。

以下配置将 rotary_pos_emb 函数映射到 RoPE 内核，其余所有操作均保留在 PyTorch 上。该函数名必须由模型注册。

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

模块融合

通过向 KernelConfig 的键传入 (class_name, path_pattern) 对的元组，将相邻模块融合为单个内核。所有模式必须共享相同的父模块。* 匹配任意单个路径段。

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 要求内核仓库在 KernelName 类旁边提供一个配套的 KernelNameLayout 类。有关如何实现该类，请参阅编写内核指南。

本地内核

在 KernelConfig 中设置 use_local_kernel=True，即可从本地文件路径加载内核。这样会从本地文件系统路径加载，而不是从 Hub 仓库加载。

本地内核使用 /abs/path:layer_name，而不是 Hub 格式 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
)

全局禁用内核

将 USE_HUB_KERNELS 环境变量设置为禁用 Hub kernels，即可在不修改代码的情况下全局禁用 Hub kernels。

export USE_HUB_KERNELS=0  # or OFF or NO

故障排除

Kernel 集成依赖于硬件、驱动程序和软件包版本的协同工作。以下部分介绍了常见的故障。

安装问题

Import 错误表示 kernels 包缺失，或其版本不在 Transformers 支持的范围内。请通过 extra 重新安装以获取兼容的版本。

pip install -U "transformers[kernels]"

内核加载失败

如果特定的内核加载失败，请尝试以下操作。

请检查您的硬件是否符合内核要求。
请确认您的 CUDA/ROCm/Metal 驱动程序已更新至最新版本。
请参阅内核仓库文档以了解已知问题。

设备兼容性

并非所有内核都支持所有设备。如果某个内核在您的硬件上不可用，该库会回退到标准的 PyTorch 操作。请查阅内核仓库的文档以了解针对特定设备的支持情况。

资源

Kernels 仓库
使用 Hugging Face Kernel Hub 在 5 分钟内增强您的模型博客文章
在 kernels-community 组织中探索内核

Update on GitHub