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 講解）

三句話總結

🎯 35B 級別的 MoE 即使採用 int4 也需要約 19.5 GB，超過消費級 RAM 容量；而傳統 SSD 卸載因層路由依賴前一層輸出，會卡在磁碟延遲上。
🔧 Edge0 透過每層的 prerouter 提前一個 token 預測路由，落實「以預測為路由」，並結合在學生路徑上蒸餾的未合併平行修復 LoRA，實現從 SSD 串流 MoE 專家。
📈 在 24 GB 的 Mac mini M4 Pro 上，僅需 2.9 GiB 峰值記憶體即可達到約 20.4 tok/s 的 35B MoE 生成速度，基準測試平均表現與 fp16 教師模型僅差 3.9 分。

摘要

在消費級硬體上執行混合專家模型（MoE）推論主要受限於權重記憶體：35B 級別的模型在 4-bit 下為 19.5GB，而稀疏性僅減少了每個 token 的計算量，並未減少必須常駐的資料位元組。單純將權重卸載（offload）至 SSD 並無助益，因為第 N+1 層的專家必須在第 N 層輸出產生前選定，導致讀取操作無法足夠提早啟動以隱藏在計算背後。我們提出了 Edge0，這是一個串流式 MoE 推論引擎，透過預測路由器（prerouter）消除了這一差距：每層的神經網路頭（head）提前一個 token 預測下一層的路由，且該預測直接作為路由決策本身，因此預備的專家集合與實際路由集合完全一致，無需丟棄任何 token。在學生路徑上訓練的未合併修復 LoRA，則彌補了 int4 量化與路由替換所損失的模型品質。在單台 24GB 機器上，Edge0 能在 3GiB 峰值動態記憶體內以 20tok/s 的速度運行 35B MoE，在五個公開基準測試中的平均表現與 fp16 教師模型僅相差數分。8B 級別的模型亦可在相同框架上運行，且本框架、檢查點（checkpoints）與適配器（adapters）均已開源。

英文原文摘要

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.

論文解讀

問題與動機

混合專家模型（MoE）在擴展總參數量時，能保持每個 token 的動態激活參數量在較低水準（例如在 35B 模型中每個 token 僅激活約 3B 參數）。然而，消費級硬體上的推論仍受限於靜態權重記憶體：35B MoE 即使量化至 int4 仍需約 19.5 GB。在 24 GB 的機器上，這會徹底排擠作業系統與 KV 快取的空間，導致嚴重的記憶體分頁（paging）並使解碼吞吐量銳減（例如常駐於 M4 Pro Mac mini 上的模型僅有 3.9 tok/s）。

將專家權重卸載（offload）至 SSD 可以減少記憶體佔用，但單純的隨選載入會引入嚴重的序列化瓶頸：第 $N + 1$ 層的專家選擇取決於第 $N$ 層的輸出。SSD 讀取無法足夠提早啟動以隱藏在計算背後，導致每個 token 在每層都會因磁碟延遲而卡頓。

核心理念

將 SSD 作為權重儲存層：MoE 專家權重以 int4 格式保存在 SSD 上，並透過記憶體映射（mmap）進入小巧的串流池。峰值記憶體取決於動態工作集大小，而非模型的總參數規模。
以預測作為路由（Prediction as Routing）：每層的 prerouter 不再等待第 $N$ 層計算完成，而是提早整整一個 token 預測第 $N + 1$ 層的路由。關鍵在於，模型在解碼時直接將該預測視為路由本身。這在架構上徹底消除了丟棄 token 或回退載入的需求，因為預先載入的集合與路由集合完全相同。
未合併修復 LoRA（Unmerged Recovery LoRA）：4-bit 量化與近似路由會帶來品質損失。Edge0 透過蒸餾在學生路由路徑上訓練了一個低秩適配器。由於將 LoRA 合併至 int4 基底模型並重新量化會抹煞大部分增量（在 logits 層級僅保留 18%），Edge0 將適配器作為平行的未合併增量執行，額外開銷微乎其微。

方法

Prerouter 頭部架構

第 $N$ 層擁有的每個 prerouter 頭部在 token $t$ 時於注意機制後的層正規化（layer norm）輸出上運行，並預測第 $N + 1$ 層在 token $t + 1$ 時的路由。輸入特徵為隱藏狀態與兩個 top- $k$ one-hot 路由向量（當前與前一個 token）的拼接：

$dim = d_{model} + 2 E$

對於 35B 等級，此維度為 $2048 + 2 \times 256 = 2560$ 。該頭部計算：

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

其中 $ℓ (x)$ 為線性殘差路徑，其權重使用下一層路由器的權重進行熱啟動（warm-started）。

未合併適配器執行

推論服務中的 MoE 層保留了量化後的基底位元組，並執行平行的 LoRA 增量：

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

三階段訓練

頭部蒸餾（Head Distillation）：在凍結網路其餘部分的情況下，訓練 prerouter 頭部以模仿教師模型下一層真實路由器的決策。
學生路徑 SFT（Student-Path SFT）：使用教師生成的文本訓練修復 LoRA 適配器（附加於注意力、線性注意力與共享專家投影，但不附加於串流專家），並強制透過 prerouter 進行路由。
策略內蒸餾（On-Policy Distillation）：學生模型在 prerouter 路由執行下生成文本，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 個 token 的 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）：
- 在五個基準測試（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 序列化方式運作；批次處理（batching）會改變專家工作集與快取動態，這是單串流設定所未涵蓋的。
CPU 計算圖構建開銷：解碼速度受限於 CPU 端框架的額外負擔（在 40 層 MLX 圖構建中約每步 44 ms），SSD 的最佳化無法緩解此問題。
推理任務的品質下降：效能劣化的情況主要集中在複雜推理基準測試（例如 35B 在 AIME 下降 6.1 分，8B 下降 10.0 分）。
硬體相依性：Prerouter 帶來的相對吞吐量提升受限於磁碟延遲；當權重快取處於熱狀態（warm）或運行在極高速儲存裝置上時，優勢會減小。

關鍵術語

Edge0: 一種串流式 MoE 推論引擎，透過路由預測將讀取與計算重疊，並利用未合併的修復 LoRA 彌補品質損失，實現直接從 SSD 提供模型推論服務。
prerouter: 一種輕量級的逐層神經網路頭，利用第 N 層的隱藏狀態與近期的路由歷史，提前一個 token 預測第 N+1 層的路由決策。
prediction-as-routing: 一種設計原則，在解碼時直接將 prerouter 預測的專家集合執行為實際的層級路由，消除了運行時 token 丟棄或未命中懲罰。
recovery LoRA: 在學生預先路由路徑上透過蒸餾訓練的低秩適配器，作為未合併的平行分支執行，以補償 int4 量化與路由近似帶來的誤差。
staged decode: 一種使用固定槽位雙緩衝區的專家執行模式，專家索引透過 GPU 上的槽位表進行映射，以避免每步重新分配記憶體。
incremental sticky-slot stack (incr_stack): 一種就地更新持久化堆疊專家張量的機制，每步僅覆寫發生變更的專家槽位，而非重新組裝所有專家張量。
unmerged delta: 將 LoRA 適配器與量化基底權重平行執行，而非直接融合至權重中，避免在重新量化過程中遺失微調適應的細粒度資訊。
memory wall: 處理器吞吐量與 DRAM 頻寬/容量之間的差距；在本地 LLM 服務中，這被拆分為動態 KV 快取狀態與靜態模型權重儲存兩部分。

常見問題

為什麼 Edge0 採用「以預測為路由」，而不是將預測視為提示（hints）？

將預測視為提示需要在預測未命中時進行運行時回退或丟棄 token。直接將預測作為實際路由，可確保預備的專家集合始終與路由集合吻合，將近似代價轉移至離線訓練階段，並藉由蒸餾修復其品質。

為什麼修復 LoRA 是以未合併方式執行，而不是直接融合到 4-bit 權重中？

LoRA 權重增量的 RMS 約為 10^-3，小於 4-bit 量化的步長。若合併並重新量化會抹去適配器的大部分效果，在 logits 層級僅能保留 18%，而平行的未合併增量則能維持 100% 的保真度。

為什麼 Edge0 是提前一個 token 預測，而不是在同一個 token 內預測？

在同一個 token 內預測需要每層進行 GPU 管線排空（pipeline flush），每步耗時 30 到 100 ms，這甚至長於其試圖隱藏的載入時間。領先整整一個 token 可以讓所有預備層在單一批次排空中完成預測，提供一整個解碼步的計算時間來隱藏 SSD I/O。

Edge0 如何在 24 GB 機器上達到比純記憶體基準線更快的解碼速度？

完全常駐的 19.5 GB 模型會佔用 18.2 GiB RAM，幾乎沒有空間留給作業系統與 KV 快取，導致嚴重的 OS 記憶體分頁並使解碼速度驟降至 3.9 tok/s。Edge0 從 SSD 串流專家時僅需約 2.9 GiB 峰值動態記憶體，完全避開了 OS 記憶體壓力，並以 20.4 tok/s 運行。

本頁解讀由 AI 根據論文生成，可能有錯誤；細節請以原文為準。

用 XYZNOT 閱讀（全文翻譯與 AI 講解）