编程 CUDA kernel 也能按版本装：Hugging Face kernels 的 v1 分支、lock 与 wheel 变体

2026-10-09 00:04:39

仓库：huggingface/kernels · 浏览 kernel：huggingface.co/kernels · 文档：huggingface.co/docs/kernels

背景：kernel 难在“能被依赖”

自定义 CUDA kernel 能拿到性能，但塞进真实项目是另一回事：构建慢、依赖地狱、和 PyTorch 版本 / CUDA 版本 / C++ ABI 绑死。Kernel Hub 想把 kernel 变成可以按版本拉取的产物——Python 库和应用直接从 Hub 加载 compute kernel。

因为要动态加载，Hub kernel 和传统 Python kernel 包不同，必须满足三点：Portable，能从 PYTHONPATH 之外的路径加载，产物落在缓存目录而不是 site-packages；Unique，同一个 Python 进程里可以同时加载同一个 kernel 的多个版本；Compatible，支持所有近期 Python 版本与不同 PyTorch 构建配置（各种 CUDA 版本、C++ ABI），并且兼容较老的 C 库版本。

消费端：get_kernel

安装：pip install kernels（要求 torch>=2.5 且有 CUDA）。

import torch
from kernels import get_kernel

# 从 Hub 下载优化好的 kernel
activation = get_kernel("kernels-community/activation", version=1)

x = torch.randn((10, 10), dtype=torch.float16, device="cuda")
out = torch.empty_like(x)
activation.relu(out, x)

这是拉取 kernels-community/activation 的 version 1。kernel 按主版本号版本化：version=1 拿到 v1 分支上最新的构建。version 分支内绝不能破坏 API、也不能删除旧 PyTorch 版本的构建，已有代码才能持续工作。Hub kernel 必须用 version 或显式 revision 之一加载。

version 0 不受 API 兼容承诺保护，留给 alpha/beta 质量、API 还会快速变化的 kernel。

get_kernel(repo_id, revision=None, version=None, backend=None, user_agent=None, trust_remote_code=False, check_arch=True)

revision 与 version 不能同时用，必须指定其一。backend 只能是 cpu 或 Torch 编译时用的后端，不传自动探测。check_arch 默认 True，检查 kernel 构建是否支持当前设备架构（如 CUDA compute capability）；有些 kernel 通过 Triton fallback 支持比声明更多的架构，这种情况可以 check_arch=False 跳过。trust_remote_code 默认 False，只允许可信组织的 kernel。

get_loaded_kernels() 返回进程内每个已加载 kernel 的快照 LoadedKernel（导入的 module、package_name、repo_infos）。repo_infos 只有 get_kernel() 加载的才有；用本地路径 get_local_kernel() 或 lockfile（get_locked_kernel()、load_kernel()）加载的为 None。

版本语义与 lock

版本就是 Git tag，形如 vx.y.z：

img2gray_lib = get_kernel("drbh/img2gray", revision="v1.1.2")

用 Git shorthash 能降低破坏风险，但难读、也不支持在版本区间内平滑升级，所以推荐 semver。上界比精确版本更有用：semver 里 1.y.z 的公共 API 不得有不兼容改动，于是可以要求「至少 1.1.2、小于 2.0.0」：

img2gray_lib = get_kernel("drbh/img2gray", version=">=1.1.2,<2")

版本说明符和 Python 依赖一样。va.b.c 形式的版本 tag 在这里也用得上——kernels 通过仓库的版本 tag 查询有哪些版本可用。在 pyproject.toml 里指定 kernel 后，kernels lock 会锁定到具体版本，生成 kernels.lock。

锁定的最后一步是把每个 get_kernel 换成 get_locked_kernel。配套的 load_kernel 只加载本地已有 kernel，绝不尝试从 Hub 下载，本地没有就抛异常。要让 kernel 在本地可用：

kernels download .

它下载 kernels.lock 里指定的 kernel。放进 Docker 构建阶段，就把这一步固化了。

打成 wheel

多数情况推荐用 kernels 包从 Hub 下载，但有些项目要求以 wheel 形式部署。一条命令可以把任意 Hub kernel 转成一组 wheel：

$ kernels to-wheel drbh/img2grey 1.1.2
img2grey-1.1.2+torch27cu128cxx11-cp39-abi3-manylinux_2_28_x86_64.whl
img2grey-1.1.2+torch26cu124cxx11-cp39-abi3-manylinux_2_28_x86_64.whl
img2grey-1.1.2+torch26cu126cxx11-cp39-abi3-manylinux_2_28_x86_64.whl
...
img2grey-1.1.2+torch27cu118cxx11-cp39-abi3-manylinux_2_28_aarch64.whl

每个 wheel 都像普通 Python wheel 一样可以 import img2grey。本地版本标识把变体编进去了：torch27 是 PyTorch 2.7，cu128 是 CUDA 12.8，cxx11 是 C++ ABI，cp39-abi3 是 Python 标签（abi3 表示稳定 ABI），manylinux_2_28_x86_64 / aarch64 是目标架构。一个 kernel 要覆盖多少构建变体，这一屏文件名就是答案。

作者侧：注册成原生算子

示例是把图像 RGB 转灰度。GPU 代码放在 csrc/img2gray.cu，用线程的二维网格定义 kernel——处理图像的自然方式。

关键一步不是绑到 Python，而是用 PyTorch 现代 C++ API 把函数注册为原生 PyTorch 算子，让它在 torch.ops 命名空间下可见：

与 torch.compile 兼容：注册后 torch.compile 能“看到”这个算子，可以把它融合进更大的计算图，减少开销。
按硬件分发：同一算子可以有不同后端。再加一个 TORCH_LIBRARY_IMPL(img2gray, CPU, ...) 指向 C++ CPU 函数，PyTorch 分发器会根据输入张量所在设备自动调用正确实现（CUDA 或 CPU）。

作者侧可以用 kernel-builder 构建 kernel；Hugging Face 在 kernels-community 组织里维护了一批 kernel。