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
cs.AI
2026-09-16
Edge0用逐层prerouter提前一个token预测专家路由。35B MoE在24GB机器上跑到20.4 tok/s、峰值2.9GiB,五项公开基准平均落后fp16教师3.9分,比常驻权重快五倍。
MoE 推理在消费级机器上卡的是权重体积,不是计算。35B 级模型 int4 仍有 19.5 GB。稀疏只缩小每个 token 算多少,不缩小必须放在某处的字节。量化到 4 bit 已经是能用的底线,再低误差会把模型毁掉。
把专家卸到 SSD 也救不了 decode:第 N+1 层选哪些专家,要等第 N 层输出出来,磁盘读藏不到计算后面。Mac mini M4 Pro 24 GB 上,整份 4-bit 权重常驻的 mlx-lm 只跑到 3.9 tok/s,还占 18.2 GiB,操作系统和 KV 几乎没地方。
Edge0 把专家权重留在 SSD,按需 mmap 进有界流式池。核心有三处。
预测即路由。 第 N 层的小头用当前 token 的 post-attention 隐状态,预测第 N+1 层在下一个 token 上的路由。decode 时 MoE 层直接吃这份预测,不再跑原 router。预取集合和真正激活集合按构造相等,没有 drop,也没有 fallback 加载。提前量必须是整 token:同一 token 内、attention 后立刻选下一层,会在他们的引擎里每步同步 30–100 ms,比能藏住的加载还贵。
头是 fc1 → erf-gelu → fc2,加一条从下一层 router 权重热启动的线性残差。输入拼上本层本步和上一步的 top-k one-hot。35B 档 33 个头(owner 6–38),32/40 层走预测。
不合并的 recovery LoRA。 int4 基座冻结,低秩适配器在学生路径上蒸馏 fp16 教师,同时补偿量化和路由替换。合并后再量化回 4-bit 会把适配器效果抹掉:注意力投影只剩 34%,稠密投影剩 2%,logits 层只剩 18%。并行 delta 只要 42 MB,decode 时间测不出差别。
训练三阶段,基座始终冻在 4-bit 反量化的 bf16 上:先蒸馏路由头,再在学生路径上对约 200 万行教师文本做 SFT,最后用学生自己的生成做 on-policy 蒸馏(reverse KL)。只蒸头、不做 SFT 的检查点会重复崩溃。
35B 发布档把原模型的 K=8 收到 K=4,换适配器文件即可,不必重训基座。
吞吐在 Mac mini M4 Pro 24 GB 上测。质量用 OpenCompass,和 fp16 基座同设置对照。
| 方法 | 吞吐 | 峰值内存 |
| mlx-lm 常驻 4-bit | 3.9 tok/s | 18.2 GiB |
| Edge0 35B K=4 | 20.4 tok/s | 2.9 GiB |
| Edge0 8B | 28 tok/s | 1.5 GiB |
24 GB 这台机器上关掉 prerouter 仍有 19.9 tok/s,页缓存已经能把热专家留下。真正拉开的是权重放不下的 16 GB MacBook M2:相对纯按需加载,K=2/4/8 分别从 4.8、3.5、1.8 tok/s 提到 8.6、6.4、3.3 tok/s,增益 +80%、+82%、+84%。K=4 时主线程堵在专家加载上的时间从 244.0 ms 降到 101.9 ms。
| 基准 | edge0-35b | Qwen3.6 fp16 | edge0-8b | Ling 3.0 tiny fp16 |
| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
| 平均 | 79.2 | 83.2 | 69.9 | 72.7 |
逐项平均差距 35B 为 3.9 分、8B 为 2.8 分。8B 的 MMLU-Pro 比教师高 4.3 分。
消费级机器跑 35B MoE 的常规答案是买内存或上云。Edge0 给出第三条:专家放 SSD,用训练好的提前路由把磁盘延迟藏进 decode,质量账单用一份不合并的 LoRA 付。框架、检查点、适配器都开源。
这是 Apple Silicon 上单机、单请求的方案。MLX 是唯一实现,CUDA 槽位还停在架构层。宽度 K 同时动速度和内存:同一缓存预算下 K=8 收到 K=4,decode 从 3.3 提到 6.4 tok/s。
页缓存热、盘快时,别指望 prerouter 再给 80%。24 GB 机器上它几乎不加吞吐,它赚钱的场景是权重根本塞不进内存。
单请求 FIFO,没有批处理。批会改变专家工作集,现有 profile 覆盖不了。decode 的地板在 CPU 构图,40 层前向每步 44 ms,优化存储搬不走这块。相邻 token 只共享约四分之一专家,预取经常读一次付两次。
质量缺口集中在长链推理:AIME 上 35B 掉 6.1 分,8B 掉 10.0 分。其余 8B 基准都在 6.7 分以内。单次基准在这套硬件上有 ±40% 的 run-to-run 波动,跨 session 能到 2.3×,所以论文只用同 session 轮换 A/B 的中位数。