AI・機械学習
MiMo v2.5 の推論最適化:ハイブリッド SWA 効率を極限まで追求
Inference Optimization for MiMo v2.5: Pushing Hybrid SWA Efficiency to the Limit (mimo.xiaomi.com)
要約
この記事は、MiMo-V2.5 モデルファミリーの推論効率を向上させるための最適化について論じています。特に、ハイブリッド Sliding Window Attention (Hybrid SWA) と sparse Mixture-of-Experts (MoE) が、計算コストとメモリ使用量を削減し、長文コンテキストやマルチモーダルシナリオでの効率を大幅に改善する仕組みを解説しています。KVCache 管理、階層型キャッシュ、スケジューリング戦略といったエンジニアリングの実践を通じて、理論上の効率を実際のプロダクション環境で実現することに焦点を当てています。
全文翻訳
English简体中文ProductMiMo CodeResearchPaperBlogJoin UsEnglish简体中文 2026年5月30日
MiMo-V2.5 シリーズのフルパイプライン推論最適化:ハイブリッド SWA 効率を極限まで追求
MiMo-V2.5 および MiMo-V2.5-Pro を含む V2.5 モデルファミリーは、いくつかのアーキテクチャ設計上の選択を組み合わせています。ハイブリッド Sliding Window Attention (Hybrid SWA) は KVCache ストレージをフルアテンションの約 1/7 に圧縮し、sparse MoE アクティベーションはモデル容量を維持しながらトークンごとの計算量を削減し、マルチモーダルエンコーダーはビジョン、オーディオ、ビデオ間のクロスモーダル理解を可能にします。これらの機能が組み合わさることで、MiMo-V2.5 シリーズは、長文コンテキストおよびマルチモーダルシナリオにおいて、大幅なパフォーマンスと効率の可能性を秘めています。
当初から、私たちの目標は明確でした。強力でありながら長文コンテキスト推論に効率的なモデルをトレーニングすることです。これら 2 つの目標は本質的に相反します。強力な推論には長距離の依存関係をモデル化する必要がありますが、これには通常、より大規模なアテンション計算と高い KVCache オーバーヘッドが必要になります。従来のフルアテンションアーキテクチャでは、アテンション計算と KVCache ストレージの両方がコンテキスト長とともに急速に増加するため、長文コンテキストのトレーニングと推論は法外に高価になります。ハイブリッド SWA は、レイヤー間でローカル Sliding Window Attention (SWA) とグローバルフルアテンションをインターリーブすることで機能します。ほとんどのレイヤーはローカルウィンドウ内でのみアテンションを計算し、少数の主要レイヤーがグローバルビューを維持します。理論上、この構造は長距離の依存関係をモデル化する能力を維持しながら、アテンションの複雑さをほぼ線形にまで低減します。
しかし、理論上のアーキテクチャ上の利点は、自動的にプロダクション効率に変換されるわけではありません。ハイブリッド SWA は、KVCache ヒット率の管理、プレフィックスマッチング、フルアテンションレイヤーと SWA レイヤー間のデュアルセマンティック整合性の維持において、新たな複雑さをもたらします。実際のエンジニアリングシステムは、マルチレベルストレージ間のデータ移動、非同期プリフェッチとスケジューリングのずれ、分散キャッシュ状態の同期の困難さといったさらなる課題に直面しており、理論上の利点が直接達成されるのを妨げています。
ハイブリッド SWA に加えて、MoE は分散スケジューリングとロードバランシングに大きな要求を課し、マルチモーダルエンコーダーは大規模画像および長尺ビデオシナリオにおけるスループットのボトルネックとなります。スケジューリング戦略と Prefill/Decode 実行パイプラインも慎重な最適化が必要です。この記事では、KVCache 管理、階層型キャッシュシステム、SWA を意識したプレフィックスキャッシュツリー、スケジューリング戦略、Prefill/Decode 実行パイプライン、マルチモーダル最適化を網羅する、MiMo-V2.5 シリーズの推論システムのエンドツーエンドのエンジニアリングプラクティスを紹介し、アーキテクチャの理論的な効率ポテンシャル(特に Hybrid SWA)をプロダクションで体系的に実現します。
1. ハイブリッド SWA:推論効率の利点
具体的な最適化に入る前に、まずハイブリッド SWA の理論的な効率の限界を定量化しましょう。これは設計上のアーキテクチャ上の根拠であり、後続のすべての最適化の基準となります。
1.1 計算分析
MiMo-V2.5-Pro を例にとると、モデルは合計 70 層で、10 層がフルアテンション、60 層が SWA で、スライディングウィンドウサイズは 128 です。フルアテンションと比較して、ハイブリッド SWA の計算コストは以下の図に示されています。SWA 層は全層の 6/7 を占めるため、ハイブリッド SWA アーキテクチャの総計算量はフルアテンションの約 1/7 です。プリフィルが主に計算バウンドであるチャンクプリフィルシナリオでは、これはプリフィルコストの比例的な削減に直接つながります。
1.2 KVCache ストレージ分析
SWA 層は、シーケンス全体ではなく、スライディングウィンドウ内の KV のみを保持する必要があるため、KVCache メモリ使用量も約 1/7 に低下します。デコードフェーズは主にメモリバウンドであり、そのレイテンシはモデルパラメータと KVCache の読み取りバイト数の合計に比例します。長尺シーケンスの場合、KVCache の量はモデルパラメータをはるかに超える可能性があるため、KVCache ストレージの削減は、長尺シーケンスシナリオにおけるデコードコストの削減にほぼ直接つながります。
KVCache ストレージはモデルアーキテクチャによって大きく異なり、アクセスパターンも異なります。以下の図に示すように、MiMo-V2.5-Pro と MiMo-V2.5 は、DeepSeek-V4-Pro と DeepSeek-V4-Flash に次いで KVCache 効率で 2 位にランクされています。
実際のコストの違いが KVCache サイズの比率に厳密に対応しないことに注意することが重要です。なぜなら、シーケンス長に依存しない固定の計算およびメモリアクセスコストが存在するからです。しかし、長文コンテキストシナリオでは、全体的な傾向は変わりません。短いシーケンスでは利点はわずかですが、シーケンスが長くなるほど、推論コストの利点は大きくなります。
2. KVCache システムのリファクタリング
MiMo-V2 および MiMo-V2.5 シリーズは、ハイブリッド SWA アーキテクチャを採用した初期のモデルの 1 つでしたが、当時、主流のオープンソース推論フレームワークやキャッシュシステムは、SWA を完全にサポートしていませんでした。MiMo API をローンチした際、私たちは SGLang v0.5.5 をサービングバックエンドコードベースとして選択しましたが、すぐに深刻な課題に直面しました。そのバージョンでは、SGLang の HiCache は SWA をサポートしていませんでした。あるいは、初期の SWA サポートは、互換性を維持するためにフル KVCache を保存することで実装されていました。SWA をより使いやすくするための回避策はいくつかありましたが、私たちはより高いパフォーマンス上限とより優れた使いやすさを持つ KVCache システムを構築したいと考えていました。
2.1 SWA KVCache 管理
KVCache デュアルプール設計
ハイブリッド SWA は、根本的なストレージの競合をもたらします。フルアテンションレイヤーはシーケンス全体の KV (O(N)) を保存する必要がありますが、SWA レイヤーはスライディングウィンドウ内の KV (O(W)) のみを維持する必要があります。従来の単一 KV プール設計では、システムはすべてのレイヤーに対して O(N) の GPU メモリを割り当てる必要があり、SWA のウィンドウ疎性を活用できず、事実上ほぼフル KVCache 実装に退化してしまいます。
自然な解決策は、フルアテンションと SWA のために KVCache を 2 つの独立したプールに分割し、システムレベルで統一された抽象化を提供することです。
物理レイヤー:個別のフル KV プールと SWA KV プールを維持します。SWA プールはウィンドウ用にのみサイズ設定され、ウィンドウに基づいた独立したエビクションをサポートし、SWA ストレージを O(W) に厳密に制約します。このメカニズムは L2 および L3 ストレージ階層にも拡張されます。
論理レイヤー:単一のシーケンスビューを上位レイヤー(プレフィックスツリー、スケジューラ、トランスポートプロトコル)に公開し、フルアテンションインデックスを権威ある参照とし、透過的な階層型ストレージのためにフル → SWA マッピングを維持します。
スケジューリング制約:システムは、リクエストを受け入れる際にフル KV と SWA KV の両方の容量制約を検証し、単次元チェックからのリソースの誤割り当てを回避します。
データ移動:階層間の転送は SWA マスクのみに基づいて実行され、有効なウィンドウデータのみが移動され、冗長な帯域幅消費が回避されます。
この設計により、SWA KVCache はシステムレベルで厳密な O(W) ストレージ制約を達成し、全体的な KVCache 容量効率を約 7 倍向上させ、ハイブリッド SWA の構造的利点を解き放ちます。主流の推論フレームワークも同様の実装アプローチを採用しています。
レイヤーごとの KVCache プリフェッチ
SWA KVCache ストレージ最適化が完了すると、SWA 層は最小限の KVCache 量のみをプリフェッチすればよくなります。これにより、ホストからデバイスへの KVCache プリフェッチと計算とのレイヤーごとのスケジューリングによるほぼ完璧なオーバーラップが可能になり、推論中のキャッシュ読み取りコストはほぼゼロになります。
SWA を意識したプレフィックスキャッシュツリー
従来の RadixAttention のヒットルールは、単純な仮定に基づいています。つまり、同じトークンシーケンスは同じ KV を持つということです。この仮定はフルアテンションでは成り立ちます。2 つのリクエストが同じトークン ID を共有している限り、それに対応する KV はプール内に存在し、直接再利用可能であることが保証されます。
しかし、この仮定は SWA では破綻します。その理由は、プレフィックスツリーの論理的なライフサイクルと SWA KV の物理的なライフサイクルが一致しないためです。プレフィックスツリーノードの長さは SWA ウィンドウによって制約されません。ノードの...