轻松使用 Hugging Face 构建和共享 ROCm 内核

Hugging Face 已发布指南和工具，以简化 ROCm 兼容内核的创建和分发。通过使用 kernels 库和 kernel-builder，开发者可以为 AMD 硬件构建高性能 GPU 操作，并通过 Hugging Face Hub 共享它们，确保可重复性和与 PyTorch 的无缝集成。

RadeonFlow GEMM 内核示例

为了演示构建过程，Hugging Face 使用了 RadeonFlow GEMM 内核，这是一个针对 AMD Instinct MI300X GPU 优化的高性能 FP8 分块矩阵乘法实现。

技术规格

精度：使用 e4m3fnuz FP8 浮点格式作为输入，以提高吞吐量并降低内存带宽。
准确性：采用每块缩放因子（a_scale 和 b_scale）来在 FP8 动态范围有限的情况下保持数值稳定性。
输入/输出：
- a: K × M 在 e4m3fnuz 中
- b: K × N 在 e4m3fnuz 中
- a_scale: (K // 128) × M 在 fp32 中
- b_scale: (K // 128) × (N // 128) 在 fp32 中
- c: M × N 在 bf16 中
认可：此内核在 2025 年 6 月的 AMD Developer Challenge 2025 中获得了大奖。

使用 kernel-builder 构建 ROCm 内核

开发自定义内核通常涉及复杂的构建标志和 ABI 问题。Hugging Face kernels 库通过结构化的项目组织和使用 Nix 来实现可重复性，从而抽象出这种复杂性。

项目结构

项目被组织到特定目录中，以帮助构建器识别文件类型：

build.toml：项目清单，用于编排构建过程。
gemm/：包含原始 HIP 源代码（.hip 用于实现，.h 用于头文件）。
flake.nix：通过锁定依赖项来确保可重复的构建环境。
torch-ext/：包含将内核作为 PyTorch 操作公开所需的 C++ 绑定和 Python 包装器。

配置和注册

build.toml：定义后端（例如 rocm）、目标架构（例如 MI300 系列的 gfx942）和源文件。
PyTorch 集成：使用 TORCH_LIBRARY_EXPAND 将内核注册为原生 PyTorch 操作。这使得内核可以通过 torch.ops 访问，并作为 PyTorch 框架的一等公民行为。
Python 包装器：__init__.py 文件提供用户友好的界面，在调用底层操作之前处理张量创建和形状验证。

可重复性和部署

基于 Nix 的构建过程

构建过程通过 Nix 处理，以确保在不同机器上的环境相同。

锁定：nix flake update 生成一个 flake.lock 文件来固定 kernel-builder 及其依赖项。
缓存：Hugging Face 缓存（通过 cachix）用于避免重新构建 PyTorch 版本的高昂成本。
多版本支持：nix build . -L 命令可以自动化构建内核，以支持所有 PyTorch 和 ROCm 的版本。

通过 Hugging Face Hub 分发

构建完成后，可以使用 kernels upload 命令或通过 Git Xet 将内核上传到 Hugging Face Hub（用于二进制文件 .so 文件）。这消除了传统安装的需求；用户可以直接从 Hub 使用 get_kernel 加载内核：

import torch
from kernels import get_kernel

# Load the kernel from the Hub
gemm = get_kernel("kernels-community/gemm

# Execute the kernel
result = gemm.gemm(A_fp8, B_fp8, A_scale, B_scale, C)