Kernels Hub 集成与使用

kernels 库允许直接从 Hub 加载经过优化的计算内核。
你可以在专用组织中找到 kernels，也可以在 Hub 中搜索 kernel 标签。

内核是经过优化的代码片段，有助于模型的构建、训练和推理。在这里，我们将重点介绍它们与 TRL 的集成，但请参阅上述资源以了解更多信息。

安装

要在 TRL 中使用 kernels，您需要在 Python 环境中安装该库：

pip install kernels

在 TRL 中使用来自 Hub 的 Kernels

Kernels 可以直接替换注意力实现，无需手动编译 Flash Attention 等注意力后端，只需从 Hub 拉取相应的注意力 kernel 即可提升训练速度。

加载模型时，您可以指定一个内核：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2"  # other options: kernels-community/vllm-flash-attn3, kernels-community/paged-attention
)

或在运行 TRL 训练脚本时：

python sft.py ... --attn_implementation kernels-community/flash-attn2

或使用 TRL CLI：

trl sft ... --attn_implementation kernels-community/flash-attn2

现在，您可以利用 Hub 上针对您的硬件配置预先优化的内核，使用更快的注意力后端，从而加速开发和训练。

选择和固定内核版本

Hub 上的 Kernel 仓库以分支形式进行版本管理（v1、v2、v3……），并且 Transformers 会为每个仓库选择一个默认版本。该默认版本可能会随着 Transformers 的每次发布而改变，因此相同的 attn_implementation 值并不总是解析到相同的构建版本。

要控制加载哪个构建版本，请在仓库 ID 后附加一个修订版本，可以是版本分支或提交 SHA：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2@v2",  # or a commit SHA to pin an exact build
)

固定版本（Pinning）有助于使训练过程可复现，并在出现新版本回归问题时，保持在已知可用的构建版本上。请注意，每个版本仅针对特定范围的 Torch 和 CUDA 版本提供构建，因此固定的版本可能没有适用于您环境的变体。

attn_implementation="flash_attention_2" 在未安装 flash-attn 包时，会回退到 kernels-community/flash-attn2 Hub 内核，因此您可能正在使用 Hub 内核，而并未显式请求使用它。

目前，针对 CUDA 12.8 的 v3 构建在使用分组查询注意力的模型中，反向传播阶段会失败。如果你的环境使用该 CUDA 版本，请在 huggingface/kernels-community#1085 关闭之前，固定使用 @v2。

比较注意力实现

我们使用 TRL 和 SFT 评估了 transformers 中可用的各种注意力实现以及不同的内核后端。
实验在单张 H100 GPU 上运行，使用 CUDA 12.9，并采用 Qwen3-8B，batch size 为 8，gradient accumulation 为 1，精度为 bfloat16。
请注意，此处展示的结果仅针对该特定配置，在不同的训练配置下可能会有所不同。

下图展示了不同注意力实现和内核后端的延迟（每个训练步骤的时间）和峰值分配内存。
基于内核的实现与自定义安装的注意力性能相当，并且增加模型的max_length可进一步提升性能。所有实现的内存消耗相似，没有显著差异。如下一节所述，我们获得了相同的性能，但摩擦更小。

Flash Attention 与 Hub 内核

从源代码构建 Flash Attention 可能非常耗时，具体所需时间从几分钟到数小时不等，这取决于您的硬件、CUDA/PyTorch 配置以及是否有预编译的 wheel 可用。

相比之下，Hugging Face Kernels 提供了更快、更可靠的工作流。开发者无需担心复杂的配置——一切都会自动处理。在我们的基准测试中，内核在大约 2.5 秒 内即可就绪，无需编译。这让你几乎可以立即开始训练，从而显著加速开发过程。只需指定所需的版本，kernels 便会处理其余部分。