AMDとメタは、AMD MI350シリーズGPU向けPyTorchにFlyDSLサポートを追加したと発表した。モデル開発者は既存の演算呼び出しを維持しながら、行列積や正規化などをより高速な実装で実行できる。発表によると、少ないビット数で値を表す低精度モデルの推論は最大3.01倍高速化した。
1. 既存の演算経路に加わった選択肢
対象は、空きを設けず値を計算する一般的な密行列積、複数のまとまりを一緒に処理するグループ化行列積、低精度スケール行列積だ。値を正規化するRMSNorm、大きい値をK個選ぶTopK、入力間の関係を計算するマルチヘッドアテンションも含まれる。
コンパイル方式ではTorchInductorが条件に合う行列積の実装を直接測定し、最速の候補を保存する。既存実装と新実装を処理ごとに競わせて選ぶ方式だ。即時実行では、任意のランタイムがインストールされ有効になると、対応条件を満たす正規化・選択演算を、実行する計算手順を記述したコードである新カーネルに送る。
アテンションはAOTritonライブラリを経由し、事前のチューニングデータに応じて実装が選択される。
FlyDSLはカーネル作成者がデータ移動と計算の配置を直接制御できるようにし、ハードウェア別の最適化を支援する。
2. 演算別の測定結果は範囲が異なる
演算レベルの測定はInstinct MI355X(gfx950)で、コンパイルと初回の自動チューニングを除外して実施した。低精度の測定は、少ないビット数で値を表す量子化済みの入力から開始した。以下の幾何平均は、各ケースに同じ重みを与えた集計だ。
発表によると、BF16密行列積は左の行列が行優先、右が列優先となるNT配置の15個の形状で、既存の2実装のうち速い方に対して幾何平均1.10倍だった。比較対象の2実装はATenとTritonだ。
発表によると、MXFP8とMXFP4のスケール行列積はNT配置のそれぞれ17個の形状で、既存のデフォルト実装に対し、それぞれ1.58倍と1.68倍だった。
グループ化行列積は、発表によると24個のケースでTriton比1.20倍、ATen比1.95倍の幾何平均を記録した。
発表によると、RMSNormは22個のケースでATen比1.17~3.66倍だった。アラインされた次元ではおおむね1.2~1.5倍で、最高値は特定のFP16入力で得られた。
発表によると、Kが小さいTopKは10個のケースでATen比の幾何平均が4.82倍だった。これは同じ入力に同じ結果を返す手順である決定的アルゴリズムを無効にした条件で、有効にした条件では4.02倍だった。
3. アテンションとモデル推論で確認された効果
モデル全体の測定は、MI355X 1台でvLLMを使って実施した。入力256トークンと出力512トークンを固定し、同時リクエスト数を8~256個に変えた。比較する双方のサービング設定とモデル別のキャッシュは同じで、行列積の候補リストだけが異なった。
発表によると、BF16推論では既存の2バックエンドに新候補を加えると、Llama-3.1-8B-Instructが最大1.13倍、Qwen3-32Bが最大1.06倍だった。Llama-3.3-70B-Instructはおおむね測定ノイズの範囲にとどまった。
低精度8ビット推論は、発表によると以下のとおりだった。基準はhipBLASLtを使うATenで、モデルと同時リクエスト数の18個の組み合わせすべてで改善した。
MI355X 1台(TP=1)、MXFP8 A8W8、入力256トークン・出力512トークン、同時リクエスト8~256個、ATen(hipBLASLt)に対してFlyDSLを追加、組み合わせごとに1回測定
| モデル | 最小値 | 最大値 |
|---|---|---|
| Qwen3-32B | 1.40 | 2.03 |
| Llama-3.1-8B-Instruct | 1.04 | 1.92 |
| Llama-3.3-70B-Instruct | 1.38 | 2.04 |
同時リクエスト数によってモデル別の改善幅は異なった。発表によると、同じリクエスト条件のMXFP4推論も18個の組み合わせで1.25~3.01倍だった。この測定では最終出力層を量子化せず、組み合わせごとに1回ずつ実行した。
アテンションはAOTritonの旧バージョンに対するスループットの集計だ。発表によると、gfx950で0.13bに対し0.14.2bは順方向が0.92~2.34倍、逆方向が0.93~2.47倍だった。
4. 利用条件と性能の代償
最近のROCm PyTorch nightlyとgfx950ベースのMI350シリーズで利用を始める。コンパイル行列積用の任意のランタイムは0.3.x系で、以下は原文の例の候補設定部分だ。
import torch
import torch._inductor.config as config
config.max_autotune_gemm_backends = "ATEN,TRITON,FLYDSL"
config.flydsl_enable_autotuning = True設定後に行列積をtorch.compile(mode="max-autotune")でコンパイルすると、自動チューニングの候補になる。アテンションは追加のパッケージや設定なしで、AOTriton 0.14.2b以上と接続されたビルドで利用する。
繰り返し実行の高速化と引き換えに、初回呼び出しにはコンパイルとチューニングのコストがかかる。低精度入力が必要な配置に合わなければ、変換コピーも追加される。対応演算がモデルの実行時間に占める割合が大きいほど、最適化の価値が高まる。
アテンションのバージョン比較には、ほかのカーネルやチューニングデータの変更も含まれ、現在の新しいアテンションカーネルはFP32に対応しない。正規化の逆方向は既存実装を使い、上位K個の選択では同じ値のインデックス順序がバックエンドごとに異なる場合がある。
5. サポート拡大は進行中
現在の統合はgfx950のMI350シリーズを対象とする。次の段階として、MI450対応、低精度・アテンションカーネルの拡張、後続の計算を1つにまとめる後処理の融合、逆方向サポートの拡大などが進行中だと原文は説明する。
FlyDSLはAMD GPU向けPyTorchに最適化の選択肢を加え、効果は特に低精度行列積と一部の選択演算および推論条件で顕著だ。