论文

SlimWise:在 Prefill 与 Decode 间解耦专家剪枝以实现高效 MoE 服务

SlimWise:在 Prefill 与 Decode 间解耦专家剪枝以实现高效 MoE 服务

Mixture-of-experts (MoE) 模型每个 token 只激活少量专家,但批量解码几乎会访问整个专家池,使专家权重流量成为主要瓶颈。传统专家剪枝虽能降低该流量,却同时剪掉了计算受限的 prefill,几乎换不来吞吐收益,还牺牲模型质量。 SlimWise 是一个按推理阶段定制专家池的服务框架:prefill 使用完整模型,decode 使用剪枝模型,并直接复用 prefill 生成的 KV cache,无需任何转换。在两个 MoE 骨干和三种剪枝准则下,这种免训练的 KV cache 交接在多种设置中显著缩小了与完整模型的精度差距。 我们还发现,benchmark 精度会掩盖剪枝引起的生成长度变化。为缓解这类失真与残余精度损失,SlimWise 引入低成本的蒸馏阶段:让 decoder 从完整模型的 KV cache 继续生成,且只更新一小部分参数。基于 vLLM 实现,支持 PD 分离与 PD 同置两种服务模式。在 Qwen3.6-35B-A3B 上,50% 专家剪枝下 decode 吞吐最高提升 1.81x,精度损失极小。

论文精读

TL;DR SlimWise 解耦 MoE 推理的 prefill/decode 专家剪枝:prefill 用全模型生成 KV cache,decode 用剪枝模型直接复用,配合低成本蒸馏,在 50% 剪枝下 decode 吞吐最高 1.81x 且精度损失极小。

问题

问题背景

MoE 推理服务 的核心挑战是在 连续批处理 下,每个解码步骤加载的专家权重集合接近全量专家池,导致 权重传输流量 成为主要瓶颈。

现有方法局限

传统 专家剪枝 方法在 prefill 和 decode 阶段统一剪枝,虽然减少了 decode 的权重流量,但同时也削弱了计算密集的 prefill 阶段模型容量,造成明显的精度损失,且 prefill 本身吞吐收益有限。更关键的是,剪枝后的 prefill 会生成 低质量 KV cache,进一步损害下游 decode 的输出质量,且不同剪枝策略下生成长度可能发生扭曲,使得 基准准确率 掩盖了实际分布偏移。

为什么难点/重要

prefill 与 decode 的资源瓶颈不同:prefill 是 计算密集,decode 是 内存/带宽密集。需要将专家剪枝仅应用在 decode 阶段,并让 decode 模型直接复用 full-model prefill 生成的 KV cache,但面临两个挑战:模型架构不匹配(剪枝后专家维度变化) 和 蒸馏成本过高(若需训练恢复精度)。SlimWise 通过训练-free 的 KV cache 直接接管和低成本的 仅更新少量参数蒸馏 来解决,在 Qwen3.6-35B-A3B 上实现了 50% 剪枝下 1.81× decode 吞吐提升且精度损失最小,体现了 分阶段差异化优化 的工程价值。

行业类比

类似 长上下文对话系统 中,prefill 一次编码全部历史,decode 逐 token 生成,对 prefill 和 decode 分别采用不同精度的专家集合,可大幅降低生成阶段带宽压力。

核心洞察

  • 阶段解耦的专家裁剪:SlimWise 的核心洞察在于 prefill 与 decode 阶段的瓶颈性质截然不同——prefill 是计算密集,decode 是内存带宽密集。传统 expert pruning 对所有阶段统一剪枝,导致 prefill 质量损失却无吞吐收益。SlimWise 保留 prefill 完整模型,仅对 decode 裁剪,实现差异化优化。这一思路挑战了“一刀切”的剪枝假设,为 MoE 推理优化提供了新的维度。
  • 训练无关的 KV cache 交接:裁剪后的 decode 模型可直接复用完整模型生成的 KV cache,无需转换或对齐,因为专家裁剪不改变隐藏维度。这一零成本交接不仅简化了部署,还大幅缩小了精度差距,表明剪枝损失主要集中在 prefill 阶段。与需要额外适配的方法相比,SlimWise 的做法展示了 MoE 剪枝中阶段隔离的有效性,为快速迭代提供了可能。
  • 蒸馏纠正生成长度失真:SlimWise 发现基准精度可能掩盖剪枝导致的生成长度变化,进而影响实际用户体验。为此引入低成本蒸馏,仅更新少量参数让解码器适应完整模型的 KV cache。这提示评估剪枝不能只看静态基准分数,还需关注生成行为一致性。该策略以极小代价恢复残余精度损失,对工程实践具有指导价值。

方法

输入与阶段划分

SlimWise 面向 混合专家 (MoE) 模型的在线推理服务,输入为批量请求的 prompt 序列。框架将推理显式拆分为 prefill 和 decode 两个阶段:prefill 阶段使用 完整 MoE 模型(所有专家可用),decode 阶段使用 专家剪枝后的模型(只保留部分专家)。这一划分基于观察:prefill 是计算密集,专家剪枝几乎不带来吞吐收益;decode 是内存密集,专家权重流量是瓶颈,剪枝能显著降低权重加载开销。

关键模块:KV Cache 直接交接

核心创新在于 KV cache 的无转换复用。prefill 由完整模型生成 KV cache,decode 的剪枝模型 直接读取这些 KV cache,无需任何投影或对齐操作。这是因为剪枝不改变 hidden dimension 或 attention 结构,只减少 FFN 专家数量,因此 KV 表示天然兼容。实验表明,这种 training-free 的交接在多数设置下大幅缩小了剪枝模型与完整模型的精度差。

蒸馏微调:弥补残余损失

针对 benchmark 精度可能掩盖的生成长度变化,以及仍存在的少量精度损失,SlimWise 引入 低成本的蒸馏阶段:

  • 训练目标:让剪枝后的 decoder 学会从完整模型生成的 KV cache 继续生成,使输出分布逼近完整模型;
  • 参数效率:只更新 一小部分参数(例如 adapter 或特定层),避免全量微调开销;
  • 实现位置:在 vLLM 中集成,兼容 prefill–decode 分离 (PD disaggregation) 和 PD 共置 两种部署模式。

输出与效果

最终输出为 decode 阶段逐 token 生成的文本。在 Qwen3.6-35B-A3B 上,50% 专家剪枝下 decode 吞吐提升最高 1.81×,且精度损失最小。

与同类工作的差异:传统专家剪枝在 prefill 和 decode 阶段统一剪枝,牺牲 prefill 质量换不来解码吞吐收益;SlimWise 首次将专家池按推理阶段解耦,并在不转换 KV cache 的前提下实现阶段间模型切换。

实验

实验设计

  • MoE backbone: 两种 MoE 架构(摘要未列出具体型号,正文提到 Qwen3.6-35B-A3B 用于吞吐评测)
  • 剪枝准则: 三种不同的专家剪枝方法
  • SlimWise 框架: prefill 使用 full model,decode 使用 pruned model 并直接复用 prefill 生成的 KV cache(无需转换)
  • 基线: 传统剪枝同时作用于 prefill 与 decode;以及 full model 不剪枝
  • 蒸馏: 低成本蒸馏阶段,仅更新少量参数,训练 decoder 从 full-model KV cache 继续生成

关键发现

  1. 训练无关的 KV cache 交接在多数设定下显著缩小与 full model 的精度差距。
  2. benchmark 精度可能掩盖剪枝导致的生成长度变化。
  3. 加入蒸馏后,残余精度损失进一步恢复。
  4. 在 Qwen3.6-35B-A3B 上 50% 专家剪枝时,decode 吞吐最高提升 1.81×,且精度损失最小。

与基线对比解读

传统剪枝把 prefill 和 decode 一视同仁,但 prefill 是 compute-bound,剪枝对吞吐帮助不大却严重伤精度。SlimWise 将剪枝限制在 decode,prefill 保持全模型从而保证 KV cache 质量;decode 阶段专家权重流量是瓶颈,剪枝直接降低访存。这一解耦比统一剪枝在精度-吞吐权衡上更优。

行业影响

落地场景

SlimWise 直接适用于大规模 MoE 推理服务,尤其适合高并发 decode 场景:对话式 AI、代码补全、内容生成平台、推荐系统特征交互等。任何采用 vLLM 或类似 serving stack 的团队,若已在用 prefill–decode disaggregation,即可无缝接入。

商业价值

核心收益在降本增效:decode 阶段 expert pruning 减少专家权重读取流量,直接降低存储带宽压力与 GPU 内存占用,使相同硬件能承载更高并发;在 Qwen3.6-35B-A3B 上 50% 剪枝可提升 decode throughput 最高 1.81×,意味着单位推理成本下降接近一半。同时训练免剪枝的 KV cache 交接 + 轻量蒸馏,保持模型质量,避免因准确率下降导致用户流失。

与现有工作流接口

SlimWise 提供 drop-in 能力:

  • 兼容 vLLM,支持 PD 分离和 PD 同置两种部署模式;
  • 无需转换 KV cache,prefill 用全量模型、decode 用剪枝模型,对上层 API 透明;
  • 剪枝标准可选(3 种),蒸馏阶段仅更新少量参数,训练成本低;

具体 use case

  1. 电商搜索/推荐:实时对话式购物助手需低延迟、高吞吐 decode。SlimWise 可在不牺牲回复质量的前提下,将 decode 吞吐提升约 1.8×,帮助平台在促销峰值期应对激增的并发查询。
  2. 企业知识库问答:企业服务中的批量文档摘要、客服机器人对长序列生成需求大,decode expert 流量是主要瓶颈。采用 SlimWise 后,相同 GPU 集群可服务更多租户,降低每 token 成本。

关键差异:与全局 expert pruning 不同,SlimWise 保留 prefill 全量计算,避免 compute-bound 阶段被无谓剪枝,从而在吞吐收益与模型质量之间取得更优平衡。

局限

  • **蒸馏阶段引入额外开销与数据依赖**:尽管作者强调低成本和仅更新少量参数,但该蒸馏仍需使用特定数据集进行训练,且这些数据可能无法覆盖所有下游任务分布。仅更新解码器子集参数的做法可能限制模型对剪枝后容量损失的补偿能力,尤其在需要强泛化能力的复杂任务上,剪枝带来的精度下降可能仍然存在。此外,该方法在无蒸馏模式下(training-free)的精度差距虽被缩小,但并未完全消除,对于对精度要求极高的场景仍需权衡。
  • **prefill 阶段开销未优化**:SlimWise 在 prefill 阶段使用完整模型,因此整体吞吐提升受限于 prefill 与 decode 的时间占比。论文仅报告 decode 吞吐最高提升 1.81×(在 50% 剪枝下),并未给出端到端延迟或吞吐的改善数据。对于长 prompt、低并发或 prefill 密集型负载,计算瓶颈可能转移至 prefill,此时本方法的收益会大幅缩水。评估仅在一个模型(Qwen3.6-35B-A3B)和有限的工作负载上展示,泛化性有待验证。
  • **部署复杂度与框架耦合**:SlimWise 依赖 prefill 与 decode 模型共享相同的隐藏维度和 KV cache 布局,这限制了剪枝策略的灵活性——激进的专家剪枝若改变模型结构(例如专家数量变化导致维度不一致)可能无法直接复用 KV cache。此外,实现紧密耦合于 vLLM,尚未展示在其他 serving 框架中的可移植性。作者也指出基准测分可能掩盖生成长度的变化,这一现象虽被识别,但论文中并未提供系统性的解决方案或评测指标,实际应用中仍需额外关注。
论文Gunho Park2026-09-28原文

相关内容