MiniMax-M3の効率的な推論サービス:100万トークンコンテキストとマルチモーダルを後悔なく実現
Together AIは、KVブロックメジャースパースアテンション、ページ化MSAデコード、最適化されたインデックススコアリングカーネル、およびRustベースのマルチモーダル前処理ゲートウェイにより、MiniMax M3の効率的なサービスを実現し、同時実行レベル全体で81~125%のスループット向上を達成しました。
Together AIはMiniMaxと協力し、最新フラッグシップモデルM3の優先クラウドパートナーとなり、オープンウェイト公開後には開発者向けエンドポイントとしてモデルを提供します。M3は、コーディング性能、エージェントワークフロー、ネイティブマルチモーダル推論を統合し、100万トークンのコンテキストウィンドウを経済的にサポートするように設計されています。Together AIの推論およびカーネルチームは、M3を効率的にサービスするために、KVブロックメジャースパースアテンションカーネル、ページ化アテンション統合、最適化されたインデックススコアリングカーネル、Rustベースのマルチモーダル前処理ゲートウェイなどの重要な最適化を実装し、異なる同時実行レベルで81~125%のスループット向上を達成しました。
M3の最も革新的なアーキテクチャ変更はMiniMaxスパースアテンション(MSA)です。これはブロックスパースアテンションメカニズムにより、各クエリが注目できるトークン数を制限し、長コンテキスト処理のコストを削減します。プリフィル段階で9倍以上、デコード段階で15倍以上の高速化を実現します。MSAの計算は、各KVグループに対して最も関連性の高いKブロックを決定するスコア計算と、それらのブロックに対するクエリトークンの密なアテンションの2部構成です。この設計により、KVグループ次元での表現力を維持しつつ、クエリトークンが注目するKVトークン数に制限を設け、アテンション計算がコンテキスト長に対してO(N^2)でスケールしなくなるため、長コンテキストワークロードに非常に適しています。
KVブロックメジャースパースアテンションカーネルでは、プリフィル中に外側ループでKVグループを反復し、内側ループでクエリトークン間のアテンションを計算することで、KVキャッシュの移動を1回に削減し、演算強度を向上させます。マッピングを{q, kv block}から{kv block, q}に再編成し、アテンションカーネルを再実装し、最後にLog-Sum-Expに基づくリダクションを行います。
ページ化アテンションとの統合では、デコード中に選択されたブロックに基づいてページテーブルを構築し、KVグループ次元をバッチ次元にフラット化し、KVキャッシュテンソルのストライドビューを活用して既存のGQA対応アテンションカーネルを使用できるようにしました。各クエリの選択ブロック数が限られているため、ブロックからページへのマッピングのオーバーヘッドは非常に低く、この設計によりデコードスループットが5%向上します。
デコードインデックススコアリングカーネルの最適化では、MSAがコストを密なアテンションからスコア/top-kインデクサーに移します。各デコードクエリに対して、エンジンはクエリ側インデックスベクトルと候補キー側インデックスベクトルを比較し、128トークンのKVブロックを単一スコアにリダクションし、実際のアテンションカーネル用にトップブロックのみを保持します。最適化されたパスでは、AB交換HMMAレイアウトを使用し、128トークンのキーインデックスブロックをMMA M次元とし、クエリ側をより小さいN次元にパディングします。カーネルは非同期コピーで128トークンのKインデックスをステージングし、次のページをプリフェッチし、bfloat16のHMMAでドット積を計算し、各ページを1ブロックスコアにリダクションします。
マルチモーダル前処理は、RustベースのSMG(サービングモデルゲートウェイ)が担当します。SMGはOpenAI互換APIと推論エンジンの間に位置し、ルーティングとトークン化に加えて、すべてのビジョン前処理をCPUで実行します。画像やビデオのダウンロード、デコード、フレームサンプリング、リサイズ、パッチテンソル変換などを行います。M3の場合、SMGはFFmpegでビデオをフェッチし、FPSに基づいてフレームを選択し、リサイズと正規化を行い、時間次元を組み込んでパッチ化します。出力はフラットなパッチテンソルとグリッドメタデータテンソルで、gRPCメッセージにパックされます。ワーカーはビジョンエンコーダを直接実行するだけで、前処理は不要です。SMGのマルチモーダルパイプラインはRustトレイトで構造化されており、モデル固有の前処理ロジックをパイプラインから分離しています。M3のサポート追加は、M3固有の定数でトレイトを実装するだけで済み、パイプライン自体は変更されません。このアーキテクチャは、ビジョン機能を持つほとんどのオープンソースモデルに適用可能で、推論エンジンのランタイム間で一般化できます。
パフォーマンス結果では、一般的なエージェント型トラフィックパターン(60Kプレフィックスキャッシュ、同時実行8、NVIDIA B200)において、MSAがイテレーションあたりのアテンション計算の壁時間割合を大幅に削減しました。全体として、さまざまな同時実行レベルで81~125%のスループット向上を達成しました。
将来の取り組みとして、スパースアテンションアーキテクチャにより導入されたより小さなカーネル(kvブロック上のtopk、q-kvマッピングのkv-qへの再マッピングなど)のカーネル融合の機会を探り、カーネルエージェント研究チームがプロダクショングレードのカーネルを記述するエージェントを開発しています。また、kインデックスと実際のKVキャッシュのCPUキャッシュオフロードを分離し、topk選択に基づいてオンデマンドでKVキャッシュをロードする作業を進めています。Together AIは、M3の推論パフォーマンスをさらに最適化し、開発者に安定した効率的なエンドポイントを提供し続けます。