AMD MI300 用のカスタムカーネルの作成

Hugging Face と AMD は、AMD MI300X 用のオープンソース最適化カーネルのセットを開発し、VLLM を使用した FP8 での Llama 3.1 405B の提供性能を向上させました。3 つの特定のカスタムカーネル—融合残差接続 / RMS 正規化 / FP8 変換カーネル、融合 SwiGLU 活性化 / FP8 変換カーネル、そして Skinny GEMM カーネル—を実装することで、デコードレジーム中のレイテンシを大幅に削減しました（入力サイズ 1、出力サイズ 128 で測定）。

カスタムカーネルの実装と性能向上

融合 RMS 正規化カーネル

RMS 正規化カーネルは、残差接続、行単位の Root Mean Square (RMS) 正規化、および FP8 量子化を単一の操作に融合することで、デコーダブロックの開始部分を最適化します。

技術的最適化:

  • ベクトル化メモリアクセス: カーネルは 128 ビット幅のロードを使用して、1 命令あたり 8 つの FP16 要素を取得し、メモリアクセスが連続かつ合流するようにして warp の効率を最大化します。
  • 共有メモリ (SMEM) の活用: 繰り返しの VRAM（グローバルメモリ）アクセスを回避するために、隠れ状態 $x$ の修正バージョンが共有メモリに格納されます。Llama 405B では、次元 $d=16384$ が各コンピュートユニットで利用可能な 64KB の共有メモリに収まります。
  • ブロックレベルのリダクション: カーネルは各行に 1 つのスレッドブロックを割り当て、RMS 正規化に必要な総和のためにスレッドを共有メモリで同期させます。

結果: 「Vectorized + SMEM」実装は、標準の PyTorch と VLLM の既存実装の両方を上回り、さまざまなバッチサイズで大幅な速度向上を提供しました。

融合 SwiGLU カーネル

SwiGLU カーネルは、MLP ブロックの「Gate / Up」投影に対する活性化関数とその後の FP8 量子化を融合します。

技術的最適化:

  • パックド命令: カーネルは FP16 の加算と乗算に対して MI300X のパックド命令を利用し、命令あたりの作業量を増やします。
  • 高速数学近似: レイテンシを削減するため、カーネルは標準の exp 命令を、入力を $\log(2)$ でスケーリングしてより高速な exp2 命令に置き換え、精度の損失はほとんどありません。
  • パックド FP32 から FP8 への変換: MI300X は FP32 からの FP8 変換のみをサポートしているため、カーネルはパックド変換命令を活用して性能を向上させます。

結果: カスタム SwiGLU カーネルは、平均で PyTorch より 14 倍以上高速で、VLLM カーネルよりも 27% から 100% 高速です。

Skinny GEMM カーネル

標準ライブラリの一般的な行列乗算 (GEMM) カーネルは、行数が非常に少ない「スキニー」行列（低バッチサイズでのデコード時に典型的）に対しては、タイル機会が限られるため GPU の利用率が低く、非効率になることが多いです。

技術的最適化:

  • Split‑K アルゴリズム: カーネルは共有 K 軸に沿って GEMM をいくつかのサブ GEMM に分割し、同時に実行します。これにより、ワークロードを分散させてアクティブなコンピュートユニット (CU) の数が増え、各 CU が K 軸で費やす時間が短縮されます。
  • パディング除去のためのスパース性トリック: 行数が最小の密テンソルコア命令サイズ（例: 16）未満の場合の無駄なパディングを回避するため、カーネルは 4:2 構造化スパース命令を使用します。密な 8 行行列を 16 行のスパース行列にマッピングし、最小の密命令の深さの 2 倍である 16x16x64 スパース命令を使用できるようにします。
  • Warp の特殊化と非同期実行: 低い算術強度に対処するため、カーネルは Warp を「プロデューサー」（VRAM から共有メモリへのデータロード専用）と「コンシューマー」（計算専用）に分離します。共有メモリ内のキューがこれらの Warp を非同期に調整し、コンシューマーが遅い VRAM ロードを待つ間にアイドル状態になることを防ぎます。

結果: Skinny GEMM カーネルは、低行数（M = 1, 8, 16）において PyTorch より顕著な速度向上を示し、特に QKV と Gate/Up 投影で効果的ですが、バッチサイズが 32 に増加すると効果は減少します。

実装と入手方法

開発されたすべてのカーネルは hf-rocm-kernels GitHub リポジトリで入手可能で、ソースコード、Python バインディング、ベンチマークスクリプト、テストスイートが含まれます。これらのカーネルは単独で使用するか、VLLM に統合して使用できるよう設計されています。結果を再現するには、Hugging Face は開発時に使用した特定の ROCm 6.3.1 コンテナの使用を推奨します。

Sources