輕鬆使用 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）、目標架構（例如 gfx942 對應 MI300 系列）以及來源檔案。
PyTorch 整合：核心透過 TORCH_LIBRARY_EXPAND 註冊為原生 PyTorch 運算子。這使得核心能透過 torch.ops 存取，並作為 PyTorch 框架的一等公民運作。
Python 包裝：__init__.py 檔案提供使用者友善的介面，在呼叫底層運算子之前處理張量建立與形狀驗證。

可重現性與部署

基於 Nix 的建置流程

建置透過 Nix 進行，以確保在不同機器上的環境保持一致。

鎖定：執行 nix flake update 會產生 flake.lock 檔案，以釘住 kernel-builder 及其依賴項。
快取：透過 cachix 使用 Hugging Face 快取，以避免重新建置 PyTorch 版本所帶來的高額成本。
多版本支援：指令 nix build . -L 可自動為所有支援的 PyTorch 與 ROCm 版本建置核心。

透過 Hugging Face Hub 進行分發

建置完成後，核心會透過 kernels upload 指令上傳至 Hugging Face Hub，或透過 Git Xet 上傳二進位檔案（.so 檔）。這樣省去了傳統安裝的需求；使用者可直接透過 get_kernel 從 Hub 載入核心：

import torch
from kernels import get_kernel

# Load the kernel from the Hub
gemm = get_kernel("kernels-community/gemm\n
# Execute the kernel
result = gemm.gemm(A_fp8, B_fp8, A_scale, B_scale, C)