llama.cppがAMD CDNAアクセラレーターでも必要な位置だけを選んでアテンションを計算する実行経路を開いた。AMD GPUでモデルを実行する開発者は、処理条件が合えばこの経路を利用でき、ソフトウェアは該当カーネルを優先的に選択する。
1. CUDA共有コードでAMDの実行を阻む条件を解除
今回の変更は10月29日にllama.cppへマージされた。CUDAと、AMD GPU向けの実行・開発環境であるHIPが共有するバックエンドで、HIPビルドのスパースアテンション経路を阻んでいた条件を取り除いたことが核心だ。既存のCUDA経路の選択条件を維持しつつ、適用範囲をAMD CDNAまで広げた。
従来はHIPでスパースアテンションのコードがコンパイル対象から外れ、必要な位置を集める関数も実行を中断していた。変更後はHIPでもこの準備処理とスパース計算を続けて実行する。カーネル選択規則も変更し、CDNAでスパース条件を満たす場合は、従来の処理サイズ別の選択規則よりスパース行列演算カーネルを優先する。
2. 必要な位置を先に集めるコストで計算量を削減
スパースアテンションは、計算対象の位置を示すマスクから参照可能なキー位置をリストに集め、選択したキーと値を処理する方式だ。マスクされた位置まで計算せず、必要な位置だけを選んで処理するため、アテンション演算を減らす利点と、位置リストを先に作る準備コストとの交換になる。
今回の実装は、GPUで一緒に動く実行単位のサイズに合わせて位置の選択結果を保存する。CDNAのGPUが並列計算を分担する64個の実行レーンを格納できるよう64ビットマスクを使い、選択した位置を数える処理もこれに合わせた。
3. 処理の形に合わせてスパースカーネルを選択
コードのコメントによると、CDNAでスパース対応をコンパイルする対象は、クエリ・キーの次元を示すDKQが576で、カーネルの列グループ設定であるncols2が16の変種だ。この変種を選択できるか先に確認し、その後で実際の処理条件を調べる。RDNAなど他のアーキテクチャへの対応は、作成者が今後の作業として残した。
実行時はマスクが存在し、計算結果に加えるバイアスと値の上限設定が0の処理をスパース経路の候補とする。CDNAのキー長は最低4096で、かつスパースリストに格納する最大キー数の2倍以上でなければならない。このようにソフトウェアが処理のサイズと形を調べ、実行経路を決める。
4. 合成テストのスループットと変動幅
pp512・tg64はテスト区分で、d16384・d32768はテスト深度の表記だ。以下の2列は原文に掲載された測定順に従う。スループットの横の±値も原文のままだ。
表：qwen4exp合成テストの2つの測定結果
条件：PR作成者の合成テスト。qwen4exp A3B Q4_K - Medium、103.68 GiB、176.94 B、ROCm、50回反復。2x mi100s、1x r9700、1x rx7900xtxに分散ロード。スパース経路を有効化する別のパッチを含む。
単位：t/s（1秒当たりの処理トークン数）
| 測定項目 | 1回目の測定 | 2回目の測定 |
|---|---|---|
| pp512 @ d16384 | 760.52 ± 55.29 | 782.11 ± 11.73 |
| tg64 @ d16384 | 35.81 ± 3.04 | 36.43 ± 2.43 |
| pp512 @ d32768 | 633.18 ± 99.63 | 736.58 ± 62.15 |
| tg64 @ d32768 | 30.78 ± 5.79 | 35.20 ± 2.19 |
1回目の測定では深度がd16384からd32768に増えると、pp512のスループットは760.52から633.18 t/sに、tg64は35.81から30.78 t/sに下がる。pp512の±値も55.29から99.63に増え、この測定では深いテストのスループットが低く、変動幅も大きい。
作成者は、このテストで性能改善が明確だと評価した。また、モデルの約半分が今回の変更の影響を受けないアーキテクチャに配置された点から、測定結果は改善効果を小さく示していると主張した。性能表には適用前後を識別する表示がないため、2列だけで改善率を計算するのは難しい。
5. 旧ROCmビルドにも配慮
ROCm 7.0未満では、レーンごとの選択結果を集める__ballot_syncの代わりに従来の__ballotを使う互換定義を追加した。作成者は、ROCm 7.0より前には__ballot_syncの組み込み関数がHIP_ENABLE_WARP_SYNC_BUILTINSの後ろに隠されていたと指摘した。原文に追加されたコードは次のとおりだ。
#if HIP_VERSION < 70000000
#define __ballot_sync(mask, var) __ballot(var)
#endif //HIP_VERSION < 70000000作成者は、この場合はこのカーネルでは__ballotをそのまま使えばよいとした。ただし、この互換マクロ自体はカーネルごとの限定ではなく、共有HIPヘッダーggml/src/ggml-cuda/vendors/hip.hにHIP_VERSION < 70000000の条件で定義されている。同時に、別のファイルでも__ballot_syncが使われており、その箇所ではwave64デバイスでマスクが実際に使われるため、ROCm 6.3-6.4のwave64ではその使用が正しくないと警告した。その使用については後続で対応すると述べ、その後続コミットはこのPRには含まれていない。したがって上述の互換対応は、残っているこの正しさの問題を解決していない。
今回の変更の核心は、CDNAで必要な位置だけを計算する経路を開き、処理条件に合わせてそのカーネルを選択するよう接続した点にある。