内存墙的另一半：使用训练路由预测从SSD服务35B MoE模型

# 内存墙的另一半：通过训练路由预测从SSD服务350亿参数MoE模型
来源：https://arxiv.org/html/2609.18063
Yiming Wang††thanks:yiming\.wang@autoark\.aiRunyuan Cai††thanks:runyuan\.cai@autoark\.aiHanze Liu††thanks:hanze\.liu@autoark\.aiXiaodong Zeng††thanks:xiaodong\.zeng@autoark\.ai所属机构：AutoArk

2026年9月

###### 摘要

在消费级硬件上进行混合专家模型推理受限于权重内存：一个350亿参数级别的模型以4位量化存储大小为19.5GB，而稀疏性仅减少了每个令牌的计算量，并未减少需要存储的字节数。单纯将权重卸载到SSD本身并无帮助，因为第N+1层专家的选择依赖于第N层输出，在输出生成前无法提前启动读取以隐藏在计算之后。本文提出Edge0，一种流式MoE推理引擎，通过*预路由器*弥合了这一差距：每层包含一个预测头，提前一个令牌预测下一层的路由决策，并将该预测直接作为路由结果使用，因此预取的专家集合与实际路由集合完全一致，不会丢弃任何专家。我们训练了一个未合并的*恢复LoRA*模块，用于补偿int4量化和路由替换带来的质量损失。在单台24GB内存机器上，Edge0能在3GiB峰值活跃内存内以约20令牌/秒的速度服务35B MoE模型，其性能在五个公开基准测试中平均仅比fp16教师模型低几个百分点。同一框架还支持8B规模模型，且框架、检查点和适配器均已开源。

## 1 引言

1995年，Wulf和McKee提出了“内存墙”概念：处理器发展速度超过DRAM，导致机器在等待内存上花费的比例越来越大[32]。三十年后，内存墙面临AI推理挑战，其算术规则对数据中心外的大型模型部署并不友好：解码阶段每读取一权重字节仅执行少量浮点运算，而硬件“脊点”要求高出几个数量级；关键瓶颈在于字节数，而数据存储在内存中。

内存中存储着行为迥异的两类数据：*状态*（KV缓存）是动态的，随生成令牌增长；*权重*是静态的，自训练结束即固定。业界着力解决动态部分：多头潜在注意力将KV状态压缩了一个数量级[5]，稀疏注意力将其限制在固定窗口[13]。而无法协商减少的部分——数十GB权重——的解决方案是将其置于数据中心，通过专家并行和分片服务来承载[25]。

在350亿参数规模下，两种标准的本地解决方案均失效。量化在4位触及下限；低于4位误差急剧增长直至模型不可用，且无后期技术能修复[11,21]。MoE稀疏性在另一维度发挥作用：35B MoE每个令牌仅激活约30亿参数，减少了*计算量*而非*存储量*。19.5GB仍需存储于某处，而在与操作系统共享的24GB桌面机上，该空间占满整个内存，挤占了系统及用户其他程序的运行空间。

Edge0改变了权重存储位置。专家权重保留在SSD上，通过mmap按需流式加载到内存，峰值内存受限于活跃集合而非参数总量。单纯按需加载速度不足：在每个解码步骤中，第N+1层的专家选择依赖第N层输出，朴素流式引擎会在每层每令牌因磁盘延迟停顿一次。Edge0通过训练的*预路由器*消除停顿：每层拥有一个小型预测头，基于第N层在当前令牌的状态预测第N+1层在下一令牌的路由，使专家读取与前向传播重叠。该预测直接作为路由使用：预取专家集合与路由集合在设计上一致，其他预门控方案在推理时通过回退加载和丢弃令牌所承受的近似误差，在Edge0中通过一次性训练补偿并恢复。

近似路由和4位量化均会损失质量。Edge0在此基础上训练*恢复LoRA*：冻结int4基础模型，从fp16教师模型中在学生路由路径下蒸馏出低秩适配器，并以*未合并*的并行增量形式服务。合并适配器至int4基础并重新量化会抹除大部分适配效果（§3.3）。

贡献：
1.  **SSD作为权重层**：实现从磁盘按需流式加载MoE专家权重的专家执行器，路由数学运算与官方模型保持比特级一致，每个执行器路径均逐元素验证与反量化参考的等价性，并设计插槽机制消除每步堆重建开销（§3.1）。
2.  **预测即路由**：跨令牌预路由器的预测*就是*解码时的路由，因此预取解码不丢弃专家且读取与计算重叠，在无法容纳检查点的机器上实现+80%至+84%解码加速，同时提出使语言模型在路由替换下工作的训练方案（§3.2, §4, §5.3）。
3.  **未合并恢复LoRA**：将适配器合并至int4基础并重新量化会破坏其大部分效果；我们证明以并行增量形式服务可避免此问题，且开销可忽略（§3.3）。
4.  **两个开源层级，端到端**：发布35B和8B模型的检查点与适配器组合，两者加载后平均性能仅比fp16基础低几个百分点（§5）。

## 2 背景与相关工作

#### MoE路由

生产级稀疏MoE每层将每个令牌路由至K个专家（共E个）[26,20,10,18]。两大路由流派占主导：Softmax-top-k（Qwen3.6）：对专家logits精确计算softmax，取top-k后重新归一化；Sigmoid-group（DeepSeek-V3[6], Ling[29]）：sigmoid评分；按每组前两个评分之和对组排名，保留top-G组；在幸存组内取top-k；权重为原始sigmoid值，重新归一化并缩放。Edge0完整实现两者，并在预路由器训练和推理中复用相同函数（§3.2）。

#### 预门控与前瞻距离

预门控MoE[17]在同一令牌内选择下一个块的专家：在块N注意力完成后生成决策，在块N+1权重加载前消费。该调度与流式引擎不兼容，Edge0改为前瞻一个完整令牌（§3.2，其中评估了逐层方案的代价）。

#### KV侧压缩

MLA[5]、稀疏注意力[7]和线性状态模型[13]压缩动态状态，服务系统采用分页而非压缩[19]。它们与Edge0正交且互补：Edge0的8B层级是MLA+MoE混合架构，压缩注意力为活跃专家集留出更多内存预算。

#### 量化

GPTQ/AWQ类方法[11,21]使4位成为训练后量化实用下限；低于4位后，后期误差增长过快，本地部署模型极少采用更低精度。Edge0使用int4仿射组64的专家权重，公开接受质量损失，随后通过蒸馏恢复大部分效果。QLoRA证明4位基础模型配合适配器可接近16位微调质量[8]；我们以4位基础为前提，研究适配器在服务时的耐受性（§3.3）。最接近的工作是专家跳跃自蒸馏[23]，表明训练后的MoE模型在专家数减半时仍能保持性能；我们的路由替换训练是将同一观察应用于预测即路由场景。

#### 卸载

专家卸载研究广泛：llama.cpp的`--cpu-moe`[12]将MoE专家存储在CPU/系统内存，PowerInfer[28]将神经元划分为跨GPU/CPU的热冷集合，Mixtral-offloading[9]和MoE-Infinity[33]按热度或复用距离缓存专家，FlexGen[27]将层级扩展至磁盘以同时存储权重和KV。它们均转移占用而非减少占用：权重仍占数十GB，卸载至磁盘的权重无法预知下一步需要的专家。

## 3 系统设计

图1：Edge0系统。专家权重以int4逐层堆叠的safetensors格式存储于SSD，通过mmap按需流式加载至有界流池（操作系统页缓存、LRU、分段固定插槽双缓冲）。冻结的int4解码器栈与并行未合并的恢复LoRA分支协同计算，输出通过求和节点结合：y=W_int4(x)+(α/r)·B·A·x。第N层拥有的预路由器头基于第N层的后注意力归一化输出，在令牌t时刻预测第N+1层在令牌t+1的路由（红色路径，“预测即路由”），使填充预取插槽的SSD读取与当前前向传播重叠；底部解码时间线显示令牌t+1的预测读取与令牌t的计算并发执行。Edge0是流式MoE推理框架：所有MLX代码位于后端抽象层之后（core/nn/io/quant），核心逻辑（模型规格、预路由器、流池、服务器）仅依赖该抽象层，其他后端可实现相同接口。模型由单个MoESpec描述（专家数K、路由族、量化、权重布局、键模板）；通用流式层服务所有模型。LoRA和预路由器权重为带溯源元数据的safetensors文件，从模型目录解析且永不合并至基础模型。图1展示系统架构与解码时数据流。

### 3.1 SSD流式专家层

#### 问题

40层、256专家的MoE模型每层专家权重在int4下占453MB——每个专家1.77MB——或在19.5GB检查点中占18GB（16.9GiB）路由专家，剩余为注意力、嵌入、共享专家和int4分组缩放。这相当于24GB机器内存的四分之三，尚未计入KV缓存、操作系统和用户其他程序。Edge0通过mmap映射量化后的专家safetensors，按需读取字节范围；操作系统页缓存承载热数据，峰值内存随*活跃*集而非参数总量变化。

#### 执行器路径

四条路径从正确性基线到预填充批量路径：
1.  **精确路径**：去重按需构建、堆叠、量化收集。所有其他路径的正确性基线。
2.  **分段路径**：固定插槽双缓冲，解码主力。路由索引通过插槽表映射并`take`；索引永不离开GPU。重复专家集复用缓存图节点；持久粘性插槽张量原地更新（`incr_stack`），使一步仅重写变化的专家，而非重建层内九个堆叠张量（门控、上投影、下投影权重及其缩放和偏置）。
3.  **热路径**：每层固定LRU常驻热专家集合，与PowerInfer的热冷分离及Mixtral-offloading、MoE-Infinity的热度缓存相同[28,9,33]；命中走堆叠收集，未命中回落至精确路径。
4.  **整层路径**：预填充时一次性加载一层所有专家，检查点逐层堆叠布局使其为九次直接读取（相同三个量化投影），而非256×9次构建；CPU加载与前一层GPU执行重叠。

#### 数学约定

所有路径使用相同量化收集内核计算down(silu(gate(x))·up(x))，测试套件逐元素比较每路径与反量化参考（相对L2误差<1%；测量残差≈0.24%为内核本身的bf16内部精度）。该约定允许执行器按层、按阶段自由切换路径而不改变输出。

#### 为何卸载能直接胜出

该硬件上的全驻留引擎并非快速基线，因为模型无法完全容纳：Mac mini M4 Pro配备24GB，原生mlx-lm服务器在驻留全部19.5GB 4位权重时解码速度为3.9令牌/秒，占用18.2GiB；而Edge0的K=4配置解码速度达20.4令牌/秒，占用2.9GiB。[1]同一机器、相同解码协议、开启思考模式。驻留路径占用18.2GiB无法回收，使KV缓存和操作系统几乎无可用空间；Edge0的代价在于图构建和张量组装，这是流式引擎而非权重的代价，是第6节讨论的主题。理想实验明确该代价：在完美路由预测和无限缓存下，剩余代价仅为每步张量组装，`incr_stack`直接处理（附录B）。

### 3.2 预路由器：预测即路由

#### 动机

第N+1层的专家选择依赖第N层输出，但在第N+1层需要开始加载时该输出尚不存在。串行等待将磁盘延迟嵌入每一步。预路由器通过双重偏移打破依赖：属于第N层的头在令牌t时刻基于第N层后注意力归一化输出运行，预测第N+1层在令牌t+1的路由。第N+1层消费提前一步的预测；填充其插槽的SSD读取与当前前向传播重叠（图1）。

前瞻距离必须为一个完整令牌，而非一层。像Pre-gated MoE那样在第N层注意力后立即选择第N+1层专家[17]，需要逐层同步和头评估，会耗尽GPU流水线（在我们的引擎中每步30–100ms，超过其可隐藏的加载时间，且我们测量的所有同令牌变体均低于朴素LRU基线），并且下一层注意力计算时尚未得到其专家选择所需信息。一个令牌的前瞻将头移动...