MiMo-V2.5推論最適化：ハイブリッドSWAを本番で実用化する

XiaomiのMiMoチームは、MiMo-V2.5およびV2.5-Proモデルの背後にある推論最適化について詳細な解説を公開した。見出しは次の通り：ハイブリッドスライディングウィンドウアテンション（SWA）はKVCacheのストレージと計算を7倍削減すると約束するが、論文から本番環境へのギャップを埋めるには、キャッシュとスケジューリングスタックの根本的なリファクタリングが必要だった。

V2.5シリーズは70層のアーキテクチャを使用している——10層のフルアテンション層と、ウィンドウサイズ128の60層のSWA層。理論上、これによりアテンション計算とKVCacheがフルアテンションモデルの約1/7に削減される。しかし実際には、チームが指摘するように、SGLang v0.5.5のようなフレームワークのナイーブな実装ではSWA層にフルKVCacheを保存しており、その利点を無効にしていた。

デュアルプールKVCache設計

核心的な修正は、KVCacheを2つの独立したプールに分割することだった：1つはフルアテンション層用（O(N)ストレージ）、もう1つはSWA層用（O(W)ストレージ）。物理層では、SWAプールはウィンドウのみのサイズに設定され、独立した削除をサポートする。論理層では、統一されたシーケンスビューがプレフィックスツリーとスケジューラに公開され、透過的な階層型ストレージのためにFull→SWAマッピングが維持される。これだけで7倍の容量効率が回復する。

SWA対応プレフィックスキャッシュツリー

従来のRadixAttentionは、等しいトークンシーケンスが再利用可能なKVを生成すると仮定する。SWAではそれが崩れる——プレフィックスツリーノードは論理的にはトークンシーケンスを表すかもしれないが、そのSWA KVはウィンドウを超えて削除されている可能性がある。チームはマッチングルールを「ウィンドウセーフ長」に改訂した：マッチ長は、末尾WトークンがSWAプールに有効なスロットを持つようにクリップされる。削除はリクエストライフサイクル（チャンク完了、終了、Nデコードトークンごと）に結び付けられ、SWAプールの使用量をシーケンス長に応じて増加させるのではなく、ウィンドウサイズで一定に保つ。ノードは現在、フルアテンションセグメントインデックスとSWAセグメントマッピングの二重インデックスを持ち、ウィンドウ外のSWAセグメントを独立して削除しつつ、フルアテンションセグメントを再利用のために保持できる。

階層型キャッシングと分散一貫性

3層のHiCacheシステム（デバイス、ホスト、ストレージバックエンド）は、階層ごとに有効なSWAインデックスを維持するようにリファクタリングされた。チームは、フルアテンションキャッシュヒットとSWAキャッシュミスが深刻なマッチ長の切り捨てを引き起こし、再計算を強制することを発見した。彼らは、デプロイメント間の分散一貫性、共有プレフィックス長、非同期データ移動を最適化して、両者を同期させた——新しいマッチングルールの下でヒット率を維持するために重要である。

MoEとマルチモーダルのボトルネック

SWA以外にも、記事はMoEスケジューリングと負荷分散に触れ、マルチモーダルエンコーダが大画像や長動画入力のスループットボトルネックのままであると指摘している。チームはこれらに対処するためにPrefill/Decode実行パイプラインとスケジューリング戦略を最適化したが、この記事は主にKVCacheの作業に焦点を当てている。

結論：ハイブリッドSWAの理論上の効率性は現実のものだが、それはキャッシュシステムをゼロから再構築した後にのみ実現する。MiMoチームのアプローチ——デュアルプール、ウィンドウセーフプレフィックスツリー、階層型一貫性——は現在、主流の推論フレームワークに採用されつつあり、この作業がどれほど必要だったかを物語っている。