清华大学と腾讯混元がMLSys2026 MoE推論チャレンジで優勝 NPU推論のスピードを4.1倍に向上

清华大学ストレージラボとテンセント・ホンユアンAI Infraチームは、国際的な機械学習システムのトップ会議MLSys2026で開催されたMoEモデル推論最適化チャレンジにおいて、世界チャンピオンを獲得しました。

トレーニングに数兆パラメータを持つ混合エキスパート（MoE）アーキテクチャが非均質なチップ（NPU）上で推論性能の限界に直面している中、共同チームは公式指定モデルおよびNPUハードウェアに対して全フロント最適化ソリューションを開発しました。エキスパートごとにタスクを分割するE-Shard戦略、PSUM三次元テンソルバッチ読み込み、出力を複数Bankに分散して並列処理するGEMVパス、そしてスカラエンジンを利用して初期データ転送遅延を低下させるなどの組み合わせにより、データ転送不足やアクティベーションの再転送などの下位演算子の問題を成功裏に解決しました。

また、アテンションモジュールに関して、チップ上でのデータレイアウトを再構築し、Transformerの重要な演算子を統合することで、ビット単位の高精度の整合性を実現しました。

図3：MoE最適化構造の概要図。E-Shardエキスパート分割、連続DMA、PSUM/GEMV並行処理、コールドスタートパイプライン、プリフェッチ制御を含む。

この競争では、チームは「Knight」というエージェントベースの推論演算子オプティマイザを開発・共同しました。これは、提案、コードの実装、反省と改善の自動化されたループを通じて、最適化の探索空間を大幅に拡大します。最終的に、このソリューションによりモデルのエンドツーエンド推論時間を14.91秒から3.56秒に短縮し、性能は4.1倍に向上しました。1ステップのデコード遅延は12.63ミリ秒から5.45ミリ秒に減少し、重みロード中のDMAエンジン利用率は約80%まで上昇しました。

スタンフォードやMITなど国際的なトップ大学を倒して優勝したことは、中国チームが大規模モデルの下層システムへの適合性と演算子最適化における深い蓄積を示しており、今後の超ノード演算力プラットフォームにおける数兆パラメータを持つMoEモデルの配置に非常に価値のある工学的モデルを提供しています。