🤖 machine learning

Lynx: Enabling Efficient MoE Inference through Dynamic Batch-Aware Expert Selection

Lynx は、トレーニングによって誘発される活性化の偏りを活用し、新規の AffinityBinning 技術を用いてトークンからエキスパートへの割り当てを動的に再マッピングすることで、エキスパートの呼び出し数を削減し、精度を大幅に損なうことなくスループットを最大 1.30 倍向上させる、ワークロードに依存しないシステムです。

原著者： Vima Gupta, Jae Hyung Ju, Kartik Sinha, Ada Gavrilovska, Anand Padmanabha Iyer

公開日 2026-05-20

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Vima Gupta, Jae Hyung Ju, Kartik Sinha, Ada Gavrilovska, Anand Padmanabha Iyer

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

「The MoE Kitchen」という巨大で高級なレストランを運営している状況を想像してください。

このキッチンでは、すべての料理を作ろうとする一人の巨大なシェフを置く代わりに、64 人の専門的な sous-chef（副料理長）からなるチーム（Experts）を持っています。そして、すべての注文（Token）を見て、その料理を作るためにどの 8 人のシェフを特定するかを決定する、ヘッドウェイター（Router）がいます。

この仕組みは素晴らしいものです。なぜなら効率的だからです。すべての 64 人のシェフにすべての注文を処理させる必要はありません。必要な 8 人だけに支払えばよいのです。これが Qwen や Llama などの現代の AI モデルが機能する方法です。これらは巨大ですが、生成する各単語に対して、その「脳」の小さな部分だけを「目覚め」させます。

問題：ラッシュアホールの渋滞

この論文は、レストランが混雑する際に発生する重大な問題（Batchingと呼ばれる現象）について説明しています。

顧客が 1 人だけの場合、ヘッドウェイターは注文を特定の 8 人のシェフに送ります。簡単です。
しかし、16 人の顧客のバッチがある場合、ヘッドウェイターはすべての 16 件の注文を見ます。すべての顧客が異なるため、ウェイターは結果的にグループを処理するために、ほぼすべての 64 人のシェフを呼び出す必要に迫られます。

ボトルネック：
キッチンは、材料（シェフの知識）が巨大で高速なパントリー（GPU Memory）に保管されるように組織されています。料理をするために、シェフたちは特定の材料を手に取るためにパントリーへ走らなければなりません。

問題点： 16 人の顧客が到着すると、パントリーが溢れかえります。シェフたちは実際に料理をする時間よりも、材料を取りに行くために往復する時間の方を多く費やします。キッチンが遅くなるのは、「料理」する（計算）ことではなく、「走る」こと（メモリ帯域幅）がボトルネックになっているためです。
結果： AI は高速で効率的であるはずですが、複数のリクエストを同時に処理する際には、データの取得に詰まってしまうため、動きが極端に遅くなります。

解決策：LYNX（賢いウェイター）

著者たちはLYNXと呼ばれる新しいシステムを作成しました。LYNX を、キッチンが混雑しているとき（「decode」フェーズ、つまりレストランが一口ずつ食事を提供しているような状況）にのみ介入する、超スマートで動的なマネージャーだと考えてください。

LYNX はシェフを解雇したり、メニューを変更したりしません。代わりに、AffinityBinningと呼ばれる巧妙なトリック（「信頼度によるグループ化」という意味の洗練された表現）を使用します。

LYNX がどのように機能するか、ステップバイステップで説明します。

「信頼度」チェック：
時々、ヘッドウェイターは「シェフ A がこの料理に最適だ」と 100% 確信しています。他の時には、ウェイターは不確実で、ルールが「8 人の異なる人を選ばなければならない」と言っているため、シェフ B を選びます。この論文では、これらの「不確実な」選択は往々にして冗長であることがわかりました。
- 比喩： 友人に映画の推薦を求めたとき、彼らが「確信はないけど、もしかしたら映画 Xか映画 Yかな」と言った場合、彼らはどちらにも本気でコミットしているわけではありません。もう一度尋ねれば、彼らは再び映画 Xを選ぶかもしれません。
「ビンニング」戦略：
LYNX はウェイターの信頼度スコアを確認します。そして、注文を「バケツ」にグループ化します。
- 高信頼度：「この注文は必ずシェフ A に渡さなければならない。」（LYNX はこれに手を出しません）。
- 低信頼度：「この注文はシェフ B になんとなく向いている。」（LYNX は「実際には、これをシェフ A に送ろう。なぜならシェフ A はすでに他の注文のためにキッチンにいるからだ」と言います）。
大再マップ：
LYNX は、その「低信頼度」の注文を取り、それらをバッチによってすでに使用されているシェフたちへリダイレクトします。
- 魔法： 64 人の異なるシェフのためにパントリーへ走る代わりに、キッチンは今や例えば 30 人のシェフのためにパントリーへ走るだけで済みます。
- 結果： シェフたちは走る時間を減らし、料理する時間を増やします。キッチンははるかに速く動きます。

これが特別である理由

この論文は、LYNX が重要である 3 つの主要な理由を強調しています。

「ワークロード非依存」であること： LYNX は特定の種類の顧客やメニューで訓練される必要はありません。それは毎回、その場でパターンを把握します。マニュアルを必要とせずに客の習慣を瞬時に学習するマネージャーのようなものです。
何も壊さないこと： LYNX はシェフを解雇したり、レシピを変更したりしません。それは単に、その特定の瞬間のために誰が何をするかを再配置するだけです。論文は、料理（AI の回答）が同じように美味しく、時にはさらに良くなることを示しています。なぜなら、シェフたちに自信のない料理を強要するのをやめるからです。
他者と協調すること： LYNX は、さらに速くするために、シェフのエプロンを縮めたり、別の建物へ送ったりするなどの他の速度向上テクニックの上に追加して使用できます。

結果

著者たちは、コーディング、数学、推論などのタスクにおいて、Qwen、Mixtral、Llama などの実世界の AI モデルでこれをテストしました。

速度： 最悪の場合でも、システムは1.3 倍速くなり（30% の速度向上）、
精度： 回答は同じくらい正確か、わずかに優れていました。最悪の場合でも、精度の低下は 1% 未満で、ほとんど気づかれないレベルです。
効率性： 遅延することなく、一度に多くの顧客を処理することを可能にしました。

まとめ

LYNXは、混雑した AI キッチンのための交通整理員のようなものです。注文のグループが入ってくると、キッチンがあまりにも多くの異なるシェフのためにパントリーへ走る時間を浪費していることに気づきます。そこで、それは巧妙に「多分」の注文を、すでに働いているシェフたちへリダイレクトし、渋滞を減らして食事をより速く提供します。すべてはメニューを変更したり、誰かを解雇したりすることなく行われます。

技術概要：LYNX – 動的なバッチ認識型エキスパート選択による効率的な MoE 推論の実現

1. 問題提起

現代の基盤モデルは、パラメータ数と計算コストを分離するために、 increasingly Mixture-of-Experts（MoE）アーキテクチャを採用しています。MoE はトークンごとにパラメータの小さなサブセットのみをアクティブ化します（例：Qwen2-57B では 570 億パラメータ中 14 億のみ）。しかし、この効率性の約束は、レイテンシとスループット間の緊張関係により、本番環境での推論サービスにおいて破綻しています。

高スループットを達成するために、推論システムはリクエストをバッチ処理します。しかし、バッチ内の各トークンが独立してトップ- $k$ のエキスパートを選択するため、高帯域幅メモリ（HBM）からロードする必要があるエキスパートのセットは、バッチサイズと構成に応じて増加します。レイテンシのサービスレベル目標（SLO）を満たすために必要な中程度のバッチサイズにおいて、この動的でデータ依存性の高いメモリアクセスパターンにより、MoE デコーディングはメモリ帯域幅に縛られる状態となります。その結果、レイテンシはバッチ全体でアクティブ化された異なるエキスパートの総数に比例して増加し、しばしばエキスパートプール全体を飽和させます（例：わずか 8 つの多様なリクエストからなるバッチでも、Qwen2-57B の全 64 個のエキスパートをアクティブ化し得る）。これは MoE のスパース性の利点を無効化し、動的なエキスパート移動のオーバーヘッドが加わった状態で、同等容量の密なモデルと同様のパフォーマンスしか発揮しないことになります。

既存の解決策、例えば静的プルーニング、量子化、またはクラスタリングは、オフライン較正に依存し、モデルを永続的に変更するため、ワークロードの変化に対して柔軟ではありません。動的なトークンごとのアプローチは、バッチレベルでフェッチされるエキスパートの総数を削減しないため、ボトルネックに対処できません。

2. 手法：LYNX

LYNX は、モデルアーキテクチャを変更したりオフライン較正を必要としたりすることなく、MoE 推論におけるメモリ帯域幅のボトルネックを緩和するように設計された、ワークロード非依存のシステムです。これは、MoE のトレーニングダイナミクスから導き出された重要な洞察を利用します：ロードバランシング損失は、セカンダリ選択への確信度が低い場合でも、ルーターがトークンをエキスパート全体に広く分散させるように強制します。これにより、ルーターが真のトークン - エキスパート親和性ではなく、トレーニング上の制約のためにエキスパートを選択するバッチレベルの冗長性が生じます。

LYNX は、この冗長性を、完全にランタイムで動作する革新的なAffinityBinning技術と 3 つのコンポーネントからなるパイプラインによって利用します。

2.1 中核原理

この設計は、3 つの実証的観察に基づいています。

フェーズ感受性：MoE はプリフィルフェーズ中は計算に縛られますが、デコードフェーズ中はメモリ帯域幅に縛られます。エキスパートの再マッピングが有益なのはデコードフェーズのみです。さらに、プリフィルフェーズはエキスパートの忠実度に対して非常に敏感であるのに対し、デコードフェーズは再割り当てに対して頑健です。
トークン確信度：ルーターの確信度スコア（トップランクのエキスパートのロジット間の差から導出される）は、どのトークン - エキスパートマッピングが重要かを確実に識別します。低確信度のトークンは、精度の低下を最小限に抑えながら再マッピング可能です。
エキスパートランクの階層性：トップランクのエキスパート（ランク 0）は出力品質を支配します。トップ- $k$ に含まれるより低いランクのエキスパートは高い冗長性を示します。再割り当てに対する感受性は、層によっても異なります。

2.2 システムコンポーネント

LYNX は、フェーズ認識型オプティマイザーと 3 つの軽量 CUDA カーネルコンポーネントを介して、サービスエンジン（例：vLLM）と統合されます。

フェーズ認識型オプティマイザー：現在のイテレーションがメモリに縛られている（デコード）のか、計算に縛られている（プリフィル）のかを判定します。不要なオーバーヘッドを避けるため、プリフェーズまたはチャンクドプリフィルフェーズ中は LYNX のロジックをバイパスします。
確信度アナライザー：各層でルーターのロジットをインターセプトします。各トークンのルーター確信度をAffinityBinningを使用してビンに離散化します。ビンの幅と数は、モデルのスパース性比率（ $k/N$ ）のみによって決定されるため、システムは自己較正可能であり、ワークロードに依存しません。これにより、割り当てられたエキスパートとの親和性が低いトークンを識別します。
適応型エキスパートスコアラー：各エキスパートに対してバッチ全体のスコアを計算します。確信度アナライザーによって割り当てられたビン値に基づいた指数重み付け方式を使用します。これにより、高確信度のトークンに好まれるエキスパートが、不均衡に高いスコアを受け取ることを保証します。これにより、バッチに必要な「重要」なエキスパートの最小セットを決定します。
エキスパートリマッパー：親和性の低いトークン - エキスパート割り当てを、生存した重要エキスパートセットへリダイレクトします。重要なのは、LYNX がトップ- $k$ アクティブ化セマンティクスを保持することです。各トークンは依然として正確に $k$ 個のエキスパートをアクティブ化しますが、バッチ全体としては、ユニークなエキスパートのより小さく統合されたセットのみをアクティブ化します。

2.3 実装

LYNX は、vLLM 上のプラグアンドプレイモジュールとして実装されています。その中核ロジックは、ルーター出力をインターセプトする 4 つの融合された Triton/CUDA カーネルによって実現されています。この設計は、中間テンソルデータの移動を排除し、カーネル起動のオーバーヘッドを削減することで、クリティカルパス内で動作しながらもシステムが追加するレイテンシを無視できるレベル（<4% のオーバーヘッド）に抑えています。

3. 主要な貢献

初のワークロード非依存システム：LYNX は、較正データや永続的なモデル変更を必要とせずに、MoE 推論におけるメモリ帯域幅のボトルネックを緩和する最初のシステムです。
AffinityBinning 技術：アーキテクチャ上のスパース性比率に基づいてルーター確信度を離散化する革新的な手法であり、任意の MoE アーキテクチャへの自動適応を可能にします。
バッチレベル最適化：従来のトークンごとの動的アプローチとは異なり、LYNX はバッチレベルで動作し、HBM からフェッチされるユニークなエキスパートの総数を削減することで、メモリ帯域幅のボトルネックに直接対処します。
相補性：LYNX は、オフロードや量子化を含む既存の最適化と併用して動作するように設計されています。

4. 実験結果

LYNX は、4 つのファミリー（Qwen、DeepSeek、Mixtral、Llama）に属する最先端モデルを、コード生成、数学、推論、視覚の 8 つのベンチマークで評価しました。

レイテンシ削減：LYNX は、ベースラインと比較して最大1.30 倍のレイテンシ削減（出力トークンあたりの時間）を達成します。プリフィルとデコードが分離された非集約型サービスシナリオでは、削減率は最大 1.35 倍に達します。
精度への影響：システムは最悪の場合でも1% 未満の精度低下しか引き起こしません。注目すべきは、LYNX が平均的に精度を向上させることが頻繁にあることです（例：DeepSeek-Coder-V2 における HumanEval で +1.9%）。これは、トレーニング時のロードバランシングが LYNX によって修正されるべき非最適なエキスパート選択を強制していることを示唆しています。
スループット：固定された SLO バジェット（例：25ms または 50ms）の下では、LYNX はシステムスループットを最大2.1 倍向上させます。
互換性：
- オフロード：オフロード技術（例：Fiddler）と組み合わせると、LYNX はパフォーマンスを**31%**向上させます。
- 量子化：INT4 量子化の上に適用すると、LYNX は追加で**10%**の高速化を提供します。
- 並列化：LYNX は、エキスパート並列化（EP）およびテンソル並列化（TP）の上に追加的なスループット向上をもたらします。
堅牢性：パフォーマンスの向上は、異なるバッチサイズ（中程度のサイズで飽和）およびバースト的な到着率を持つ実世界のトレースにおいて一貫しています。

5. 意義

本論文は、GPU 計算能力がメモリ帯域幅の成長を上回り続けるにつれて、MoE デコーディングにおける算術強度のギャップが拡大し、メモリに縛られるボトルネックがますます重要になると主張しています。LYNX は、バッチレベルでのエキスパート選択のあり方を根本的に変えることで、この課題に対処します。トレーニング時のロードバランシングによって導入された冗長性を活用することで、LYNX はモデルの再トレーニングやオフラインプロファイリングを必要とせず、MoE アーキテクチャが将来のハードウェア上でその効率性の約束を実現することを可能にします。このシステムは、動的でランタイム認識型のエキスパート管理が、現在の MoE サービスパラダイムに内在するメモリ帯域幅の制約からスループットを大幅に分離できることを実証しています。