Edge0 项目推荐：把 35B MoE 卸到 SSD，用"提前一个 token 猜路由"把读盘时间藏起来

Edge0（Edge0-AI/Edge0）是一个用 Python + MLX 写的流式 MoE 推理框架。它想做的事一句话讲完：让 35B 级的稀疏 MoE 模型在消费级机器上跑起来，而且峰值内存只由"当前激活的专家"决定，而不是参数量。

它的技术报告标题把这道题说得很准——《The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction》。所谓"memory wall 的另一半"，指的是大家习惯只算容量那一半（权重要占多少内存），而忽略了延迟那一半：把权重卸到 SSD 上，“装不下"确实解决了，可是每一层都要等上一层的输出算完才知道该读哪些专家——读盘时间根本来不及藏在计算后面。

Edge0 的答案是：训练一个小的路由预测头（prerouter），在 token t 就提前一个 token 预测第 N+1 层的路由，于是专家加载与当前 step 的计算重叠；而且这个预测直接当作路由本身使用，所以不存在"猜错了要补"的问题。

关键数据概览

有两个数字比 star 更值得看。第一个是 2.9 GiB：23GB 的 checkpoint、40 层、每层 256 个专家，峰值活跃内存不到 3GiB——说明它的内存上限确实由"激活集"而不是"参数量"决定。第二个是那句写在 README 里的量化边界：int4 + 训好的 LoRA 适配器 + prerouter 路由，在五个基准上平均只掉 3.9 分，其中 MMLU-Pro 还高于 fp16 原模型。（AIME 2026 从 92.7 掉到 86.6，是掉得最狠的一项，README 如实列了出来。）

起源故事：MoE 让"显存墙"变成了两堵墙

先讲清楚它踩中的那个技术转折。

MoE 之前，本地推理的路线是"量化”：7B 压到 4bit 就 3.5GB，皆大欢喜。MoE 之后，模型看起来很大、实际算得很少，于是产生一个诱惑——既然每个 token 只用一小撮专家，那把专家放到 SSD 上按需读就行了。

但这条路一上手就会撞墙。MoE 的每一层是这样的：上一层的输出 → 路由（router）→ 选出 top-k 专家 → 用这些专家算。也就是说，第 N+1 层的专家清单，必须等第 N 层算完才能知道；而知道清单之后，还要从 SSD 把这几百 MB 读进来（这期间 GPU 在等），读完了才算。 读盘延迟被完全暴露在关键路径上，一步都不能省。

这就是论文说的"另一半"：只要"先路由、后加载"这个顺序不变，SSD offload 就只是把一次性的容量问题换成了每步都付一次的延迟问题。已有的 offload 方案（llama.cpp 的 --n-cpu-moe、各家的 expert offloading）大多停在这里：能跑，但对存储延迟极其敏感。

Edge0 的破题点是把顺序反过来：既然专家清单要等上一层算完才知道，那就提前猜。用小网络在 token t 预测 token t+1 时第 N+1 层会选谁，让读取与计算并行起来。

技术亮点：三处"必须做对否则整个方案不成立"的取舍

亮点一：prerouter——用"双 shift"把预测变成路由本身

这是整个项目最核心、也最巧的一处。先看问题：预测本身不可能 100% 准，如果预测只是"预读提示"，那猜错的部分照样要在关键路径上补读，收益就被吃掉了。Edge0 的解法是接受预测、消灭"猜错"这个概念：

• 层内 shift：第 N 层自己拥有的预测头，输入是第 N 层的 MoE 输入（post-attention norm 的输出），输出预测的是第 N+1 层的路由；
• token 内 shift：第 N+1 层消费的，是第 N 层的头在上一个 token t-1 时给出的预测。

两个 shift 叠起来的效果是：这一步要 staged 的专家集，就正好是这个预测本身。 于是 staged_replace 语义下"路由集 == staged 集"，零丢弃、零补读——预测不是提示，它就是路由。

docs/prerouter.md 把 shift 关系画得很清楚，实现落在 python/src/edge0/prerouter/ 五个文件里（heads.py 数学、install.py 注入、stager.py 跨 token 提交、state.py 状态）。头的结构出人意料地小（伪代码，来自 heads.py）：

# 第 N 层自己的头：输入 = [隐状态, 本 token 实际 top-k 的 one-hot, 上一 token 实际 top-k 的 one-hot]
feats = concat([h, executed_topk_onehot, prev_token_topk_onehot])   # 维度 512 + 2E

# fc1 -> 精确 erf 形式的 GELU -> fc2，再叠一条 linear_init 的直连
logits = linear_init(feats) + fc2(gelu_erf(fc1(feats)))             # -> num_experts

# 必须与导出权重逐位一致；训练精度与生产精度都是 fp16

几个能看出"不是玩具"的细节：

• 头的输入里显式带上本 token 与上一 token 实际路由到的专家（feature_topk="executed"，即真正被路由到的 top-k，而不是用原始 gate 重算的"teacher"语义）。预测一个带记忆的路由过程，历史选择是最强的特征；
• gelu_erf 特意实现成精确 erf 版本去对齐 torch 的 F.gelu 默认行为，heads.py 的文档字符串写着"bit-identical to the training export"——因为权重是训练完直接搬过来的，一点数值偏差都会累积；
• 预测头的规模：35b 用 33 个头（第 6~38 层），hidden 512、fp16；8b 用 16 个头。全部打包成 prerouter_edge0_35b.safetensors 与模型放在同一个目录里。

收益量级：解码吞吐最多 +59%，README 写明了增益随"存储延迟 / 模型大小 / 路由宽度 K"增长——这正好符合它的原理：读得越慢、专家越大，重叠省下来的就越多。

亮点二：SSD 流式专家层——四条执行路径 + “索引永不离开 GPU”

prerouter 决定"什么时候读"，streaming/layer.py 里的 StreamingSwitchGLU 决定"怎么读得快"。它按"快慢档位"提供了四条路径，用一套 typed 配置（LayerOptions）切换，而不是散落的 env 开关：

staged 这条路是整个设计的重心，也最能说明作者对"每步一次主机同步"有多过敏：

• 路由索引通过 mx.take 查槽表拿到数据，per-layer per-step 零 host sync——如果有一次 indices.tolist() 之类的操作，GPU 就要停下来等 CPU，解码速度立刻掉一个档；
• asm_cache 把"槽表 + 已堆叠的 wargs"按专家集合缓存，重复集合不再重建计算图节点；
• incr_stack 用 put_along_axis 的原地行写替换掉每步 9 次 mx.stack（持久张量），配合 incr_writeback 把被换出的 bundle 写回共享 LRU；
• 缺专家映射到"全零溢出槽"，贡献被丢弃——但在生产 profile 里由构造决定不会缺，因为 staged 集就是 prerouter 的预测；
• 共享缓存是跨层的一个全局 LRU 预算（SharedExpertCache，key 是 (layer_idx, expert)）。LayerOptions 的注释把这件事讲得很直白：解码一步会摸到每层 K 个专家（8b 是 23 × 8 = 184 个不同 key），所以预算必须按整个模型来配——仓库里 64 个槽的默认值实测命中率恰好是 0，连一个 token 的工作集都放不下。这种"把实测结论写进注释"的地方，在热门项目里并不多见。

第三块是 warm_willneed：在构建 bundle 之前，对将要访问的专家字节区间发一次 madvise(MADV_WILLNEED)，让内核做批量异步预读，而不是让加载器一页一页 demand-fault。注释给出了原因和实测——在内存吃紧的机器上，缺页会串行化在 VM map lock 上：

staged_sync + incr_stack + incr_writeback：峰值 MLX 2.92 → 2.60 GiB，
   major/minor 缺页 5162 → 1925 次/step，CPU 143 → 109 ms/step
warm_willneed：staging 墙钟 162 → 27 ms/step
两者叠加：244.25 → 187.58 ms/step，4.09 → 5.33 tok/s（+31.0%，配对 4/5）；
  prerouter 相对关闭时的增益也从 +8.0% 升到 +37.8%（配对 5/5），输出逐位一致

这组数字（含机型、checkpoint、回放步数、配对轮次）直接写在 options.py 的 docstring 里，是"工程结论跟着代码走"的写法。

亮点三：Recover-LoRA——用"不合并的适配器"把 int4 的损失买回来

量化到 4bit 一定会掉分，Edge0 的补偿方式不是后训练量化（PTQ）调参，而是蒸馏：冻结 int4 的基座，用 FP 教师模型去训 LoRA 适配器，把量化损失补回来。这件事在这个框架里有几个明确的工程后果：

• 适配器永不合并。adapters/lora.py 走的是"平行 delta 路径"，基座只读，一份基座可以服务多套适配器；升级适配器只换 .safetensors 文件，不动基座；
• 适配器格式统一成 safetensors + 溯源元数据（lora_edge0_35b.safetensors / prerouter_edge0_35b.safetensors），和基座 checkpoint 放在同一个模型目录里，AutoEngine.from_pretrained(<dir>) 直接从 config.json 认出厂牌并自动加载；
• 没有适配器也能跑：--no-lora / --no-prerouter 退化成"朴素 int4 基座"，方便做消融对照；
• 顺带一个舒服的用法——模型目录自包含：一次 huggingface-cli download 下来就是"基座 + LoRA + prerouter 头"三者齐备、开箱可跑的目录，不用手工拼装。

质量表（OpenCompass 同一套参数）是这样：

另外，框架本身有几个"像是被人用过"的设计：后端隔离——所有 MLX 代码锁在 backends/mlx/ 里，核心逻辑（模型 spec、prerouter、流式专家池、server）只依赖 facade（backends/base.py），backends/cuda/ 是预留槽位，换后端不改核心；transformers 风格 API——AutoConfig / AutoModel / AutoEngine 按模型名解析 tier；MoE 抽象——MoESpec 描述路由族（SOFTMAX_TOPK / SIGMOID_GROUP）、量化布局与权重排布，所以同一个 StreamingSwitchGLU 类服务所有模型，加新模型只需补一份 spec。

优点与要注意的地方

优点：

• 一套机制解决两堵墙：SSD mmap 解决容量，prerouter 解决延迟，而且"预测即路由"让零丢弃成为构造性质而非工程努力；
• 工程数据是带机型和方法论的：速度表标了机器（Mac mini M4 Pro / 24GB）、冷热态定义、prompt 长度、采样步数；内存与缺页的优化都是"改前改后 + 配对轮次"；
• 质量与代价都摊开写：五个基准的逐项分数（含掉分最多的 AIME）、int4 与 fp16 的对比、以及 --no-lora/--no-prerouter 的消融开关；
• 适配器不合并 + 一份基座多套适配器：升级适配器不动基座，safetensors 带溯源元数据；
• 接口标准：本地 edge0 serve 提供 OpenAI 兼容的 /v1/chat/completions，并已合入 tool/function calling 的完整往返，任何能改 base URL 的客户端都能接；
• 一份代码五个平台：python/（框架）、macos/（Rust CLI/守护/桌面）、ios/（Swift + MLX Swift）、android/（Kotlin）、windows/（C++ + Vulkan）都在同一个仓库里开源，而"统一访问层"排在 Q4 的路线图上；
• 文档是给实现者看的：docs/streaming.md 连"bundle 顺序必须与 _swiglu(up, gate) 的参数位置一一对应，顺序错了相对 L2 从 0.2% 跳到 100%“这种坑都写进了文档，还配了 tests/test_streaming_math.py 逐条对齐 exact 路径。

要注意的地方：

• 硬门槛是 macOS + Apple Silicon：Python 框架这一路要求 Python ≥ 3.10 与 MLX（mlx==0.30.6 / mlx-lm==0.31.0），CUDA 后端目前还只是 backends/cuda/ 的预留槽，排在 Q4 路线图；
• 磁盘和下载量不小：35b 的 4bit checkpoint 约 23GB，8b 约 4.2GB；专家权重靠 mmap 按需读，不会预先载入内存，但磁盘读取性能直接决定体验；
• 速度不是"丝滑"级别：M4 Pro 上 15~18 tok/s（冷 prefill 113 tok/s，热 prefill 140 tok/s）——它是"能跑"而不是"跑得快”，追求吞吐的场景不合适；
• 平台引擎与统一框架之间还有一段距离：iOS/macOS/Android/Windows 是平台原生栈各自实现，README 明确说统一推理框架在 Q4 2026 交付；在此之前"一套 API 通吃所有平台"还不成立；
• 本地缓存的容量要按整个模型算：默认 cache_slots=64 实测命中率为 0，生产 profile 另有取值——自己调参时别把它当"每层槽数"；
• 没有 git tag / release 语义：以 main 为发布面，跟着升级要接受默认值会变（LayerOptions 的预设仍在演进）；
• 预览版模型：两个 tier 都带 preview 后缀，训练数据与适配器溯源在各自模型页上，用于关键业务前建议按自己的任务实测。

适合谁

• 手上有 Apple Silicon、想本地跑"真·35B"的人：24GB 统一内存就能把 35B MoE 跑起来，峰值内存不到 3GiB，代价是速度在十几 tok/s 量级；
• 在做端侧/离线推理、被"显存墙"卡住的人：SSD offload + 路由预测 + 全局 LRU 这一整套，是可以直接搬到自家场景的参考架构；
• 想把本地模型接进现有工具链的人：OpenAI 兼容接口 + 工具调用往返，Claude Code / 各类客户端改个 base URL 就能指过去；
• 研究或复现"内存墙"问题的人：有论文、有可复现的 bench 脚本（python/examples/bench.py）、有逐项的消融开关（--no-lora、--no-prerouter、各种 LayerOptions），是比较少见的"论文 + 可跑代码 + 实测数据"三件套；
• 想学 MoE 推理系统怎么写的工程师：moe/spec.py 的抽象、streaming/layer.py 的四条路径、prerouter/ 的双 shift 提交、后端 facade 的切法，每一处都能当教材读。

一句话总结：Edge0 的价值不在于"把大模型塞进小机器"这句口号——MoE + SSD offload 谁都会说——而在于它正面解决了这条路上最容易被略过的那个顺序问题：第 N+1 层的专家清单必须先于第 N 层的结果存在。它的回答是训练一个头去猜，并且让猜的结果直接成为路由，于是"猜错"这个失效模式在构造上就不存在。再加上一份把每项优化都标上机型和配对轮次的实测记录，这是那种"方法听起来简单、但每一步都看得出真在做工程"的项目。

项目地址：https://github.com/Edge0-AI/Edge0 技术报告：《The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction》https://arxiv.org/abs/2609.18063 模型权重：https://huggingface.co/Edge0/Edge0-35B-A3B-preview ｜ https://huggingface.co/Edge0/Edge0-8B-A1B-preview（含 ModelScope 镜像） 快速上手：cd python && pip install -e '.[dev,fetch]' → edge0 demo edge0-35b / edge0 serve edge0-8b