PuzzleMoE：通过稀疏专家合并和位打包推理实现大型混合专家模型的高效压缩

# PuzzleMoE: 通过稀疏专家合并与位打包推理实现大型混合专家模型的高效压缩 来源：https://supercomputing-system-ai-lab.github.io/projects/puzzlemoe/ ## 方法：PuzzleMoE PuzzleMoE通过精心设计的成对双掩码合并算法以及与位级打包技术的系统协同设计，实现了高效的细粒度稀疏专家合并。这两种设计能够在压缩后有效保持MoE模型的性能，同时最小化推理阶段的掩码开销。 ### 算法设计：成对双掩码专家合并 SFPO算法图2：细粒度稀疏专家合并算法概述我们合并一对专家*i*和*j*的步骤为：(i) 计算每个专家的显著性掩码 \\(\mathbf{M}^{\mathrm{sal}}_i\\) 和 \\(\mathbf{M}^{\mathrm{sal}}_j\\) 以保留高重要性权重；(ii) 构建跨专家相似性掩码 \\(\mathbf{M}^{\mathrm{sim}}\\) 以保留对齐的权重；(iii) 提取符号矩阵 \\(\mathbf{S}_i\\) 和 \\(\mathbf{S}_j\\)。合并后的幅值 \\(\mathbf{W}_{\mathrm{merged}}\\) 按元素定义为： \\[ \mathbf{W}_{\mathrm{merged}} = \mathbf{M}^{\mathrm{sim}} \odot \tfrac{|\mathbf{W}_i| + |\mathbf{W}_j|}{2} \;+\; \bigl(\mathbf{1} - \mathbf{M}^{\mathrm{sim}}\bigr) \odot \bigl(\mathbf{M}^{\mathrm{sal}}_{i}\odot|\mathbf{W}_i| + \mathbf{M}^{\mathrm{sal}}_{j}\odot|\mathbf{W}_j|\bigr), \\] 其中 \\(\odot\\) 表示逐元素乘法，\\(\mathbf{1}\\) 表示形状匹配的全1矩阵。最后，我们存储 \\(\{\mathbf{M}^{\mathrm{sal}}_{i}, \mathbf{M}^{\mathrm{sal}}_{j}, \mathbf{S}_{i}, \mathbf{S}_{j}, \mathbf{W}_{\mathrm{merged}}\}\\) 用于推理。 ### 内核设计：基于位打包的高效推理 位打包内核设计图3：掩码打包过程示意图。(a): Mixtral-8x7B中Bfloat16权重指数的分布。经过移位后，指数可用5位编码。(b): 打包后的Bfloat16格式中掩码和符号的位级组织。**观察。** Bfloat16数据格式（1位符号、8位指数、7位尾数）是LLM推理的标准格式。我们发现MoE大语言模型并未充分利用可用的8位指数范围：专家权重的指数集中在一个狭窄的频带内，导致大量动态范围未被使用。对于Mixtral-8×7B（图3(a)），大多数指数落在112到128之间。 **位打包操作。** 利用集中的指数分布，我们应用固定移位将所有指数映射到5位范围内，如图3(b)所示。具体而言，任何小于112的指数都向上取整为112，然后将所有指数减去112，得到的值落在0到31的范围内。此移位操作释放了原始8位指数字段中的3位，可用于打包二进制掩码位和一位符号位。因此，结果 \$\mathbf{W}_{merged}\$ 以标准Bfloat16数据格式存储，有效地将掩码和符号嵌入其中，无需额外存储。 位打包解码设计**即时解码。** 为了将位打包方案转化为实际的推理加速，我们在Triton中实现了一个专门的decode-GEMM内核。按照算法1，每个权重在用于乘法之前立即从其打包格式中即时解码。这避免了在内存中实例化解码后的矩阵：解码逻辑是一种轻量级的原地操作，它寄生在内核现有的数据加载路径上，增加了可忽略的开销。 **融合双专家计算。** 该内核还利用合并专家的结构来进一步节省计算。如图4所示，传统GEMM需要两次独立的内存事务来加载每个专家的权重。我们的内核则只需一次访问打包后的 \$\mathbf{W}_{merged}\$，即可从拼接的输入 \$A_i, A_j\$ 产生两个输出 \$O_i, O_j\$，从而在LLM解码的内存受限GEMM中实现理论上的2倍延迟降低。 位打包解码GEMM设计图4：PuzzleMoE decode-GEMM内核概述。(a): 原始工作流程需要独立访问两个专家。(b): PuzzleMoE decode-GEMM工作流程只需一次访问合并后的专家即可计算两个输出。 ## 更高效的压缩与推理 效率图6：效率分析。(a): Mixtral和Deepseek-MoE的压缩时间对比。(b): Mixtral-8x7B在不同批大小下的吞吐量对比。**压缩成本。** (a) PuzzleMoE压缩Mixtral-8x7B仅需2分钟，而D2由于SVD操作计算成本高需要55分钟。对于Deepseek-MoE，PuzzleMoE需要10分钟进行压缩，这表明PuzzleMoE在具有大量专家的MoE模型上具有高效性。请注意，NAEE依赖穷举搜索，使得压缩时间非常昂贵。例如，将其应用于具有64个专家的DeepSeek-MoE需要大约 \$10^{18}\$ 量级的前向传播，在实践中不可行。 **端到端吞吐量。** 我们在一个配备8块RTX5090 GPU的集群上评测了PuzzleMoE的端到端吞吐量，Mixtral-8x7B在50%压缩比下的结果如图6(b)所示。在16个预填充令牌和128个解码令牌的配置下，对于不同批大小，PuzzleMoE相比未压缩基线实现了高达1.80倍的加速。这一效率提升主要得益于我们专门的PuzzleMoE decode-GEMM内核，该内核相比标准cuBLAS实现提供了2倍的吞吐量提升。 ## BibTeX `` @misc{zhao2025puzzlemoeefficientcompressionlarge, title={PuzzleMoE: Efficient Compression of Large Mixture-of-Experts Models via Sparse Expert Merging and Bit-packed inference}, author={Yushu Zhao and Zheng Wang and Minjia Zhang}, year={2025}, eprint={2511.04805}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2511.04805}, } ``