概要
GoogleがQwen 3.5-397B MoEモデルを最新のTPU v7x (Ironwood)上で最適化し、JAX/Pallasカーネルや新並列化手法により推論スループットを最大4.7倍向上させた技術報告。
詳細内容
Googleのパフォーマンスエンジニアリングチームが、巨大なMoEモデルであるQwen 3.5-397BをTPU v7xで効率的に動作させるための手法を公開しました。Qwen 3.5は75%のGated DeltaNet（線形注意）と25%のGQA（Grouped Query Attention）を組み合わせた特殊なハイブリッド構造を持ちます。チームは、GQAのKVヘッド数制限を克服するためにAttention DP（データ並列）とEP（エキスパート並列）を組み合わせたハイブリッド・シャーディングを導入しました。さらに、JAX/Pallasを用いて、SparseCoreを活用したトークンルーティングのオフロード、Conv1Dと再帰更新を統合したGDN融合カーネル、ページサイズを最適化したRagged Page Attentionなどを実装。その結果、理論的なルーフライン限界の約80%に達する効率を実現し、デコード性能を3.1倍、プリフィル性能を4.7倍向上させました。これらの成果はvLLMなどのオープンソーススタックにも統合され、大規模エンタープライズ推論の効率化に寄与します。