The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin; Yiming Wang; Runyuan Cai; Hanze Liu; Xiaodong Zeng

論文 · arXiv 2609.18063

メモリの壁のもう半分：学習済みルーティング予測を用いたSSDからの35B MoE推論

The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng

公開日 2026-09-16cs.AI arXiv PDF

XYZNOT で読む（全文翻訳と AI 解説）

3行まとめ

🎯 35BクラスのMoEはint4でも約19.5 GBを必要とし民生用RAMを超過する一方、レイヤールーティングが前層の出力に依存するため、単純なSSDオフロードではディスク遅延でストールします。
🔧 Edge0は、1トークン先を予測する各層のprerouterにより「予測をルーティングとして採用」し、スチューデント経路で蒸留したマージを行わない並列復元LoRAと組み合わせて、SSDからMoEエキスパートをストリーミングします。
📈 24 GBのMac mini M4 Pro上で、2.9 GiBのピークメモリ内に抑えつつ35B MoEを約20.4 tok/sで推論し、fp16教師モデルとのベンチマーク平均差を3.9ポイント以内に留めました。

概要

民生用ハードウェアにおけるMixture-of-Experts（MoE）の推論は、重みメモリによって制約されます。35Bクラスのモデルは4-bitでも19.5GBに達し、スパース性はトークンあたりの計算量を削減するものの、保持すべきバイト数は減らしません。SSDへの単純なオフロードだけでは解決しません。なぜなら、レイヤーN+1のエキスパートはレイヤーNの出力が存在する前に選択される必要があり、計算の裏で隠蔽できるほど読み込みを早期に開始できないためです。本研究では、prerouterによってこのギャップを埋めるストリーミングMoE推論エンジン「Edge0」を提案します。各層に配置されたヘッドが次の層のルーティングを1トークン先回りして予測し、その予測をルーティングそのものとして利用するため、ステージングされたエキスパート群とルーティングされる群が完全に一致し、トークンのドロップが発生しません。また、スチューデントの経路で学習されたマージを行わない（unmerged）復元LoRAが、int4量子化やルーティングの置換によって失われた品質を補償します。単一の24GBマシンにおいて、Edge0は3GiBのピークアクティブメモリ内で35B MoEを20tok/sで提供し、5つの公開ベンチマークの平均でfp16の教師モデルから数ポイント差の性能を維持します。8Bモデルも同一フレームワークで動作し、フレームワーク、チェックポイント、アダプターはすべてオープンソース化されています。

英語の原文概要

Mixture-of-experts (MoE) inference on consumer hardware is bounded by weight memory: a 35B-class model is 19.5GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Naive offloading to SSD does not help on its own, because layer N+1's experts must be chosen before layer N's output exists, so the reads cannot start early enough to hide behind compute. We present Edge0, a streaming MoE inference engine that closes the gap with a prerouter: a per-layer head predicts the next layer's routing one token ahead, and the prediction is consumed as the routing itself, so the staged expert set equals the routed set and nothing is dropped. An unmerged recovery LoRA, trained on the student path, pays back the quality lost to int4 quantization and routing replacement. On a single 24GB machine, Edge0 serves a 35B MoE at 20tok/s inside 3GiB of peak active memory, within a few points of its fp16 teacher on average across five public benchmarks. An 8B tier runs on the same framework, and the framework, checkpoints, and adapters are open source.

論文解説

問題と動機

Mixture-of-Experts（MoE）モデルは、トークンあたりのアクティブパラメータ数を低く抑えつつ（例：35Bモデルで1トークンあたり約3Bをアクティブ化）、総パラメータ数をスケールさせます。しかし、民生用ハードウェアでの推論は依然として静的な重みメモリによって制限されます。35B MoEはint4に量子化しても約19.5 GBを必要とします。24 GBのマシンでは、これがOSやKVキャッシュの領域を完全に圧迫し、深刻なメモリページングとデコードスループットの低下（例：M4 Pro Mac mini上の常駐モデルで3.9 tok/s）を引き起こします。

エキスパートの重みをSSDにオフロードすればメモリ使用量は削減されますが、単純なオンデマンド読み込みではハードウェア的な直列化のボトルネックが生じます。レイヤー $N + 1$ のエキスパート選択はレイヤー $N$ の出力に依存するためです。SSDの読み出しを計算の裏に隠せるほど早期に開始できず、トークンごと・レイヤーごとにディスク遅延による実行ストールが発生します。

主なアイデア

重み層としてのSSD: MoEエキスパートの重みはint4形式でSSD上に保持され、コンパクトなストリーミングプールにメモリマップ（mmap）されます。ピークメモリは、モデル全体のパラメータサイズではなく、アクティブなワーキングセットによって決まります。
ルーティングとしての予測（Prediction as Routing）: レイヤー $N$ の完了を待つ代わりに、各レイヤーのprerouterがレイヤー $N + 1$ のルーティングを丸1トークン先回りして予測します。極めて重要なのは、デコード時にモデルがこの予測をルーティングそのものとして扱う点です。ステージングされたセットとルーティングされたセットが構造上完全に一致するため、トークンのドロップやフォールバック読み込みは発生しません。
非マージの復元LoRA（Unmerged Recovery LoRA）: 4-bit量子化とおおよそのルーティングにより品質の低下が生じます。Edge0は蒸留を介してスチューデントのルーティング経路上で低ランクアダプター（LoRA）を学習します。LoRAをint4ベースモデルにマージして再量子化すると差分の大部分が消失するため（ロジットレベルの保持率は18%のみ）、Edge0はわずかなオーバーヘッドでアダプターを並列の非マージ差分として提供します。

手法

Prerouterヘッドのアーキテクチャ

レイヤー $N$ が保持する各prerouterヘッドは、アテンション後のLayerNorm出力に対してトークン $t$ で実行され、トークン $t + 1$ におけるレイヤー $N + 1$ のルーティングを予測します。入力特徴量は、隠れ状態と2つのtop- $k$ ワンホットルーティングベクトル（現在および前のトークン）を結合したものです:

$dim = d_{model} + 2 E$

35Bクラスの場合、これは $2048 + 2 \times 256 = 2560$ となります。ヘッドは以下を計算します:

$head (x) = fc_{2} (erf - gelu (fc_{1} (x))) + ℓ (x)$

ここで、 $ℓ (x)$ は次のレイヤーのルーターの重みからウォームスタートされた線形残差パスです。

非マージアダプターの実行

提供されるMoEレイヤーは、量子化されたベースバイトを維持し、並列LoRA差分を実行します:

$y = W_{int4} (x) + \frac{α}{r} B A x$

3段階の学習

ヘッドの蒸留: ネットワークの他の部分を固定したまま、教師モデルの次のレイヤーにおける真のルーター決定を模倣するようにprerouterヘッドを学習します。
スチューデント経路のSFT: ルーティングをprerouter経由に強制しながら、教師が生成したテキストを使用して復元LoRAアダプター（アテンション、線形アテンション、共有エキスパートの射影層に接続し、ストリーミングエキスパートには接続しない）を学習します。
オンポリシー蒸留: 事前ルーティングされた実行環境下でスチューデントがテキストを生成し、fp16ベースの教師モデルがその生成をスコアリングします。そして、裾野を考慮した（tail-aware）top- $k$ 確率目的関数を用いた逆Kullback–LeiblerダイバージェンスによりLoRAパラメータが最適化されます。

結果

スループットとメモリ（Mac mini M4 Pro, 24 GB）:
- edge0-35b（ $K = 4$ ）: 19.5 GBのint4チェックポイントからわずか2.9 GiBのピークアクティブメモリを使用して、20.4 tok/sのデコードを達成（全量常駐時の3.9 tok/s、18.2 GiBとの比較）。
- edge0-8b（ $K = 8$ ）: 1.5 GiBのピークメモリ内で28.0 tok/sのデコードを達成。
プリフェッチによる高速化（MacBook M2, 16 GB）:
- 単なるリアクティブなオンデマンドストリーミングと比較して、1トークンのprerouterプリフェッチはルーティング幅全体でデコード速度を**+80%〜+84%**向上（ $K = 2$ : 4.8 → 8.6 tok/s、 $K = 4$ : 3.5 → 6.4 tok/s、 $K = 8$ : 1.8 → 3.3 tok/s）。
- メインスレッドのブロックされたロード時間は大幅に減少（例: $K = 4$ でステップあたり244.0 msから101.9 msへ）。
モデル品質（OpenCompass）:
- 5つのベンチマーク（AIME 2026、HumanEval、GPQA-Diamond、MMLU-Pro、IFBench）全体で、edge0-35bの平均スコアは79.2であり、fp16のQwen3.6ベースの83.2に対して3.9ポイント差にとどまります。
- edge0-8bは69.9を達成し、Ling 3.0 tiny fp16ベースの72.7に対して2.8ポイント差です。

限界と未解決の課題

サービングの同時実行性: Edge0は単一リクエストに対してFIFO直列化で動作します。バッチ処理を行うとエキスパートのワーキングセットとキャッシュのダイナミクスが変化しますが、本稿の単一流プロファイルでは対応していません。
CPUグラフ構築のオーバーヘッド: デコード速度はCPU側のフレームワークオーバーヘッド（40レイヤーにわたるMLXグラフ構築でステップあたり約44 ms）によって制限されており、これはSSDの最適化では緩和できません。
推論タスクにおける品質低下: 性能低下は複雑な推論ベンチマークに集中しています（例: 35BではAIMEで6.1ポイント低下、8Bでは10.0ポイント低下）。
ハードウェア依存性: prerouterの相対的なスループット向上はディスク遅延に依存し、重みキャッシュが温まっている場合や超高速ストレージ層で実行する場合には効果が縮小します。

重要用語

Edge0: ルーティング予測によって読み込みと計算をオーバーラップさせ、非マージの復元LoRAで品質損失を補償することにより、SSDからモデルを提供するストリーミングMoE推論エンジン。
prerouter: レイヤーNの隠れ状態と最近のルーティング履歴を使用して、レイヤーN+1のルーティング決定を1トークン先回りして予測する軽量なレイヤーごとのニューラルヘッド。
prediction-as-routing: デコード時にprerouterの予測したエキスパートセットを実際のレイヤールーティングとして直接実行し、実行時のトークンドロップやミスペナルティを排除する設計原理。
recovery LoRA: スチューデントの事前ルーティング経路上で蒸留によって学習され、int4量子化とルーティング近似の誤差を補償するために非マージの並列ブランチとして実行される低ランクアダプター。
staged decode: ステップごとのメモリ確保を回避するために、GPU上のスロットテーブルを介してエキスパートインデックスがマッピングされる、固定スロットダブルバッファリングを用いたエキスパート実行モード。
incremental sticky-slot stack (incr_stack): すべてのエキスパートテンソルを再構築するのではなく、ステップごとに変更されたエキスパートスロットのみを書き換えることで、永続的なスタックドエキスパートテンソルをインプレースで更新するメカニズム。
unmerged delta: LoRAアダプターを重みに融合させるのではなく、量子化されたベース重みと並列に実行することで、再量子化時の微細な適応情報の損失を防ぐ手法。
memory wall: プロセッサのスループットとDRAM帯域幅・容量との間の格差。ローカルLLMサービングにおいては、動的なKVキャッシュ状態と静的なモデル重みストレージの2つに分かれます。

よくある質問

なぜEdge0は予測を単なるヒントとして扱うのではなく、ルーティングそのものとして利用するのですか？

予測をヒントとして扱うと、予測が外れた場合に実行時フォールバックやトークンのドロップが必要になります。予測をルーティングとして直接利用することで、ステージングされたエキスパートセットが常にルーティング対象と一致するようになり、近似に伴うコストをオフライン学習（蒸留による品質回復）へと移行できます。

なぜ復元LoRAは4-bitの重みに融合されず、非マージのまま提供されるのですか？

LoRAの重み差分のRMSは約10^-3であり、これは4-bit量子化のステップサイズよりも小さいためです。マージして再量子化するとアダプターの効果の大部分が失われ、ロジットレベルの保持率は18%に留まりますが、並列の非マージ差分として実行すれば100%の忠実度が維持されます。

なぜEdge0は同一トークン内ではなく、1トークン先を予測するのですか？

同一トークン内で予測を行うと、ステップごとに30〜100 msを要するレイヤーごとのGPUパイプラインフラッシュが必要になり、隠蔽しようとしているロード時間よりも長くなってしまいます。丸1トークン先行させることで、ステージングされるすべてのレイヤーを単一のバッチフラッシュで予測可能になり、SSDのI/Oを隠蔽するための計算時間をデコード丸1ステップ分確保できます。

24 GBのマシンで、Edge0が全量メモリ常駐のベースラインよりも高速にデコードできるのはなぜですか？

19.5 GBのモデルを完全に常駐させると18.2 GiBのRAMを消費し、OSやKVキャッシュ用の空き領域がほとんど残らないため、深刻なOSメモリページングが発生してデコード速度が3.9 tok/sに低下します。Edge0はエキスパートを約2.9 GiBのピークアクティブメモリ内でSSDからストリーミングするため、OSのメモリ圧迫を完全に回避し、20.4 tok/sで動作します。

このページの解説は AI が論文から作成したもので、誤りを含む場合があります。詳細は原文をご確認ください。

XYZNOT で読む（全文翻訳と AI 解説）