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

三句话总结

🎯 即使采用 int4 量化，35B 级 MoE 仍需约 19.5 GB 内存，超出消费级 RAM 容量；而原生的 SSD 卸载方案因层路由依赖前一层输出而受制于磁盘延迟瓶颈。
🔧 Edge0 使用每层预路由器提前一个 token 预测路由以实现“预测即路由”，从 SSD 流式传输 MoE 专家，并配合在学生路径上蒸馏得到的未合并并行恢复 LoRA。
📈 在配备 24 GB 内存的 Mac mini M4 Pro 上，以 2.9 GiB 峰值内存实现 35B MoE 约 20.4 tok/s 的解码速度，基准测试平均得分与 fp16 教师模型差距仅 3.9 分。

摘要

消费级硬件上的混合专家（MoE）模型推理受限于权重内存：一个 35B 级别的模型在 4-bit 量化下仍需 19.5GB，且稀疏性仅缩减了每 token 的计算量，并未减少必须驻留的字节数。单纯将权重卸载到 SSD 并不能解决问题，因为第 N+1 层的专家必须在第 N 层输出产生前选定，导致读取操作无法足够提前启动以掩盖在计算之后。我们提出了 Edge0，一个通过预路由器（prerouter）弥合这一差距的流式 MoE 推理引擎：每个层的专用预测头提前一个 token 预测下一层的路由，并将该预测直接作为实际路由使用，从而使暂存专家集与路由专家集完全一致且无任何 token 丢弃。在学生路径上训练的未合并恢复 LoRA 则弥补了由于 int4 量化和路由替换带来的质量损失。在单台 24GB 机器上，Edge0 在仅占用 3GiB 峰值活跃内存的情况下，以 20tok/s 的速度运行 35B MoE，在五个公开基准测试中的平均得分与 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.

论文解读

问题与动机

混合专家（MoE）模型在扩大总参数规模的同时，保持每个 token 的激活参数量处于较低水平（例如在 35B 模型中每个 token 仅激活约 3B 参数）。然而，消费级硬件上的推理仍然受到静态权重内存的制约：即使量化为 int4，35B MoE 模型也需要大约 19.5 GB。在 24 GB 的机器上，这会完全挤占操作系统和 KV 缓存的空间，引发严重的内存分页换页并导致解码吞吐量骤降（例如在 M4 Pro Mac mini 上全驻留模型的解码速度仅为 3.9 tok/s）。

将专家权重卸载到 SSD 可以减少内存占用，但原生的按需加载会引入硬性串行化瓶颈：第 $N + 1$ 层的专家选择取决于第 $N$ 层的输出。SSD 的读取操作无法足够早启动以重叠掩盖在计算时间之后，导致每个 token 在每一层都会因磁盘延迟而停顿。

核心思想

SSD 作为权重存储层：MoE 专家权重以 int4 格式保存在 SSD 上，并通过内存映射（mmap）加载到紧凑的流式池中。峰值内存由活跃工作集决定，而非模型总参数量。
预测即路由（Prediction as Routing）：每个层的预路由器不再等待第 $N$ 层完成，而是提前整整一个 token 预测第 $N + 1$ 层的路由。至关重要的是，模型在解码阶段将该预测直接作为实际路由使用。由于暂存集与路由集天然一致，该设计从原理上杜绝了 token 丢弃或回退加载。
未合并恢复 LoRA（Unmerged Recovery LoRA）：4-bit 量化和近似路由会导致质量损失。Edge0 通过蒸馏在学生路由路径上训练低秩适配器。由于将 LoRA 合并到 int4 基底并重新量化会抹去大部分增量（在 logits 层面仅保留 18%），Edge0 将适配器作为平行的未合并增量运行，且开销几乎可忽略不计。

方法细节

预路由器头结构（Prerouter Head Architecture）

属于第 $N$ 层的每个预路由器头在 token $t$ 时运行于注意力层归一化（post-attention 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)$ 是从下一层路由器的权重热启动的线性残差路径。

未合并适配器执行

推理运行的 MoE 层保留了量化后的基底字节，并运行平行的 LoRA 增量：

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

三阶段训练

预测头蒸馏：保持网络其余部分冻结，训练预路由器头以模仿教师模型下一层的真实路由器决策。
学生路径 SFT：利用教师生成的文本训练恢复 LoRA 适配器（挂载在注意力、线性注意力和共享专家投影上，但不包括流式专家），同时强制通过预路由器进行路由。
策略内蒸馏（On-Policy Distillation）：学生模型在预路由执行下生成文本，fp16 基底教师模型对生成内容进行评分，并通过带有尾部感知 top- $k$ 概率目标的反向 KL 散度（reverse Kullback–Leibler divergence）优化 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 的预路由器预取在不同路由宽度下将解码速度提升了 +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 串行处理；多并发批处理会改变专家工作集和缓存动态，单流性能分析尚未涵盖这部分情况。
CPU 计算图构建开销：解码速度受限于 CPU 侧的框架开销（在 40 层跨度上 MLX 图构建耗时约 44 ms/步），SSD 优化无法缓解这一问题。
推理能力下降：质量衰减主要集中在复杂推理基准上（例如 35B 在 AIME 上下降 6.1 分，8B 下降 10.0 分）。
硬件依赖性：预路由器的相对吞吐收益取决于磁盘延迟，当权重缓存命中较高或在超高速存储层上运行时，收益会有所收窄。

关键术语

Edge0: 一种流式 MoE 推理引擎，通过路由预测将数据读取与计算相重叠以实现从 SSD 服务模型，并利用未合并的恢复 LoRA 补偿模型质量损失。
prerouter: 一种轻量级的逐层神经预测头，利用第 N 层的隐藏状态和近期路由历史，提前一个 token 预测第 N+1 层的路由决策。
prediction-as-routing: 一种设计理念，在解码时将预路由器预测的专家集直接作为实际的层路由执行，从而消除运行时的 token 丢弃或未命中惩罚。
recovery LoRA: 在学生预路由路径上通过蒸馏训练的低秩适配器，作为未合并的平行分支执行，用于补偿 int4 量化和路由近似带来的误差。
staged decode: 一种基于固定槽位双缓冲区的专家执行模式，通过 GPU 上的槽位映射表映射专家索引，避免每步重复分配内存。
incremental sticky-slot stack (incr_stack): 一种就地更新持久化堆叠专家张量的机制，每一步仅重写发生变动的专家槽位，而非重新组装所有专家张量。
unmerged delta: 将 LoRA 适配器与量化基底权重并行执行，而非直接融合到权重中，以防止重新量化时丢失细粒度的微调增量。
memory wall: 处理器吞吐量与 DRAM 带宽/容量之间的差距。在本地 LLM 推理场景中，该矛盾被拆解为动态 KV 缓存状态与静态模型权重存储两大挑战。

常见问题

为什么 Edge0 采用“预测即路由”，而不是将预测作为提示候选？

将预测仅作为提示候选需要运行时回退机制，或者在预测失误时丢弃 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 延迟。

在 24 GB 机器上，Edge0 为何能比全内存驻留基线解码更快？

全内存驻留的 19.5 GB 模型需要占用 18.2 GiB RAM，几乎没有为操作系统和 KV 缓存留下空间，导致严重的系统内存分页换页，解码速度骤降至 3.9 tok/s。Edge0 在约 2.9 GiB 的峰值活跃内存内从 SSD 流式加载专家，完全避免了操作系统的内存压力，运行速度达到 20.4 tok/s。

本页解读由 AI 根据论文生成，可能有错误；细节请以原文为准。

用 XYZNOT 阅读（全文翻译与 AI 讲解）