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ブロックワイズ行列乗算の実装です。

技術仕様

Precision: 入力にe4m3fnuz FP8浮動小数点フォーマットを使用し、スループットを向上させ、メモリ帯域幅を削減します。
Accuracy: FP8の限られたダイナミックレンジにもかかわらず、数値の安定性を維持するためにブロックごとのスケーリングファクター(a_scaleとb_scale)を使用します。
Inputs/Outputs:
- a: K × M in e4m3fnuz
- b: K × N in e4m3fnuz
- a_scale: (K // 128) × M in fp32
- b_scale: (K // 128) × (N // 128) in fp32
- c: M × N in bf16
Recognition: このカーネルは、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 Integration: カーネルはTORCH_LIBRARY_EXPANDを使用してネイティブPyTorch演算子として登録されます。これにより、torch.opsを介してカーネルにアクセスでき、PyTorchフレームワークの第一級部分として振る舞うことができます。
Python Wrapper: __init__.pyファイルは、ユーザーフレンドリーなインターフェースを提供し、基礎となる演算子を呼び出す前にテンソルの作成とシェイプ検証を行います。

再現性とデプロイ

Nixベースのビルドプロセス

ビルドはNixを介して処理され、異なるマシン間で環境が同一であることを確保します。

ロック: nix flake updateは、kernel-builderとその依存関係を固定するためにflake.lockファイルを生成します。
キャッシュ: Hugging Faceキャッシュ（cachix経由）は、PyTorchバージョンの高コストな再ビルドを避けるために使用されます。
マルチバージョンサポート: nix build . -Lコマンドは、サポートされているすべてのPyTorchとROCmのバージョンに対してカーネルのビルドを自動化できます。

Hugging Face Hubを介した配布

ビルドが完了したら、カーネルはkernels uploadコマンドまたはバイナリファイル（.soファイル）のためのGit Xetを使用してHugging Face Hubにアップロードされます。これにより従来のインストールの必要がなくなり、ユーザーはget_kernelを使用してHubから直接カーネルをロードできます：

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)