ExFold：统一专家折叠框架实现无训练MoE预填充-解码加速

# ExFold：用于免训练混合专家模型预填充-解码加速的统一专家折叠技术
来源：https://arxiv.org/html/2608.24938  
刘逸飞 陈俊逸 范思琪 冯超然 李明浩 张刘杰 程维航\\*通讯作者 袁立\\*通讯作者

###### 摘要

混合专家模型通过稀疏激活专家机制在保持每词元计算量可控的同时，实现了模型容量的规模化扩展，从而获得优异性能。然而，高服务效率的混合专家推理正面临日益严峻的挑战，因为其横跨两个具有根本不同瓶颈的推理阶段：预填充阶段主要受逐词元专家计算支配，而解码阶段则受限于批量激活专家集引发的访存开销。现有的免训练加速方法仅优化单一资源代理指标——要么是每个词元执行的专家，要么是每个批次激活的专家——且这些方法要么完全舍弃被排除专家的贡献，要么仅对其隐式近似。本文提出ExFold，一个统一的免训练专家折叠框架，可联合加速混合专家模型的预填充与解码。该框架将两个阶段统一表述为带预算约束的输出近似问题：仅执行特定阶段的受限专家集，同时利用经过校准的标量投影器将预算外专家的贡献投影到保留专家上。基于许多专家输出方向对齐但幅度各异的观察，ExFold在无标签数据上校准逐对标量投影矩阵，并在推理时使用该矩阵将被排除专家的贡献折叠到保留专家中。在此视角下，预填充加速转化为词元级Top-K折叠，解码加速转化为批次级专家池折叠。两个阶段仅在保留专家选择方式上存在差异，而被排除贡献的恢复由统一的折叠机制完成。ExFold作为即插即用模块实现在vLLM中，配合轻量级专家折叠CUDA内核，在保持约99%原始平均性能的同时，实现了最高1.41倍首词元时间（TTFT）和2.45倍每输出词元时间（TPOT）的加速。ExFold代码见于 https://github.com/Time-Rune/ExFold-MoE。

## 引言

图1：混合专家模型预填充与解码阶段的瓶颈差异。图2：ExFold预填充与解码加速的质量-速度权衡。混合专家模型通过数量级的参数规模扩展模型容量，同时仅激活每个词元的稀疏参数子集来保持计算预算可控（文献8、16）。这一特性使混合专家模型成为众多追求更高能力的大型模型的主流设计选择（文献16、26、32）。与此同时，它也对生产环境中的低延迟服务提出了严格要求（文献4、3、7）：预填充延迟（首词元时间TTFT）和解码延迟（每输出词元时间TPOT）都至关重要，且两者具有截然不同的计算特征。如图1所示，预填充阶段主要受词元级专家计算支配，而解码阶段主要受批次级专家访存开销支配。

现有免训练混合专家加速方法根据优化的预算不同可分为两类：*词元级专家稀疏化*减少每个词元执行的专家数量，包括Dynamic-MoE（文献11）、MoDES（文献12）、NAEE（文献22）和MC-MoE（文献20）；*专家集合并*则通过静态剪枝合并（如REAP（文献18）、HC-SMoE（文献2）和Sub-MoE（文献19））或动态批次级限制（如Lynx（文献9）和SERE（文献31））来缩减活跃专家集。

然而，这两类方法都将加速视为资源削减问题：它们仅回答执行哪些专家，却未解决更关键的问题——被排除的专家去向何处。这些方法对排除贡献的处理充其量是附带性的。剪枝跳过直接丢弃贡献（文献22），静态合并将其固化为单一永久压缩模型，无法适应推理时实际面临的逐词元或逐批次预算（文献20）。基于相似性的重路由将被排除专家替换为相近专家，但从未校准替代输出偏离原始贡献的程度（文献31）。被排除的专家信息无论丢失、冻结还是仅隐式近似，都从未被显式重建——因此当预算收紧、这些方法被迫承受更大压力时，近似误差会不断累积。

我们提出*专家折叠*技术，将每个预算外专家的贡献投影到执行专家上而非丢弃，并基于此构建ExFold框架。该设计基于两项实证观察（图3）：许多源专家至少存在一个输出方向对齐的目标专家，同时不同专家间的输出幅度存在显著差异。前者使得专家替代成为可能，后者解释了为何直接重路由不足，并推动我们设计能够校正源-目标幅度失配的方向性标量投影器。具体而言，通过一次无标签文本的冻结前向传播校准每层的两个矩阵：*标量投影矩阵*和*投影损失矩阵*。推理时，损失矩阵将每个被排除专家路由到损失最小的目标，标量矩阵则将排除专家的贡献折叠到该目标的路由器权重中。

由于折叠定义在单个被排除专家层面，其在两个阶段的应用保持不变。预填充阶段每个词元选择Kpre个主导专家，解码阶段每个批次选择D个主导专家——两者复用相同的标量矩阵、损失矩阵和折叠操作符来恢复保留集之外的所有贡献。这是ExFold的核心设计：两个阶段仅在保留专家选择方式上存在差异，而被排除贡献的恢复由统一机制完成。

ExFold作为即插即用模块实现在vLLM中，配合轻量级专家折叠CUDA内核，在保持约99%原始平均性能的同时，实现了最高1.41倍TTFT和2.45倍TPOT加速。我们的贡献如下：
- • 阶段感知的公式化：我们表征了预填充与解码阶段不同的专家预算，并将两者统一为特定阶段选择约束下的单一输出近似问题。
- • 专家折叠：我们提出一种免训练折叠机制，通过方向性投影器和重建损失矩阵将预算外专家贡献投影到保留专家上，其标量形式直接融入路由器权重。
- • 评估验证：我们在真实vLLM服务负载下，针对多种混合专家架构验证了近乎无损的预填充、解码及联合加速效果。

## 背景与动机

### 混合专家模型推理

图3：（a-c）混合专家模型的专家输出可能在不同专家和层级间对齐，（d）其L2范数存在显著差异。#### 稀疏混合专家层。

稀疏混合专家层用N个门控前馈网络专家和路由器替代Transformer块中的稠密前馈网络（文献8、16）。对于维度为d的隐藏状态x，专家e的计算如下：

Ee(x)=Wdown(e)(φ(Wgate(e)x)⊙Wup(e)x)      (1)

其中φ(⋅)通常为SiLU激活函数，⊙表示逐元素乘法。路由器生成logits r(x)=Wrx，对N个路由专家选择Top-K集合SK(x)，并在该集合内归一化门权重：

αe(x)=exp(re(x))∑j∈SK(x)exp(rj(x)), e∈SK(x)      (2)

混合专家输出则为：

MoE(x)=∑e∈SK(x)αe(x)Ee(x)      (3)

可能包含额外的共享专家。该计算使前馈网络路径逐词元稀疏化，但也引入了动态路由、专家分发与专家输出聚合。

#### 阶段特定预算。

尽管相同的混合专家层（公式3）在两个阶段都运行，但其瓶颈随词元批的形状变化而转移（图1）。预填充阶段一次处理T个词元的提示，每个词元独立激活自身的支持集SK(x)，因此层执行O(TK)次专家前馈计算。当T较大时，每个加载的专家被众多词元复用，权重加载得以分摊，该阶段受限于逐词元专家浮点运算；直接减少K可降低成本。解码阶段每步为B个批处理请求各前进一步，仅生成B个词元，但其支持集⋃xSK(x)通常覆盖大部分N个专家。因此层必须加载该批次专家并集以服务极少词元，导致算术强度低下，成为受限于访存的阶段，其成本取决于不同激活专家的数量而非浮点运算。高效服务因此受两个不同预算制约：预填充阶段的每词元专家数，和解码阶段的每批次专家数。

### 混合专家加速方法

根据优化的预算不同，现有免训练混合专家加速方法可分为两类：词元级专家稀疏化和专家集合并。

#### 词元级专家稀疏化。

该家族减少每个词元执行的专家数量。Dynamic-MoE为每个词元保留累积路由器概率超过p的最小专家集，使较容易的词元使用更少专家（文献11）。MoDES基于校准估算的重要性跳过逐词元专家（文献12），NAEE根据路由器或激活统计剪枝低重要性执行（文献22），MC-MoE合并从校准中识别的共激活专家（文献20）。这些方法缓解了预填充计算预算压力，但未控制解码时的专家并集。

#### 专家集合并。

该家族缩减活跃专家集。静态方法根据校准信号为每个专家评分并永久缩减全局池：REAP通过路由器权重和激活范数剪枝最不重要专家（文献18），而HC-SMoE（文献2）、REAM（文献15）和Sub-MoE（文献19）合并相似专家。动态方法限制解码批次内的专家并集：Lynx将次级请求重路由至活跃专家（文献9），SERE离线校准逐对专家相似性，并在解码时用相似保留专家替代每个被排除专家（文献31）。这些方法缓解了解码端访存压力，但其静态变换无法匹配逐词元预填充预算。

#### 差距所在。

两类方法都决定*执行哪些*专家，但都未能保留被排除专家的贡献：剪枝直接丢弃，静态合并将其固化为永久压缩模型，相似性重路由仅改变目标而不校准替代偏离程度。随着预算收紧，误差在每个阶段内增长并在预填充与解码间累积，因为两者针对不同目标进行近似而非统一目标。

核心问题：*如何设计统一框架，以最小损失加速混合专家模型的预填充与解码阶段？*

## 设计洞见

这一差距使得从预填充和解码两侧恢复都看似困难——两者指向冲突的预算，忠实恢复每个被排除专家似乎需要重计算或盲目替代——但两项观察消解了这两重困难。

#### 被排除贡献共享单一恢复目标。

预填充与解码受不同预算约束（每词元专家数 vs 每批次专家数），但无论预算排除哪些专家，其留下的残差在混合专家输出（公式3）中都具有相同形式αe(x)Ee(x)。因此两个阶段构成单一问题——从保留专家重建被排除贡献——仅在选择保留集的约束上存在差异。

#### 专家冗余具有幅度可分性。

图3显示专家间的冗余集中在一个易于修正的自由度上——输出幅度。在二维主成分投影（a）中，原始专家输出形成扩散云团，但归一化到单位尺度后它们坍缩到一条细长共享轴上：消除幅度后，专家指向几乎相同方向。逐对相似性证实了这一点：归一化使第5层的平均专家间相似性从0.335提升至0.529，第17层从0.251提升至0.471（b）——且所有48层都存在相同差距（c），表明这是结构性对齐而非少数层现象。但被丢弃的幅度很大：原始输出范数在单层内跨越超过3倍（所示两层最大/最小分别为3.04倍和3.27倍，d）。因此排除专家与保留专家的失配几乎是纯径向的——兼容方向已存在，仅幅度不匹配。因此，单个逐对标量重缩放保留专家——而非重建向量或盲目替代——原则上能限定替代偏离程度并实现廉价忠实恢复。完整层可视化见附录D。

这些观察共同为动机问题提供了肯定答案：通过廉价的逐对标量校正到达单一恢复目标，该目标在两个阶段共享，每个阶段仅改变保留集的选择方式。下一节将ExFold转化为具体机制。

图4：ExFold概述：免训练投影器校准与统一专家折叠。

## 方法

### 统一专家折叠

ExFold将混合专家加速视为受限输出近似而非特定阶段专家丢弃。令xi表示第i个输入词元，SK(xi)为其原始Top-K支持集。我们将路由器权重扩展为：对于e∉SK(xi)，αe(xi)=0；对于保留专家，其权重通过折叠操作重计算。具体而言，被排除专家的贡献被投影到保留专家上：每个被排除专家e的贡献αe(x)Ee(x)被映射到损失最小的目标专家e'∈SK(x)，并应用标量校正因子s_{e→e'}进行幅度调整。折叠后的输出为：

ExFold(x) = ∑_{e∈SK(x)} [αe(x) + ∑_{e'∉SK(x), e'=argmin Loss(e→)} s_{e→e'}αe'(x)] Ee(x)

其中s_{e→e'}来自校准的标量投影矩阵。该机制保持输出维度不变，并显式重建被排除贡献。

### 折叠机制校准

校准阶段包含两个步骤：
1. 损失矩阵构建：对无标签数据执行前向传播，记录每个专家e到所有其他专家e'的输出方向对齐程度，构建方向性投影损失矩阵L(e→e')。
2. 标量投影器校准：在损失最小的投影方向上，求解标量缩放因子s_{e→e'}，使s_{e→e'}Ee'(x)最接近Ee(x)的幅度，得到标量投影矩阵S。

该过程仅需一次冻结前向传播，无需更新模型参数。

### 阶段特定实现

在预填充阶段，每个词元独立选择Kpre个保留专家，折叠操作应用于该词元视角下的被排除专家。在解码阶段，为批次选择D个保留专家，折叠应用于批次级专家并集视角。两种情况均使用相同的校准矩阵和折叠操作符。

### 系统集成

ExFold以即插即用模块形式集成到vLLM推理框架中，通过轻量级CUDA内核实现：
1. 在专家分发前插入折叠计算层
2. 实时应用标量投影器调整路由器权重
3. 保持原有计算图其余部分不变

这种设计最小化集成开销，同时实现显著的端到端加速。