动态量化优化MoE模型服务
研究背景
大语言模型（LLM）的规模化趋势中，混合专家模型（Mixture-of-Experts, MoE）凭借其“条件计算”特性，在保持高精度的同时显著降低了推理成本，成为主流架构之一。然而，在实际服务场景中，MoE 面临一个日益尖锐的矛盾：模型权重与 KV Cache 之间的 GPU 显存竞争。KV Cache 是自回归生成中存储中间 Key-Value 张量的缓存结构，其大小随序列长度和并发请求数线性增长。在长上下文或高吞吐服务中，KV Cache 可能占据显存大头，挤压模型权重的驻留空间，导致频繁的显存换出或精度下降。
传统解决方案包括权重静态量化（如 INT8、INT4）或 KV Cache 动态卸载，但这些方法存在固有局限：静态量化往往对 MoE 中不同专家（Expert）的重要性一视同仁，导致关键专家精度损失不可控；而 KV Cache 卸载则受限于 PCIe 带宽，带来显著的延迟开销。MoE 模型的核心特性——专家之间的异构性和稀疏激活模式——为更精细的权衡提供了可能。PagedWeight 正是针对这一痛点，提出一种运行时动态量化策略，在精度、显存占用和吞吐量之间实现帕累托最优。
核心思想
PagedWeight 的核心洞察是：MoE 模型中不同专家对任务精度的影响存在显著差异，且这种重要性在推理过程中会动态变化。因此，与其对所有权重采用统一量化精度，不如根据实时 KV Cache 压力和专家重要性，动态调整每个专家的量化比特数。这类似于操作系统的页面置换策略——将“珍贵”的显存资源分配给更关键的专家权重，而牺牲非关键专家的精度以腾出空间给 KV Cache。
具体而言，PagedWeight 将 MoE 模型的权重视为可动态分页管理的资源：每个专家权重被划分为固定大小的“页”，每页可以独立选择量化精度（如 FP16、INT8、INT4）。系统运行时监控 KV Cache 占用率，当显存紧张时，自动降低低重要性专家的量化精度（即从 FP16 降为 INT8 或 INT4），从而释放显存给 KV Cache；反之，当显存充裕时，恢复高精度以维持模型质量。这一过程无需离线重训练，完全在推理时在线执行。
方法与技术路线
PagedWeight 的实现包含三个关键组件：
专家重要性评估：传统方法往往基于专家激活频率或梯度幅值评估重要性，但 PagedWeight 引入了一种“质量感知”度量——通过测量专家权重对最终输出对数概率（logits）的贡献方差来量化其重要性。具体地，对于每个专家 $E_i$，定义其重要性分数 $I_i$ 为：
$$I_i = \mathbb{E}_{x \sim \mathcal{D}} \left[ \text{Var}\left( \frac{\partial \mathcal{L}}{\partial E_i(x)} \right) \right]$$
其中 $\mathcal{L}$ 是当前任务损失。该指标能捕捉专家对生成质量的边际影响，且计算开销通过采样近似控制。动态精度分配机制：系统维护一个显存预算 $M_{\text{total}}$，包含模型权重和 KV Cache。设当前 KV Cache 占用为 $M_{\text{kv}}$，则权重可用显存为 $M_{\text{w}} = M_{\text{total}} - M_{\text{kv}}$。PagedWeight 将权重分配问题建模为整数线性规划：
$$\max \sum_{i=1}^{N} I_i \cdot Q_i \quad \text{s.t.} \quad \sum_{i=1}^{N} S(Q_i) \leq M_{\text{w}}$$
其中 $Q_i \in {FP16, INT8, INT4}$ 是专家 $i$ 的量化选择，$S(Q_i)$ 为对应的显存占用。由于专家数量较大（通常数十到数百），PagedWeight 采用贪心近似：按重要性降序排列专家，优先为高重要性专家分配高精度，直到显存耗尽。这一决策在每次 KV Cache 大小变化时触发（如每生成一个新 token）。分页与零拷贝量化：为避免频繁的精度切换导致额外显存带宽开销，PagedWeight 采用页式存储：每个专家权重被分割为 1MB 的固定页，每页独立存储为 FP16 格式，但量化操作通过“虚拟页表”实现——仅修改页表的精度标记，实际数据在 GPU 内核中通过指针转换进行“零拷贝”读取。例如，一个 INT4 页在内存中仍占用 FP16 空间，但计算时按 INT4 解释并累加。这避免了数据重排，将量化开销降至几乎为零。
意义与影响
PagedWeight 的实验结果极具说服力：在多种显存敏感的 MoE 服务场景下（如长文本生成、批量推理），它实现了 FP16 等效精度下高达 72.0% 的显存节省和 1.94 倍的吞吐量提升；与同等显存预算的静态量化方法相比，模型质量提升最高达 39.3%，而吞吐损失仅不超过 4.1%。这一性能突破源于其动态权衡机制：当 KV Cache 压力大时，牺牲低重要性专家的精度（通常占专家总数的 60%-80%）来释放显存，而关键专家的精度得以保持。
从工程角度看，PagedWeight 提供了一种即插即用的解决方案：无需修改模型架构或重新训练，仅需在推理框架中集成其调度模块。这使其对现有 MoE 服务系统（如 Mixtral 8x7B、DeepSeek-MoE）具有直接迁移价值。此外，其“质量感知”思想可推广至其他显存动态分配场景，例如注意力头的自适应精度控制。
局限性
尽管 PagedWeight 表现出色，仍存在若干局限。首先，专家重要性评估依赖采样数据分布 $\mathcal{D}$，若服务中的任务分布剧烈偏移（如从代码生成突然切换为科学文献），重要性分数可能滞后，导致精度分配次优。其次，贪心分配算法虽高效，但无法保证全局最优——当专家重要性差异不大时，可能错过更优的组合。此外，零拷贝量化依赖于 GPU 对非对齐内存访问的支持，在部分低端硬件上可能引入额外的访问延迟。最后，PagedWeight 目前仅针对 MoE 架构，无法直接用于稠密模型；但作者指出，其思想可扩展至层间或注意力头间的精度分配，这将是未来工作的方向。