4.5 MoE推理优化：条件计算与批处理策略
文档摘要
4.5 MoE推理优化：条件计算与批处理策略 引言 MoE模型的推理优化与标准稠密模型有根本性不同。虽然MoE的每token计算量与等效小模型相当，但MoE引入了独特的性能挑战：路由决策的分支性、专家调度的复杂性、KV Cache的额外开销等。本章将系统性地介绍MoE推理优化的核心技术——条件计算和批处理策略。 MoE推理的性能特征 理论计算量分析 MoE推理的计算量由两部分组成： 1. 非条件计算（所有token共享）： 注意力层：$O(T^2 d)$（自注意力） 门控网络：$O(TdN)$（路由计算） LayerNorm + Residual：$O(Td)$ 2.
4.5 MoE推理优化：条件计算与批处理策略
引言
MoE模型的推理优化与标准稠密模型有根本性不同。虽然MoE的每token计算量与等效小模型相当，但MoE引入了独特的性能挑战：路由决策的分支性、专家调度的复杂性、KV Cache的额外开销等。本章将系统性地介绍MoE推理优化的核心技术——条件计算和批处理策略。
MoE推理的性能特征
理论计算量分析
MoE推理的计算量由两部分组成：
1. 非条件计算（所有token共享）：
- 注意力层：O(T^2 d)（自注意力）
- 门控网络：O(TdN)（路由计算）
- LayerNorm + Residual：O(Td)
2. 条件计算（仅对被选中专家）：
- 专家FFN：O(KTh^2)（K个专家各处理T个token）
总计算量对比：
当 T \gg d（长序列）时，注意力计算占主导，MoE优势不明显。当 T \approx d（短序列）时，MoE的计算效率优势显著。
内存访问模式
MoE推理的内存访问模式比稠密模型更复杂：
- 分散读取：不同token需要读取不同专家的参数
- 动态调度：每个token的路由结果在运行时才能确定
- 非连续存储：专家参数在显存中可能不连续
这导致GPU的内存带宽成为瓶颈，而非计算能力。
条件计算优化
专家参数分块加载
将每个专家的参数存储为独立的张量块：
仅加载被选中专家的参数到SRAM，避免加载全部专家参数。
实现要点：
# 伪代码：条件专家计算 for batch_idx in range(B): for seq_idx in range(T): experts = router(x[batch_idx, seq_idx]) # Top-K选择 for expert_id in experts: output[batch_idx, seq_idx] += weight * expert_forward( expert_params[expert_id], x[batch_idx, seq_idx] )
算子融合优化
将路由计算与专家计算融合为单个算子，减少GPU核启动开销：
融合后的算子内部完成：
- 计算门控分数 g = G(x)
- 确定Top-K专家
- 加载对应专家参数
- 执行专家计算
- 加权求和输出
性能提升：减少2-3次GPU核启动，减少中间结果的显存读写。
专家共享与缓存
当连续请求的路由模式相似时，可以缓存最近使用的专家参数：
缓存命中时直接使用缓存参数，避免重新加载。
批处理优化策略
按专家分组的批处理
将一个批次中路由到同一专家的所有token分组，形成"专家批次"：
对每个专家批次执行批量矩阵乘法，充分利用GPU的并行计算能力。
性能分析：
- 当专家批次大小足够大时（|\text{ExpertBatch}_i| \geq 32），GPU利用率高
- 当专家批次很小时（|\text{ExpertBatch}_i| \leq 4），GPU利用率低
批次内填充（Intra-batch Padding）
当专家批次大小不一时，用零填充到最大批次大小：
填充后可以并行计算所有专家，但浪费了部分计算。
动态批次调度
根据历史请求模式动态调整批次组成，最大化专家批次的利用率：
Expert Choice推理的特殊优化
Expert Choice路由在推理时天然适合批处理，因为每个专家处理固定数量的token：
所有专家批次大小相同，无需填充，GPU利用率最高。
KV Cache优化
MoE的KV Cache问题
MoE模型的KV Cache与稠密模型相同大小（每token O(d)），因为KV Cache只来自注意力层。但MoE的生成速度更快（每步计算量更少），导致KV Cache的读写频率更高。
分页KV Cache
使用分页机制管理KV Cache，避免碎片化：
每个token的KV存储为一个固定大小的"页面"，可以分散在显存的不同位置。
注意力层量化
MoE的注意力层计算量占比较高，可以对其进行量化优化：
- KV Cache使用INT8/FP8存储
- 注意力分数计算使用BF16
- Softmax输入使用高精度，输出使用低精度
量化与蒸馏优化
专家参数量化
由于每个专家只处理部分输入，可以激进地对专家参数量化：
- 权重量化：INT4甚至INT2（每个专家独立量化）
- 激活量化：FP8或INT8
- 混合精度：热门专家使用高精度，冷门专家使用低精度
知识蒸馏
将MoE模型蒸馏为更小的稠密模型或稀疏MoE：
教师模型（大型MoE）的知识通过软标签传递给学生模型。
推理框架选型
主流MoE推理框架对比
| 框架 | MoE支持 | 性能优化 | 部署难度 | 推荐场景 |
|---|---|---|---|---|
| vLLM | ✅ 成熟 | ★★★★☆ | ★★★☆☆ | 服务端部署 |
| TensorRT-LLM | ✅ 成熟 | ★★★★★ | ★★☆☆☆ | 高性能服务端 |
| llama.cpp | ✅ 基础 | ★★★☆☆ | ★★★★★ | 本地/CPU推理 |
| DeepSpeed | ✅ 良好 | ★★★★☆ | ★★★☆☆ | 分布式推理 |
| SGLang | ✅ 良好 | ★★★★☆ | ★★★☆☆ | 多模型服务 |
服务端部署最佳实践
- 预热：先处理若干dummy请求，确保所有专家参数已加载到GPU
- 连续批处理：使用continuous batching最大化GPU利用率
- 专家缓存：保持最近使用的专家参数在快速缓存中
- 监控：跟踪每请求的延迟和专家利用率
- 自动扩展：根据负载动态调整实例数量
本章小结
MoE推理优化需要在条件计算的高效性和批处理的并行性之间找到平衡。按专家分组的批处理、算子融合、量化压缩和KV Cache优化是四个核心技术方向。在实践中，结合框架的内置优化（如vLLM的MoE支持）和自定义的调度策略，可以在保持MoE模型质量的同时实现接近小模型的推理速度。