[論文レビュー] Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU

Muhammad Osama, Duane Merrill+3 more|arXiv (Cornell University)|Jan 9, 2023

Parallel Computing and Optimization Techniques被引用数 5

ひとこと要約

Stream-Kは、GPU上の密行列行列乗算（GEMM）に対して、出力行列をタイリングするのではなく、ストリーミングマルチプロセッサ（SM）にマクロープロセスの反復処理を均等に分散させることで、作業中心の並列分解を導入する。これにより、多様な問題形状においてほぼ100％のプロセッサ使用率を達成し、cuBLASおよびCUTLASSに対して最大14倍のピークスループット向上と平均6.7倍の性能向上を実現。1精度あたり1つのカーネルを採用することで、複雑なカーネルエンsemblesの必要性を排除した。

ABSTRACT

We introduce Stream-K, a work-centric parallelization of matrix multiplication (GEMM) and related computations in dense linear algebra. Whereas contemporary decompositions are primarily tile-based, our method operates by partitioning an even share of the aggregate inner loop iterations among physical processing elements. This provides a near-perfect utilization of computing resources, regardless of how efficiently the output tiling for any given problem quantizes across the underlying processing elements. On GPU processors, our Stream-K parallelization of GEMM produces a peak speedup of up to 14$ imes$ and 6.7$ imes$, and an average performance response that is both higher and more consistent across 32,824 GEMM problem geometries than state-of-the-art math libraries such as CUTLASS and cuBLAS. Furthermore, we achieve this performance from a single tile size configuration per floating-point precision, whereas today's math libraries employ complex kernel-selection heuristics to select from a large ensemble of kernel variants.

研究の動機と目的

• 広いGPUにおけるタイルベースGEMMワークロードの量子化非効率性が引き起こすパフォーマンス劣化を是正すること。
• 数学ライブラリにおける複雑なカーネル選択ヒューリスティクスの必要性をなくし、多様な問題形状において一貫した高性能実行を可能にすること。
• MACループレベルでのワークロード量子化を再定義することで、問題サイズや形状に関係なくほぼ100％のプロセッサ使用率を達成すること。
• 1精度あたり1つの高パフォーマンスでポータブルなカーネルに置き換えることで、ライブラリの肥大化と保守の複雑さを軽減すること。

提案手法

• Stream-Kは、GPUのストリーミングマルチプロセッサ（SM）に、合計の乗算加算（MAC）ループ反復回数を均等に割り当てることで、負荷のバランスを保証する。
• ワークロード量子化の原子的単位としてMACループ反復を用い、これは出力タイル全体と比較して一定かつ最小限のコストである。
• 動的にp個の分割シーム（pはSMの数）を生成し、問題サイズではなくプロセッサの幅に応じて負荷分散をスケーリングする。
• 1浮動小数点精度あたり1つの固定ブロッキング要因を採用し、複数のカーネルバージョンや複雑な選択ヒューリスティクスの必要性を回避する。
• 通信および同期のオーバーヘッドを最小限に抑え、これは問題サイズではなくSMの数に比例して増加する。
• CUTLASS 2.11に統合され、FP64、FP16→FP32、その他の精度フォーマットをサポートし、一貫したパフォーマンスを発揮する。

(a) Data parallel decomposition with grid size $g$ =9 CTAs, large $128\times 128\times 128$ CTA work volumes, and 75% processor utilization ceiling

実験結果

リサーチクエスチョン

• RQ1 MACループレベルで量子化する作業中心の分解は、現代のGPUにおけるタイルベースGEMMよりも高いかつ一貫性のあるパフォーマンスを達成できるか？
• RQ2 1精度あたり1つのカーネルが、パフォーマンスを損なわせることなく、cuBLASのような数学ライブラリにおける大規模なカーネルエンsemblesを置き換えることができるか？
• RQ3 多様な問題幾何形状において、Stream-Kのワークロードバランスはタイルベース手法と比較して、量子化効率の面でどの程度優れているか？
• RQ4 一貫した作業中心の分解に置き換えることで、カーネル選択ヒューリスティクスを排除した場合のパフォーマンスへの影響はいかほどか？

主な発見

• Stream-Kは、32,824個のGEMM問題幾何形状において、cuBLASに対して最大14倍のピークスループット向上と平均6.7倍のパフォーマンス向上を達成した。
• FP64では、cuBLASよりも6％高いスループットを達成し、FP16→FP32では13％高いスループットを実現した。
• Stream-Kのパフォーマンススプレッドは、cuBLASおよび理想化されたCUTLASSオラクルの両方を下回っており、優れた一貫性と使用率を示している。
• 計算制限領域（FP64で1バイトあたり150以上の演算、FP16→FP32で1バイトあたり400以上の演算）では、Stream-KはcuBLASのエンセムを一貫して上回った。
• Stream-Kは、cuBLASと比較して配布サイズを最大20倍まで削減し、1精度あたり1つのカーネルでピークパフォーマンスを達成した。これにより、数百の特化カーネルの必要性が排除された。
• 問題の形状に関係なく、Stream-Kはほぼ100％のプロセッサ使用率を達成し、タイルベース手法の量子化非効率性に起因する75–90％の理論的上限を解消した。

(b) Data parallel decomposition with grid size $g$ =18 CTAs, smaller $128\times 64\times 128$ CTA work volumes, and 90% processor utilization ceiling