RedFuser: AIアクセラレータ上の階層的リダクションの自動演算子融合フレームワーク

本文: arXiv:2603.10026v1 アナウンス種別: cross 要旨：演算子フュージョンは、AIモデルのデプロイ時の主要なパフォーマンス最適化技術であり、実行効率を大幅に向上させ、現代のAIコンパイラで広く採用されています。しかし、ループ間データ依存性を伴う複数ループからなる階層的リダクションを含むケース、例えばアテンション機構における数値的に安定なsoftmaxの後にGEMMを実行するケースでは、既存のコンパイラには自動化された融合とカーネル生成の能力が十分には備わっていません。いくつかの研究は手作業による融合戦略で特定のケースに対処していますが、それらの解法は一般性に欠け、他の類似構造へ拡張するのは難しいです。このような計算パターンが深層学習モデルで広く見られる現状を踏まえると、一般的かつ自動化された融合最適化を実現する潜在能力にはまだ未開拓の部分が残っています。 本論文では、階層的リダクションを分析し、それらを単一ループへ融合し増分計算形を導入する正式な理論的方法を提示します。これに基づき、Reduction Fuser (RedFuser) という、サポートされる階層的リダクションパターンを自動的に識別し、最適化された融合カーネルを生成するフレームワークを設計します。実験の結果、RedFuserは多様なワークロードを効果的に融合し、最先端のAIコンパイラに対して2×から5×のスピードアップを達成し、高度に最適化された手書きカーネルの性能に匹敵します。コードは https://github.com/alibaba/redfuser で利用可能です。