397BパラメータのQwen 3.5 Mixture-of-Experts（MoE）モデルをIronwood TPUで提供するために、エンジニアはモジュラーJAX/Pallas最適化スタックを開発し、プリフィル負荷の高いワークロードで最大4.7倍の推論速度向上を達成しました。チームは、ハイブリッドデータ並列化とエキスパート並列化（DP+EP）トポロジーを展開することで、厳しいハードウェアシャーディングの制約を回避しました。これは、カスタム低レベル通信フュージョン（階層的なreduce-scatterなど）と組み合わせて、クロスデバイスのトークンルーティングを最適化しました。最後に、ハードウェアを意識したカスタムカーネル（Batched Ragged Page Attentionや完全にフューズされたGated DeltaNet（GDN）ブロックなど）を実行することで、HBM帯域幅とTensorCore MXUを効果的に活用し、システムスループットを理論上のルーフラインリミット近くまで押し上げました。