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

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计算重叠。

相关工作

现有卸载与推理系统分析：
* 密集模型卸载系统: 如FlexGen【21, Sheng et al., ICML 2023】, DeepPlan【12, Jeong et al., EuroSys 2023】, 和SwapAdvisor【10, Huang et al., ASPLOS 2020】等系统主要针对密集神经网络设计，不原生支持MoE模型的部署。
* 通用卸载系统: DeepSpeed-Inference【1, Aminabadi et al., SC 2022】和HuggingFace-Accelerate等系统虽然支持MoE模型，但它们将MoE层视为密集的层，未能利用专家在推理过程中的条件性稀疏激活特性。
* 内存交换系统: Sentinel【20, Ren et al., HPCA 2021】和DeepUM【14, Jung et al., ASPLOS 2023】等系统可以追踪深度学习模型的内存访问，但其追踪粒度不在专家级别，导致高错误率和性能下降。
* MoE专用推理系统: 近期一些系统如Mixtral-Offload【7, Eliseev & Mazur, 2023】, Ollama (Llama.cpp)【18, Ollama, 2024】, vLLM【15, Kwon et al., SOSP 2023】, 和BrainStorm【3, Cui et al., OSDI 2023】虽然支持MoE，但如表1所示，它们的缓存性能仍然不佳。
* 多GPU系统: InfiniGen【16, Lee et al., OSDI 2024】和TensorRT-LLM【17, NVIDIA, 2024】主要为多GPU云服务器设计，缺乏针对个人设备的优化缓存设计，性能甚至不如vLLM等系统。

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-7）：首先检查专家是否已在缓存中。如果命中，则直接返回。如果缓存未满，则按需获取专家，存入缓存并返回。
* 缓存替换（步骤8-19）：当缓存满时，触发替换机制。
1. 预测（步骤8）：调用PredictEAM函数，根据当前iEAM和EAMC计算出pEAM，以获取每个专家在未来被重用的似然度。
2. 计算优先级（步骤9-16）：为缓存中的每个专家计算一个优先级分数$p$。该分数综合了pEAM中的似然度以及一个层衰减因子 $(1 - e.layer_idx/L)$。
3. 驱逐与加载（步骤17-19）：驱逐优先级最低的专家，为新专家腾出空间。然后按需获取新专家，将其加入缓存并返回。预取机制可以集成在FetchOnDemand函数中。

算法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缓存占用增加，留给专家缓存的空间变小。
* MOE-INFINITY的延迟增长平缓。当缓存空间不足时，它会退化为按需获取，但由于KV缓存保留在GPU中，其性能仍优于vLLM和Mixtral-Offloading。
* vLLM的延迟增长更为显著，因为它同时卸载了KV缓存，导致KV缓存与专家参数的获取在PCIe总线上产生竞争，进一步降低了性能。


图9: EAMC容量。


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

微基准测试:
* EAMC容量: 如图9所示，一个较小的EAMC容量（120，约占总请求数的3%）足以在包含290个任务的混合工作负载中捕捉大部分MoE激活模式，使所有模型达到最低平均延迟。
* 工作负载变化下的鲁棒性: 如表3所示，当任务发生切换时，MOE-INFINITY能快速适应。在同一数据集内切换任务，约需50个请求恢复最优延迟；在不同数据集间切换，平均仅需30个请求，因为不同数据集的相似任务会重用激活模式。

A5 结论

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

A6 附录

A. 现有的后验预测器

现有的后验预测器（不改变模型架构和微调）无法完全捕捉MoE模型的稀疏激活特性。
* (1) 基于依赖的预测：这类预测器根据内存依赖关系估算专家激活【10, Huang et al., ASPLOS 2020; 11, HuggingFace, 2024; 1, Aminabadi et al., SC 2022】。由于同一MoE层中的所有专家都依赖于同一个路由器，这种方法无法捕捉到激活的选择性（S）和分组（G）特性。重用（R）特性也未在同一范围内考虑。
* (2) 基于计数的预测：这类预测器使用每个专家的聚合频率计数器来估算激活【3, Cui et al., OSDI 2023; 14, Jung et al., ASPLOS 2023】。由于专家从长远来看倾向于表现出均匀的激活，这无法捕捉到稀疏性（S）。此外，单个计数器无法指导层内和层间的分组激活（G）。
* (3) 基于局部性的预测：这类预测器根据LFU和LRU等启发式方法估算专家重用【7, Eliseev & Mazur, 2023; 14, Jung et al., ASPLOS 2023; 3, Cui et al., OSDI 2023; 1, Aminabadi et al., SC 2022】。虽然只考虑了被激活的专家（S, R），但重用预测并未跨迭代应用，因此在解码阶段会失效。

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. 系统实现