WiSP: A Working-Set View of Mixture-of-Experts Serving on Extremely Low-Resource Hardware

技術要約：WiSP – 極めてリソースの乏しいハードウェアにおけるMixture-of-Expertsサービングのワーキングセット・ビュー

問題提起
Qwen3-30B-A3BやMixtral-8×7Bといった現代的なMixture-of-Experts（MoE）モデルは、パラメータの大部分をエキスパート層に配置しているが、任意のトークンに対してアクティブになるエキスパートはそのごく一部である。VRAMが限られた汎用GPUにおいて、これらのモデルはメモリ内に完全に収まらないことが多い。標準的な解決策である静的なレイヤーレベルのCPUオフロード（例：vLLMの--cpu-offload-gb）は、メモリ使用量を低く抑える一方で、フォワードパスのたびにレイヤー内のすべてのエキスパートをPCIeバス経由で転送する。このアプローチはMoEモデルの疎性（sparsity）を活かすことができず、その効率性のメリットを事実上打ち消してしまう。

著者らは、低リソース環境でのMoEサービングをワーキングセット管理問題として定義している。このレジームにおいて、GPUメモリは限られた高速なストアであり、ホストメモリはバックイングストアとなる。2つの異なるメモリ需要ストリームが同じVRAM予算を競い合う：

1. エキスパート重み： ルーティングに依存するストリームであり、トークンごとに一部のエキスパートのみが必要となる。
2. KVキャッシュ： 自己回帰的生成に必要となる、トークンに依存するストリーム。

既存のシステムはこれらを別々の問題として扱ったり、計算とデータ転送をオーバーラップさせることができる高バッチのデータセンター設定を対象としたりすることが多い。しかし、低コンカレンシーのインタラクティブなサービング（例：ローカルのコーディングアシスタント）では、ステップあたりの計算量は極めて小さいため、計算量や予測精度よりもPCIe帯域幅が主要なボトルネックとなる。

手法：WiSP (Working-Set Paging)
本論文では、居住するエキスパート集合をルーティング挙動に対するキャッシュとして扱い、KVキャッシュと共同で管理するシステムであるWiSPを提案している。WiSPは主に2つのコンポーネントで構成される：

1. ルーティング認識型エキスパート・ページャー (Routing-Aware Expert Pager):

   • メカニズム： WiSPは、既存のエキスパートマップ間接参照フック（$\pi$）を利用することで、修正なしのサービングエンジン（例：vLLM）にプラグインできる。フルレイヤーをロードする代わりに、再利用される可能性の高い「ワーキングセット」のエキスパートのみを含む小さなGPUスクラッチスペースを維持する。
   • 動作： 各ステップにおいて、ページャーは不足しているルーティング済みエキスパートのみをホストメモリからフェッチ（H2D）し、空き容量を作るために最も最近使用されていない（LRU）エキスパートを退避させる。
   • 保証： 基盤となるカーネルや重みは変更されず（マッピングのみがリダイレクトされる）、WiSPはベースラインのエンジンとバイト単位で同一の出力を生成する。新しいカーネル、量子化の変更、またはモデルの再学習を必要としない。

2. MV-WSA (Marginal-Value Working-Set Allocation):

   • 割り当て問題： エキスパート重みとKVキャッシュは同じVRAM予算を共有するため、居住するエキスパートセットを増やすとKV容量が減り、その逆も同様である。
   • 戦略： MV-WSAは、エキスパート予算とKV予算の間の最適な分割比（$f$）を決定する。これは、両方のプールにおけるバイトあたりの限界レイテンシ・ベネフィットを等しくすることを目指す。
   • アドミッション・フロア： 決定的なことに、アロケーターは現在のコンカレンシーレベルを許容するために必要な最小限のKVキャッシュサイズ（$\kappa_{min}$）を強制する。最適化された分割がKVキャッシュを枯渇させる場合、スケジューラが初期化および実行を確実にできるように、分割をクランプ（固定）する。
   • 実装モード：
     • オフライン・コンフィギュレーター： 代表的なトレースを分析し、起動時に固定の分割比を設定する。
     • オンライン・コントローラー： サービング中の「ドレイン・バリア（drained barriers）」（アクティブなリクエストがメモリブロックを保持していない時点）において、両方のプールを動的にリサイズする。これにより、総予算を超過することなく、エキスパート・スクラッチとKVプールの間で物理的にバイトを移動させる。

プリフェッチに関する主要な知見
本論文の重要な経験的知見は、ルーティングに基づくプリフェッチは、シングルストリームのデコードにおいてほとんど、あるいは全くメリットをもたらさないということである。

• ボトルネック： パフォーマンスは予測精度ではなく、PCIe帯域幅によって制限されている。バッチ1のデコードでは、ステップあたりの計算量が小さすぎるため、PCIe転送を計算の背後に隠蔽することができない。
• 結果： 推測プリフェッチは、帯域幅を競合させ、直後のステップで必要となるエキスパートを退避させてしまうため、しばしばパフォーマンスを低下させる（テストではスループットが約50%低下）。
• 再定義： ルーティング信号は、レイテンシを隠蔽するためではなく、メモリ割り当てのために価値がある。それは、最小限の「安全な」居住セットサイズを特定するのに役立ち、解放されたVRAMをKVキャッシュに割り当てることを可能にする。

結果
著者らは、NVIDIA H100（および小規模なメモリ予算をエミュレートした実験）を用い、様々なMoEアーキテクチャ（LLM、VLM、ハイブリッドSSM-MoE、および拡散モデル）でWiSPを評価している。

• スループットの向上： モデルが収まらない同一のVRAM予算条件下において、WiSPは静的なCPUオフロードと比較して最大1.95倍のデコード・スループットを達成した。この利得は、ベースラインがすべてをストリーミングしている一方で、WiSPが再利用されるエキスパートセットのかなりの部分をキャッシュできる中間的なメモリ予算において最大となる。
• KV容量： エキスパートのフットプリントを削減することで、WiSPはKVキャッシュ容量を直接的に増加させる。実験によれば、エキスパートのキャップ（上限）を縮小することで、同じ総メモリ予算において最大コンカレンシーを40%以上向上させることができる。
• 割り当て性能：
  • オフライン： MV-WSAコンフィギュレーターは、プリフィル（prefill）とデコードの両方で良好に機能し、シミュレーションにおいてワークフローごとのオラクル（理想的な解）の数パーセント以内の差に収まった。一方、固定分割（例：50/50）はこれより約20%劣る。
  • オンライン： 動的なデュアルリサイズ・コントローラーは、ワークロードの変化（例：長文脈のターンからツール呼び出しのバーストへの切り替え）に適応することで、最適な静的オフライン分割よりも1.20倍高いスループットを実現した。
• 汎用性： WiSPは、単一のGPUに収まらないモデル（例：Jamba-v0.1, MiniMax-M2）を正常にサーブでき、Vision-Language Model（Kimi-VL）や拡散モデル（DiT-MoE, LLaDA-MoE）を含むモダリティを横断して、バイト単位で同一の出力を維持する。

意義と主張
本論文は、エキスパートの居住性とKVキャッシングを単一のワーキングセット割り当て問題として扱うことで、WiSPが低リソースのMoEサービングのための統一されたフレームワークを提供すると主張している。

• 設計パラダイムの転換： 低コンカレンシーのサービングにおいては、焦点は「いかにエキスパートを速くプリフェッチするか」から、「いかにエキスパートとKVキャッシュの間でVRAMを最適に割り当てるか」へとシフトすべきであると論じている。
• 実用的な影響： WiSPは、出力の正確性を損なったり、特殊なハードウェア（Intel AMXなど）を必要としたりすることなく、汎用ハードウェア上で最先端の規模のMoEモデルのサービングを可能にする。
• 限界： 著者らは、WiSPはモデルがVRAMに収まらないレジーム向けに設計されていることを明記している。モデルが収まる場合、ページングによる不要な間接参照が発生するため、静的な居住（residency）の方が高速である。また、プリフェッチの最適化はシングルストリームのデコードでは効果的ではないことが示されたが、計算のオーバーラップが可能な高バッチまたはマルチターン・エージェントのセッションにおいては関連がある可能性がある。

結論として、メモリ制限のあるシングルストリームのデコードにおいて、ルーティング信号は「速度のレバー（レイテンシを隠蔽するもの）」ではなく、「メモリのレバー（ワーキングセットのサイズを決定するもの）」であり、インテリジェントな割り当てこそが主要な推進力となる。