SGLang即座にMiMo-V2-Flashモデルをサポート

はじめに

XiaomiMiMo/MiMo-V2-Flashモデルは総パラメータ数309B、アクティベーションパラメータ数15Bで、推論最適化のために特別に設計された新型モデルであり、デコード効率の最大化を目指しています。その2つの重要な設計はスライディングウィンドウアテンション(SWA)多層MTPです。MiMo-V2-Flashは実際のサービスシナリオに合わせてカスタマイズされており、異なるハードウェア上でスループットとレイテンシを柔軟にバランスすることをサポートしています。SGLangが最適化したSpec v2ランタイムと組み合わせることで、ほぼゼロオーバーヘッドで多層MTPと効率的なSWA実行をサポートし、H200上でバランスの取れたTPOTとスループットを実現しています。本文では、モデルアーキテクチャとSGLangの高効率サポートについて紹介します。

推論効率的モデリング

MiMo-V2-Flashは推論効率の原則に従い、2つのコア設計を採用しています:

  • スライディングウィンドウアテンション(SWA):各トークンの受容野を固定サイズの定数ウィンドウ内に限定し、アテンションの複雑度をシーケンス次元のO(N²)からO(Nw)に削減します。ここでwはウィンドウサイズです。
  • MTP:多層MTPは予測ヘッドチェーンを使用し、層ごとに順次次のトークンを予測し、その後拡張クエリを使用してドラフトトークンを並列検証します。

下図はMiMo-V2-Flashの全体アーキテクチャを示しています:

MiMo-V2-Flash全体アーキテクチャ概要図

MiMo-V2-Flash全体アーキテクチャ

SWA

MiMo-V2-Flashは5層のSWAアテンション層ごとに1層の密なGQAを交互に配置しています。SWAは多角的に推論効率を向上させます:プリフィル段階では、計算が主要なコストとなり、長いシーケンスではO(N²)のアテンションがボトルネックになりますが、SWAは線形複雑度に削減し、TTFTを大幅に短縮します。同時に、KVキャッシュの複雑度を定数レベルに削減し、より大きなバッチをサポートするためのリソースを解放し、KVロード操作を削減し、TPOTを向上させます。

下図はプリフィルベンチマークテストの結果です:

MiMo-V2-Flashプリフィルベンチマークテスト(Radix Cache無効)

MiMo-V2-Flashプリフィルベンチマークテスト(Radix Cache無効)

MTP

MiMo-V2-Flashの重要な設計は3層の多層MTPです。デコードシナリオでは、ほとんどのカーネルはメモリバウンドで、クエリ長は常に1であるため、並列デコードトークンを増やすことがスループットを向上させる最も直観的な方法です。しかし、バッチサイズが増大すると、KVキャッシュアクセスが線形に増加してボトルネックとなり、この時点では計算能力はまだ飽和しておらず、単純にバッチを増やすだけでは限界があります。

MTPは残りの計算を活用します:複数のトークンが順次予測ヘッドによって同時に生成され、同じクエリで並列検証され、KVアクセスを増やすことなくクエリ長を延長し、算術強度を向上させます。メモリバウンドでバッチ効果が限界に達した場合、積極的なMTP戦略(高い受理率)はデバイスの潜在能力を十分に活用し、TPOTを最適化できます。

ハードウェア認識MTP設定

MTPは飽和していない算術強度から恩恵を受け、MiMo-V2-FlashのGQAアテンションは自然に適合します。しかし、デプロイ時には最適な計算-メモリバランスを実現するために、適切なバッチサイズとMTP深度を選択する必要があります。高い屋根線デバイス(トレーニング向けGPUなど)はより積極的なMTPに適しており、豊富な計算を活用できます。推論アクセラレータ(H20など)はFLOPsが限られているため、MTPは慎重に行う必要があり、計算バウンドに転換してスループットが低下するのを避ける必要があります。

H200ベンチマークは、MiMo-V2-FlashがスループットとリクエストあたりのTPSを両立していることを示しています。64Kの長いコンテキストで、各DP rankのバッチが16であっても、デコードスループットは150 TPSに達しており、これはSWAとMTPのおかげです。

MiMo-V2-Flashデコードベンチマーク(DP 2, TP 4, EP 8, MTP受理長3.6, 入力16k, 可変バッチ)

MiMo-V2-Flashデコードベンチマーク(DP 2, TP 4, EP 8, MTP受理長3.6, 入力16kトークン, 可変バッチ)

MiMo-V2-Flashデコードベンチマーク(DP 2, TP 4, EP 8, MTP受理長3.6, 各DPバッチ16, 可変入力長)

MiMo-V2-Flashデコードベンチマーク(DP 2, TP 4, EP 8, MTP受理長3.6, 各DP rankバッチ16, 可変入力長)

SGLang Spec v2の高速MTPサービス

MiMoの多層MTPはSGLang Spec v2にネイティブ統合されており、完全オーバーラップMTP機能を活用してスループットとレイテンシを向上させています。Spec v2では、オーバーラップスケジューリングと投機的デコードが融合しています:遅延出力の同期/処理、次のバッチカーネルの早期発行により、CPU バッチ/同期オーバーヘッドをGPUフォワードに隠蔽し、GPUバブルを削減します。

Spec v2オーバーラップ投機的デコードパフォーマンスプロファイル

Spec v2オーバーラップ投機的デコードパフォーマンスプロファイル

さらなる議論

LLMサービスのデコード段階は多くがメモリバウンドで、主流のトレーニングGPUでは計算リソースが深刻に遊休しています。推論専用アクセラレータは帯域幅が高くFLOPsが低いため経済的ですが、速度は限られています。MiMo-V2-Flashはモデル側から推論効率を最適化し、多層MTPは汎用ソリューションになる可能性があります:受理率を最適化し、GPU計算を活用してデコードを加速します。より適合したアーキテクチャはハードウェア選択を柔軟にし、同じタイプのハードウェアでトレーニングと推論の両方に対応でき、デプロイを簡素化しコストを削減します。

MiMo-V2-FlashサポートはすでにSGLang PR(#15207, #15208)で実装されており、まもなくメインブランチにマージされます。本文のベンチマークはMiMo最適化ブランチに基づいており、最適化はSGLangメインラインにアップストリームされます。

クイックスタート

MiMo-V2-Flashは現在、SGLang Dockerイメージとpipインストールで利用可能です。以下はSGLangサーバー起動ガイドです。

Docker
# Pull the docker image
docker pull lmsysorg/sglang:dev-pr-15207

# Launch the container
docker run -it --gpus all \
  --shm-size=32g \
  --ipc=host \
  --network=host \
  lmsysorg/sglang:dev-pr-15207 bash

# Start the server
SGLANG_ENABLE_SPEC_V2=1 python3 -m sglang.launch_server \
        --model-path XiaomiMiMo/MiMo-V2-Flash \
        --dp-size 2 \
        --enable-dp-attention \
        --tp-size 8 \
        --trust-remote-code \
        --mem-fraction-static 0.75 \
        --max-running-requests 128 \
        --chunked-prefill-size 16384 \
        --reasoning-parser qwen3 \
        --tool-call-parser mimo \
        --model-loader-extra-config '{"enable_multithread_load": "true","num_threads": 64}' \
        --attention-backend fa3 \
        --speculative-algorithm EAGLE \
        --speculative-num-steps=3 \
        --speculative-eagle-topk=1 \
        --speculative-num-draft-tokens=4 \
        --enable-mtp
Pipインストール
# On a machine with SGLang dependencies installed or inside a SGLang nightly container
# Start an SGLang nightly container
docker run -it --gpus all \
  --shm-size=32g \
  --ipc=host \
  --network=host \
  lmsysorg/sglang:nightly-dev-20251215-4449c170 bash

# If you already have SGLang installed, uninstall the current SGLang version
pip uninstall sglang -y

# Install the PyPI Package
pip install sglang==0.5.6.post2.dev8005+pr.15207.g39d5bd57a \
  --index-url https://sgl-project.github.io/whl/pr/ \
  --extra-index-url https://pypi.org/simple

#Launch the server
SGLANG_ENABLE_SPEC_V2=1 python3 -m sglang.launch_server \
        --model-path XiaomiMiMo/MiMo-V2-Flash \
        --dp-size 2 \
        --enable-dp-attention \
        --tp-size 8 \
        --trust-remote-code \
        --mem-fraction-static 0.75 \
        --max-running-requests 128 \
        --chunked-prefill-size 16384 \
        --reasoning-parser qwen3 \
        --tool-call-parser mimo \
        --model-loader