MOE-INFINITY: Efficient MoE Inference on Personal Machines with Sparsity-Aware Expert Cache

发表时间: 2024-01 · arXiv:2401.14361 (University of Edinburgh)

原文: https://arxiv.org/abs/2401.14361

作者/机构: Leyang Xue 1, Yao Fu 1, Zhan Lu 1, Chuanhao Sun 1, Luo Mai 1, Mahesh Marina 1

速读

一句话结论 本文针对个人设备显存受限导致的大型专家混合模型推理极慢的问题，提出了一种基于请求级稀疏激活追踪的专家缓存系统 MOE-INFINITY，在单 GPU 环境下将各类大型模型的单 Token 生成延迟降低了 3.1 到 16.7 倍。

要解决什么问题 现有支持参数卸载的推理系统在处理专家混合模型（MoE）时面临严重性能瓶颈。个人设备的 GPU 显存通常仅 24 到 48GB，运行上百 GB 的 MoE 模型必须将参数卸载到主机内存并按需加载。然而，现有系统大多基于计算图执行顺序进行预测，错误假设所有潜在专家都需要加载到 GPU。这种机制忽略了 MoE 的稀疏激活特性，导致大量 PCIe 总线 I/O 拥塞和 GPU 空闲等待，极大增加单 Token 延迟。更深层原因是，传统缓存替换策略无法适应 MoE 的动态路由。在批处理大小为 1 的个人解码场景下，专家重用模式在单请求内部极度偏斜（少数专家被反复调用），但观察多个请求时，激活分布又趋于全局均匀。这种全局均匀性使传统策略在缓存满载时，无法准确判断当前解码阶段哪个专家最不可能被重用，从而做出错误驱逐决策，引发频繁缓存未命中。

怎么做的 核心思路是利用 MoE 模型在单请求解码阶段高度偏斜的稀疏激活特性，将专家激活预测转化为在线模式匹配问题，构建稀疏性感知的专家缓存。该方法能绕开 PCIe I/O 瓶颈，是因为每次解码迭代仅有极小的专家“工作集”被激活，精准预测并将其锁定在 GPU 缓存即可消除绝大部分按需加载开销。关键设计由三个部件构成。首先是专家激活矩阵（EAM），用于形式化追踪激活状态。对于 $L$ 层、每层 $E$ 个专家的模型，EAM 是一个 $L \times E$ 的矩阵 $M$，$M[i][j]$ 记录路由到特定专家的 Token 数量。系统会在单次迭代（iEAM）和单请求（rEAM）层级构建该矩阵。其次是专家激活矩阵集合（EAMC），用于在线保存历史请求级 EAM。由于不同请求间的模式转换不可预测，但历史请求中存在有限数量的相似激活组，新请求到来时，系统将当前 iEAM 展平为向量，通过计算余弦距离在 EAMC 中寻找最匹配的历史 EAM。最后是结合预取的缓存替换引擎。引擎聚合匹配到的历史 EAM 并按行归一化，计算专家未来的激活概率。缓存满时，系统为缓存中的专家计算优先级分数 $p$：$$p = \frac{p\_eam[e.layer\_idx] + \epsilon}{n\_token} \cdot (1 - \frac{e.layer\_idx}{L})$$ 引擎驱逐分数最低的专家。此外，系统利用模型逐层执行顺序，提前将下一层专家预取到缓存；并为初始几层专家分配更高缓存优先级，以抵消请求初期预测置信度低带来的预取失败风险。

效果如何 实验在单张 NVIDIA RTX A5000 GPU（通过 PCIe 4.0 连接主机内存）上搭建，评估了 Switch Transformers、DeepSeek-V2-Lite、Mixtral-8x7B 及 900GB 的 Snowflake-Arctic 等模型，测试数据覆盖 BIG-Bench 等 290 个任务。对比基线包括代表主流高效卸载路线的 DeepSpeed-Inference、代表资源受限 GPU 优化路线的 Llama.cpp、代表 MoE 专用卸载路线的 Mixtral-Offloading，以及代表动态神经网络统计计数路线的 BrainStorm。在提示词长度 512、解码长度 32 的设置下，MOE-INFINITY 将单 Token 延迟降低了 3.1 到 16.7 倍。对于 Mixtral 模型，延迟低至 836 毫秒，比次优的 BrainStorm 快 1.4 倍；对于激活率更低的 DeepSeek 等模型，延迟降至 155 毫秒，几乎媲美完整模型常驻 GPU 的性能；它更是唯一能在单 GPU 上流畅运行 900GB Arctic 模型的系统。在高达 128K 的长上下文任务中，随 GPU 内 KV 缓存占用增加、专家缓存空间被压缩，系统会平滑退化为按需获取，但因避免了 KV 缓存卸载导致的 PCIe 竞争，延迟表现依然显著优于 vLLM。代价方面，在线匹配 EAM 需消耗少量 CPU 资源，但在 1000 个历史记录下每次查询仅耗时 21 微秒；此外，面临全新工作负载切换时系统存在短暂适应期，平均需处理 30 到 50 个请求才能通过更新 EAMC 恢复到最低延迟。

A1 主要贡献

本文旨在解决在个人设备上高效运行大型专家混合（MoE）模型的问题。个人设备通常GPU显存有限（24-48GB），而大型MoE模型（如超过100GB的DeepSeek-MoE）必须依赖将模型参数卸载（offloading）到主机内存的方案。

核心问题: 现有支持卸载的推理系统（如DeepSpeed, Mixtral-Offloading, vLLM）在处理MoE模型时性能不佳。其根本原因在于缓存设计未能考虑MoE模型在推理过程中的稀疏激活特性。这些系统通常基于计算图中的执行顺序进行预测，假设所有潜在的专家都需要被加载到GPU，这导致了大量的PCIe总线I/O瓶颈和GPU空闲时间，从而增加了每个输出Token的时间（TPOT）。如表1所示，与MOE-INFINITY相比，其他系统延迟显著更高。即便是为动态神经网络设计的BrainStorm系统，也因无法适应MoE模型的路由机制而表现不佳。

研究目标与核心思想: 本文的目标是为个人设备上的MoE推理设计一种高效的缓存策略。核心思想是：在个人设备这种通常批处理大小为1的单用户场景下，MoE模型在解码阶段表现出高度的激活稀疏性，即一小部分专家会被频繁重用。利用这一特性，可以设计一个有效的专家缓存，通过预测哪些专家最有可能被重用，来优化缓存的替换和预取。

主要创新点:

  1. 对MoE模型激活稀疏性的实证分析: 本文通过广泛的轨迹分析，首次提供了在个人部署场景下设计高效缓存可行性的证据。研究发现，在批处理大小为1的情况下，消费级GPU的内存足以缓存解码阶段最常用的专家。此外，研究观察到专家重用模式在单个请求（prompt）内呈偏斜分布，但在处理多个请求后，这种偏斜消失，所有专家趋于均匀激活。
  2. 专家激活的在线预测问题建模与方法: 本文将专家激活的偏斜重用模式形式化为在线预测问题。通过统计建模和分析，开发了一种专家激活预测方法。该方法通过追踪专家在请求级别的稀疏激活，并选择能代表稀疏模式的轨迹来指导未来的预测。此方法被集成到一个高性能的专家缓存中。
  3. 高性能MoE推理系统MOE-INFINITY: 本文实现了MOE-INFINITY系统，其核心是稀疏性感知的专家缓存。在与vLLM、Ollama、DeepSpeed等多个先进系统的对比测试中，MOE-INFINITY在处理不同MoE模型和LLM任务时，将每个Token的延迟降低了3.1-16.7倍。该系统已在GitHub开源。


图1: 在GPU上进行MoE推理，完整模型已卸载至CPU内存。E[0,1]指第0层索引为1的专家模块。

A3 背景与相关工作

背景知识：MoE推理与卸载

MoE模型结构与卸载推理流程：如图1所示，在个人设备上运行MoE模型时，一份完整的专家参数副本保存在主机内存中，而密集激活的参数（如注意力层和KV缓存）则常驻GPU内存。每个MoE层包含一个路由器和一组专家（即前馈网络FFN）。路由器负责将每个输入Token分配给特定的专家进行处理。不同MoE模型的配置各异，例如，Mixtral-8x7B每层有8个专家，路由器为每个Token选择2个；而Switch-128x0.2B每层有128个专家，路由器为每个Token选择1个。MoE模型的推理过程包括预填充（prefilling）阶段和解码（decoding）阶段。在每个Token生成过程中，当注意力和路由器确定了需要激活的专家后，卸载系统会检查GPU内的专家缓冲区。如果专家参数不在缓冲区中，系统需要按需从主机内存中获取。为了提升性能，卸载系统通常会配备一个专家激活预测器，以辅助路由器决策并提前预取专家，从而将I/O操作与GPU计算重叠。

相关工作

现有卸载与推理系统分析：

A2 方法细节

4. 稀疏性感知的专家缓存

本节阐述了稀疏性感知的专家缓存的设计原理。首先解释为何小型专家缓存对加速LLM解码中的卸载是有效的，然后识别了优化缓存性能所需的稀疏激活模式，并介绍了如何进行在线预测以改善缓存替换。接着，对在线预测问题进行形式化，并指出现有追踪方法的局限性。最后，引入了本文提出的请求级稀疏激活追踪方法，并利用该追踪信息提升专家缓存效率。


图2: 两个样本序列及合并1000个序列后，专家在解码迭代中的重用次数。颜色越深表示归一化后的重用次数越高。样本来自Mixtral-8x7B的最后一层（上图，20次解码迭代）和DeepSeek-V2-Lite的最后一层（下图，256次解码迭代）。

4.1. 为什么专家缓存可能是有效的

小型工作集假设：专家缓存之所以有效，关键在于LLM解码的每次迭代中只使用一小部分专家。这意味着存在一个小的“工作集”，其大小可以容纳在GPU有限的内存中。如果工作集过大，将导致频繁的I/O操作，反而拖慢推理速度。

工作集大小的实证研究：为了估算工作集的大小，本文对多种主流MoE模型在不同LLM任务中进行了广泛的轨迹研究。研究发现：对于拥有约100个专家的MoE模型（如DeepSeek, QWen-MoE, NLLB, Switch-MoE），在为单个请求解码Token时，只有不到5%的专家被重复激活。即使对于专家数量较少的模型（如Mixtral），每个请求也仅观察到25%的专家被激活。这些结果表明，专家经过充分训练，已能专门处理不同类型的请求。

4.2. 为什么专家缓存必须是稀疏感知的

缓存替换策略的挑战：专家缓存的核心挑战是在缓存满时决定替换哪个专家。最优策略依赖于预测哪个专家在近期最不可能被激活。一种直接的方法是追踪专家随时间的激活频率，如BrainStorm【3, Cui et al., OSDI 2023】和DeepUM【14, Jung et al., ASPLOS 2023】所采用的。

传统方法的局限性：然而，本文发现这种方法对于专家缓存是不够的，因为它没有考虑到单个请求的稀疏激活模式。如图2所示，对于Mixtral模型的两个不同序列，专家的重用模式截然不同：序列1中专家2被重用了超过15次，远超其他专家；而序列2中专家0和专家5的重用次数最高。同样，在DeepSeek模型中，尽管解码迭代次数多达256次，专家激活仍然是稀疏的。

请求级偏斜与全局均匀分布的矛盾：虽然专家重用模式在单个请求内是偏斜的，但在多个请求之间会变得更加均匀。当分析超过1000个序列时，重用计数趋于平稳。在这种均匀的专家使用分布下，缓存将无法判断在解码阶段哪些专家更可能被重用，从而损害缓存性能。

4.3. 解码过程中激活预测的形式化

问题形式化定义：本文将LLM解码过程中的专家激活在线预测问题进行了形式化。问题的输入定义为专家激活矩阵（Expert Activation Matrix, EAM）。对于一个有L个MoE层、每层E个专家的模型，EAM是一个 $L \times E$ 的矩阵 $M$，其中 $M[i][j] \in \mathbb{Z}$ 表示路由到专家 $E[i,j]$ 的Token数量。

在线预测的目标：在线预测问题利用EAM来找出未来专家的激活可能性和所有专家的重用可能性。预测器在每个MoE层路由决策后被触发，并提供一个预测的EAM（pEAM），其每个条目代表重用或激活的似然度。

EAM的构建层次：

  1. 迭代级EAM (iEAM)：记录模型单次前向传播中每个序列的轨迹。在预填充阶段，Token数n为序列长度；在解码阶段，n=1。
  2. 请求级EAM (rEAM)：累积每次迭代的iEAM计数。rEAM分别对预填充和解码阶段进行追踪。每次迭代结束后，iEAM会更新到累积的rEAM中，该rEAM记录了自当前请求开始以来专家的使用频率。

表2: K-means聚类中肘部点确定的组数。组大小按总序列数归一化。


图3: 使用K-means对激活矩阵进行聚类，同一组内的矩阵具有相似的值。激活状态由马尔可夫链建模。


图4: 使用MMLU数据集的Arctic模型的所有组的马尔可夫链转移矩阵（左）和其中一个组的概率（右）。

4.4. 有效预测方法的直觉

预测方法的核心观察：准确预测专家激活需要通过匹配已观察到的专家激活模式组来实现。本文从两个方面验证了这一点：(i) 专家激活组的存在性；(ii) 不同组之间的转换难以预测。

激活组的存在性：首先，请求级EAM（rEAM）之间存在强相似性，这使得基于已知模式进行激活预测成为可能。本文利用K-means算法对一组rEAM进行聚类，其中每个EAM代表一个序列的激活模式。如表2所示，对于每个数据集的1000个样本，仅形成了10到30个激活模式组，这表明可以用一个EAM来推断同一聚类组内的其他EAM。

组间转换的不可预测性：其次，组间转换的预测是困难的，因为所有组在请求级别都表现出稀疏的激活模式。如图4所示，组间的转移概率很低（最高约0.3，大多数低于0.12），这意味着利用现有组信息来推断其他组是不可行的。这排除了基于后验学习的预测方法，因为数据分布的转变本质上是组间的转换，而这种转换是难以预测的。

核心结论：在请求级别，激活的专家可以被归类为有限数量的EAM组。然而，不同请求下EAM之间的转换是不可预测的。因此，一个合理的专家激活预测方法是利用组匹配来处理新的请求。


图5: 计算激活可能性的示例。

4.5. 请求级稀疏激活追踪

EAMC数据结构：为了在请求级别追踪专家的稀疏激活，本文设计了一种名为专家激活矩阵集合（Expert Activation Matrix Collection, EAMC）的新数据结构，用于在线保存历史的请求级EAM。当系统处理一个新请求时，它会将当前请求的EAM与EAMC中存储的近期EAM进行比较。通过找到一个匹配的先前EAM，可以更有效地指导预取和缓存决策。匹配方法是：将每个EAM展平为向量，并计算它们之间的余弦距离。与当前EAM余弦距离最小的先前EAM即为最接近的匹配。

激活似然度计算（PredictEAM）：给定一个EAMC和一个当前的iEAM，激活似然度的计算过程如图5所示。

  1. 匹配：当模型处理到某个中间层（如图中R2完成Token分发后），MOE-INFINITY利用当前迭代的iEAM与EAMC中的历史EAM进行匹配。
  2. 聚合与归一化：可能会返回多个匹配的EAM。此时，将这些匹配EAM的单元格值相加，并在每行上进行归一化，从而计算出每个可能激活的专家的激活概率。
  3. 层邻近度调整：为了优先考虑靠近当前层的未来专家，通过公式 $(1 - (i - l)/L)$ 调整每个单元格的值，其中 $l$ 是当前层ID，$i$ 是未来层ID。

4.6. 缓存优化

结合预取的缓存增强：为了进一步提升性能，本文将预取机制集成到专家缓存中。利用MoE模型逐层执行的顺序性，可以使用pEAM来预测下一层可能被激活的专家，并提前将它们预取到缓存中，从而减少因按需获取专家而导致的GPU停顿。

结合专家位置信息的缓存增强：在决定替换哪个专家时，本文还考虑了一个观察：MoE模型的初始几层通常从预取中获益较少，因为在请求开始时对激活组模式的预测置信度较低。通过为这些初始层的专家分配更高的缓存优先级，不仅可以抵消潜在的预取失败，还能利用模型的逐层执行特性：后续层执行较晚，更可能从预取中受益，因此对缓存的需求较低。

4.7. 稀疏性感知的专家缓存算法

算法1：专家缓存检索：算法1展示了专家缓存的检索过程。

算法1: 专家缓存检索
Require: cur EAM – 当前迭代级EAM, id – 请求的专家ID, eamc – 历史rEAM列表, cache – 存储缓存专家的字典, cache_size – 最大允许缓存大小, m – 包含L个层的模型实例
Output: expert – 检索到的专家实例

1: if id ∈ cache then
2:   return cache[id]
3: end if
4: if |cache| < cache_size then
5:   cache[id] ← FetchOnDemand(id)
6:   return cache[id] {缓存未满}
7. end if
8: p_eam ← PredictEAM(eamc, cur_EAM)
9: evict_expert ← None, p_min ← ∞
10: for id, e in cache do
11:   n_token ← Σ p_eam[e.layer_idx]
12:   p ← (p_eam[e.layer_idx] + ϵ) / n_token · (1 − e.layer_idx/L)
13:   if p < p_min then
14:     p_min ← p, evict_expert ← e
15:   end if
16: end for
17: delete cache[evict_expert.id]
18: cache[id] ← FetchOnDemand(id)
19: return cache[id]

与现有方法的对比示例：图6展示了MOE-INFINITY与传统LRU（如vLLM）和统计计数（如BrainStorm）方法的对比。在示例场景中，LRU和统计计数方法由于缺乏对请求级稀疏模式的感知，会做出错误的驱逐决策，导致缓存未命中。而MOE-INFINITY的方法能准确预测并保留即将被使用的专家，同时预取后续层的专家，从而实现更高的缓存命中率和更低的延迟。


图6: 集成缓存与预取的示例。LRU是vLLM、Llama.cpp、DeepSpeed等SOTA系统中最常用的技术，统计计数在BrainStorm中实现。

A4 实验

实验环境

实验结果


图7: 解码延迟。vLLM和Llama.cpp不支持Switch、NLLB和Arctic，因此省略了结果。

端到端性能:
如图7所示，在prompt长度为512、解码长度为32的设置下，MOE-INFINITY表现出显著优势。


图8: 长上下文下的解码延迟。使用最大上下文长度为128K的DeepSeek-V2-Lite。我们展示了在长上下文生成中排名前4的系统。

长上下文性能:
如图8所示，在长上下文生成任务中（上下文长度从4096到131072），随着上下文增长，GPU中KV缓存占用增加，留给专家缓存的空间变小。


图9: EAMC容量。


表3: 处理工作负载变化。每个单元格中的数字表示工作负载变化后恢复低延迟所需的最小、平均和最大请求数。

微基准测试:

A5 结论

MOE-INFINITY是首个使个人计算机在运行如DeepSeek等大型MoE模型时能达到有竞争力性能的系统。其核心设计——稀疏性感知的专家缓存——在多种MoE模型和LLM任务上进行了广泛评估，相较于强大的基线系统，实现了2.7-13.7倍的性能提升。我们期望MOE-INFINITY能使AI开发者更轻松、更高效地部署本地MoE推理服务，其开源特性也有望推动社区的快速采纳和进一步创新。

A6 附录

A. 现有的后验预测器

现有的后验预测器（不改变模型架构和微调）无法完全捕捉MoE模型的稀疏激活特性。

B. 实践中的考量

B.1. 预测器运行时效率

EAMC的设计与替换策略：EAMC被设计为固定容量，以限制内存成本和查找匹配EAM的时间。当EAMC达到容量时，会触发替换策略。该策略旨在记录最新的EAM以快速适应工作负载变化，并保持EAM集合的多样性。因此，我们选择替换与新进入的EAM最相似的EAM。如图10所示，当新提示P4的轨迹完成后，计算其EAM4与EAMC中所有EAM的余弦距离，并替换掉距离最近的EAM3。

运行时开销：一个适度小的EAMC容量是足够的。MoE路由器经过预训练后会创建专门的专家组，限制了组的数量。实验表明，几百到几千的EAMC容量足以应对各种LLM任务。查找最相似EAM本质上是CPU上的矩阵乘法【5, Douze et al., 2024】，在1K个EAM下每次查询耗时21µs，开销远小于模型推理延迟（通常>120ms/token）。

容量上限：将EAMC构建问题形式化为球体覆盖问题，理论分析【19, Rankin, 1947; 6, Dumer, 2007】表明，覆盖所有专家激活模式所需的聚类总数是有限的，并且是关于专家数量的多项式复杂度。对于SOTA模型，这可能需要高达40K个EAM，占用约160MB内存。

运行时优化：由于在线聚类算法的计算复杂度高，目前采用了上述简单而有效的设计，将聚类算法的增强留作未来工作。


图10: EAMC替换示例。

B.2. 系统实现