Progressive Mixed-Precision Decoding for Efficient LLM Inference

标题：用于高效 LLM 推理的渐进式混合精度解码
作者/机构：Hao (Mark) Chen, Fuwen Tan, Alexandros Kouris, Royson Lee, Hongxiang Fan, Stylianos I. Venieris (Samsung AI Center, Cambridge, UK; Imperial College London, UK)

A1 主要贡献

核心问题与研究目标：大型语言模型（LLMs）因其巨大的计算和内存需求，在资源受限的设备上部署面临挑战。量化是一种有效的解决方案，但极低精度（如2/3-bit）的训练后量化（PTQ）会导致严重的性能下降。现有方法未能探索 LLM 推理不同阶段在计算模式、冗余度和对近似的敏感性方面的多样性，而是采用统一的量化策略。本文旨在通过一种新颖的阶段感知方法，在 LLM 推理的不同阶段选择性地分配精度，以实现预填充（prefill）阶段强大的上下文提取和解码（decoding）阶段高效的内存带宽利用。

创新点：
本文基于一个核心洞察：LLM 推理的预填充阶段以及解码阶段的早期部分，比自回归生成过程的后期部分对量化等近似误差更敏感。基于此，本文提出了一种新颖的、与现有统一量化方法不同的推理方法，其主要贡献如下：

  • 一种新颖的阶段感知精度分配策略：该策略针对预填充和解码阶段的不同特性进行精度优化，利用每个阶段独特的错误容忍度，为移动级 LLM 实现极低的平均位宽。
  • 一种渐进式混合精度解码方案（PMPD）：该方案随着解码过程生成更长的 token 序列而逐步降低精度，从而有效提高内存受限的 LLM 解码阶段的硬件性能。
  • 两种精度切换调度器：为实现灵活部署，设计了一种与提示（prompt）无关的静态调度器和一种与任务无关的学习型调度器，以适应不同的数据可用性约束和生成质量要求。

该方法通过在预填充阶段使用高精度保证生成质量，在解码阶段使用逐步降低的精度来提升吞吐量。

图1：传统量化LLM推理范式（左）与本文方法（右）的图示，本文方法包括阶段感知精度分配和渐进式混合精度解码。我们的方法基于两个动机：i) 在预填充和解码阶段观察到的不同错误容忍度，以及 ii) 随着解码向更长的 token 发展，容错能力增强。
图1：传统量化LLM推理范式（左）与本文方法（右）的图示，本文方法包括阶段感知精度分配和渐进式混合精度解码。我们的方法基于两个动机：i) 在预填充和解码阶段观察到的不同错误容忍度，以及 ii) 随着解码向更长的 token 发展，容错能力增强。

A3 背景知识与关键观察

背景与相关工作

低精度 LLM 推理：多种量化方法已被用于降低 LLM 的权重、激活值和 KV 缓存的精度以提升计算效率。
* 纯权重（Weight-only）量化：这类方法通过减轻解码过程中的内存事务量（主要由权重传输主导），在商用平台上部署时能显著提速。例如，Any-Precision LLM【【24】，Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs，2024，ICML】可以在无额外存储开销的情况下，将单组权重进行多精度量化。DNS（Dense-and-Sparse Quantization）【【17】，SqueezeLLM: Dense-and-Sparse Quantization，2024，ICML】通过以全精度存储一小部分离群值权重来提高平均位宽。BitNet【【23】，The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits，2024，arXiv】引入了三元权重，但需要昂贵的量化感知训练。本文则专注于低开销的训练后量化。现有方法大多在所有阶段和提示中采用统一精度，这在对小型模型进行激进量化（如2或3位）时难以维持性能。本文方法利用 Any-Precision LLM 提供的低内存占用方案，但采用动态降低精度的方法，考虑了 LLM 推理阶段和具体输入提示，以突破精度限制。

精度自适应方法：本文方法通过动态的、阶段感知的渐进式精度降低方法来利用权重化。
* 精度自适应训练：MUPPET【【26】，Multi-Precision Policy Enforced Training (MuPPET): A Precision-Switching strategy for Quantised Fixed-Point Training of CNNs，2020，ICML】、CPT【【13】，CPT: Efficient Deep Neural Network Training via Cyclic Precision，2021，ICLR】和 AdaPT【【19】，Adaptive Precision Training (AdaPT): A dynamic quantized training approach for DNNs，2023，SDM】等方法通过在训练过程中动态调整精度来提高训练效率，主要关注分类任务。本文方法则针对 LLM 的特性进行了优化，专注于提高推理效率。
* 混合精度推理：HAQ【【28】，HAQ: Hardware-Aware Automated Quantization with Mixed Precision，2019，CVPR】和 HAWQ【【10】，HAWQ: Hessian Aware Quantization of Neural Networks with Mixed-Precision，2019，ICCV】允许在神经网络的不同层使用不同精度，但所选精度对所有输入样本保持固定。Bit-mixer【【4】，Bit-mixer: Mixed-precision networks with runtime bitwidth selection，2021，ICCV】允许跨层混合精度，并能根据每个样本调整位宽。CascadeCNN【【18】，CascadeCNN: Pushing the Performance Limits of Quantisation in Convolutional Neural Networks，2018，FPL】采用级联的不同精度模型变体。这些方法主要为基于 CNN 的分类任务设计，不直接适用于 LLM 的自回归解码过程。与本文更相关的是 LLM-PQ【【34】，POSTER: LLM-PQ: Serving LLM on Heterogeneous Clusters with Phase-Aware Partition and Adaptive Quantization，2024，PPoPP】，它根据不同服务器的硬件支持调整模型精度，并进行阶段感知的模型分区，目标是分布式 LLM 推理服务。相比之下，本文方法是在不同时间步而非不同机器上调整精度，并为边缘部署进行了优化。

关键观察

观察一：预填充（Prefill）和解码（Decoding）阶段存在不同的错误容忍度。受近期研究【【3】，Tandem Transformers for Inference Efficient LLMs，2024，ICML】启发，即为预填充阶段分配更多模型容量可以减少解码阶段的容量需求，我们研究了 LLM 推理中预填充和解码阶段表现出的不同冗余度和容忍度。为此，我们针对三个不同的语言任务，比较了 Vicuna-7B 模型的两个变体：i) 一个统一使用2位量化权重的低精度变体，和 ii) 一个在预填充阶段使用3位权重、在解码阶段使用2位权重的双精度变体。实验结果（图2）表明，在预填充阶段使用3位模型显著增强了统一2位模型的生成能力。我们认为这是因为预填充阶段对输出生成过程具有重要的引导作用。具体来说，在预填充阶段从输入提示中提取的高质量 KV 缓存提供了对任务的良好理解和信息丰富的上下文，从而增强了低精度模型的生成能力。这一洞察激励我们为预填充和解码阶段探索不同的精度策略。

图2：阶段感知精度分配。此处预填充由2位（上）或3位（下）Vicuna-7B执行，而解码对两种情况都采用了2位模型。我们观察到三个关键改进。i) 增强的推理能力。模型在数值任务（如计算斐波那契数）上表现出更高的准确性（左）。ii) 改进的指令遵循能力。模型理解并响应了用户指令（中）。iii) 强化的新兴能力，如创造力。模型成功处理了开放式任务，如写诗（右）。更多量化结果见附录A.2和图6。
图2：阶段感知精度分配。此处预填充由2位（上）或3位（下）Vicuna-7B执行，而解码对两种情况都采用了2位模型。我们观察到三个关键改进。i) 增强的推理能力。模型在数值任务（如计算斐波那契数）上表现出更高的准确性（左）。ii) 改进的指令遵循能力。模型理解并响应了用户指令（中）。iii) 强化的新兴能力，如创造力。模型成功处理了开放式任务，如写诗（右）。更多量化结果见附录A.2和图6。

观察二：在解码过程中渐进式降低精度。为了进一步提高降精度生成的输出质量，我们探索了混合精度解码的优化空间。图3展示了我们应用不同低、高精度权重调度方法时的发现。具体来说，我们探索了在生成 token 的前半部分、中间部分和后半部分使用高精度权重。与我们的假设一致，在解码阶段的前半部分采用高精度产生了最佳性能，其表现与在整个生成过程中使用高精度的性能相匹配。一个直观的解释是，使用高精度生成最初的几个 token，可以最大限度地减少会影响后续推理的误差累积。这一观察也与“注意力池”（attention sink）现象【【29】，Efficient streaming language models with attention sinks，2024，ICLR】相符，该现象观察到注意力分数严重偏向初始 token，表明它们的重要性更高。作为参考，我们还测试了在每个 token 交替使用两种模型的性能，结果也很好。然而，频繁的精度切换会引入硬件开销，包括权重加载时间和额外的内存流量，因此这种调度方案在实际部署中并不可行。

图3：混合3位和4位模型进行解码的不同调度方案。我们使用Phi-1.5在CNN/DM上进行了文本摘要测试。为确保公平比较，4位基线生成了相同数量的 token。此外，所有调度方案的预填充都使用了4位模型。
图3：混合3位和4位模型进行解码的不同调度方案。我们使用Phi-1.5在CNN/DM上进行了文本摘要测试。为确保公平比较，4位基线生成了相同数量的 token。此外，所有调度方案的预填充都使用了4位模型。

A2 方法细节

方法论总览。基于第3节的发现，即解码序列中位置更深的 token 对近似的容忍度更高，我们采用权重化作为近似机制，并提出了一种新的混合精度解码方法，包含两个关键技术：i) 阶段感知精度分配和 ii) 渐进式混合精度解码（PMPD）。为了在保持生成质量的同时最大化推理效率，这两项技术互为补充，应用针对 LLM 预填充和解码阶段各自特性的精度选择和调度策略。为了实现对我们方法的细粒度控制，我们进一步引入了两种运行时调度器，分别以与提示无关和与任务无关的方式驱动精度切换决策。

端到端方法流程。算法1展示了我们的端到端方法，包括：i) 离线阶段（1-8行）和 ii) 部署阶段（9-22行）。离线阶段负责生成一个给定 LLM 的可变精度量化变体（第1行），进行阶段感知精度分配（第2行），并为目标用例配置所选类型的精度切换调度器及优化的 PMPD 参数（第3-8行）。在部署时，调度器为预填充和解码阶段分配相应的精度，并为每个遇到的提示触发精度切换事件（第20行）。

算法 1: 渐进式混合精度解码

输入: LLM m，最大输出长度 OL
      精度集 P, 量化器 Q(·)
      任务无关校准集 Dcalib, 任务相关验证集 V (如果使用与提示无关的静态调度器则为可选)
      参考质量 qref, 质量下降容忍度 ϵ
输出: 输出 token 序列 (t0, t1, ...)

/* --- 离线阶段 --- */
1  {mp} ← Q(m, p) ∀p ∈ P                                  ▷ 获取可变精度量化的模型变体
2  pprefill, pdecode ← PAPAlloc(P, Dcalib, qref, ϵ)          ▷ 阶段感知精度分配
3  if type(PMPScheduler) is prompt-agnostic static then
4     PMPDScheduler.S ← OfflineSolver({mp}, V, qref, ϵ)    ▷ 离线优化公式(1)
5  end
6  if type(PMPScheduler) is task-agnostic learned then
7     PMPDScheduler.model ← SchedulerTrainer({mp}, Dcalib) ▷ 训练学习型调度器
8  end

/* --- 部署阶段 --- */
9  while not prompt queue is empty do                        ▷ 处理传入的提示
10    d0, K0V0 ← m_pprefill(prompt)                          ▷ 预填充阶段
11    t0 ← Sampler(d0)
12    if type(PMPScheduler) is task-agnostic learned then
13       PMPDScheduler.S ← PMPDScheduler.model(K0V0)         ▷ 在解码前生成精度切换调度
14    end
15    pnew ← pdecode
16    for i ← 0 to OL − 1 do                                ▷ 解码阶段
17       di+1, Ki+1Vi+1 ← m_pnew(ti, KiVi)
18       ti+1 ← Sampler(di+1)
19       if ti+1 is EOS then break                           ▷ 序列结束
20       pnew ← PMPDScheduler(i + 1)                         ▷ 精度切换调度器
21    end
22 end

4.1 阶段感知精度分配

核心思想。所提出的阶段感知精度分配利用了每个 LLM 推理阶段的独有特性，并相应地定制了各阶段的权重精度。具体来说，我们的推理方案规定在预填充阶段使用高精度模型，这反过来使得在内存受限的解码阶段能够使用更低精度的权重，从而更有效地利用内存带宽。由于预填充阶段是计算密集型的，而 LLM 推理延迟通常由解码阶段主导，因此在预填充阶段使用高精度模型引入的延迟开销可以忽略不计。

形式化定义。我们的方案通过以下参数进行定义：1) 量化器 $Q(·)$，通常是一种训练后量化（PTQ）方法；2) 一组精度 $P$，表示支持的权重位宽；3) 一个特定任务质量指标中的参考值 $q_{ref}$；以及 4) 质量下降容忍度 $\epsilon$。我们方法的目标是找到用于预填充和解码阶段的最小权重精度对，使其达到的算法性能不低于 $q_{ref} - \epsilon$。

离线校准步骤。我们通过方法论离线阶段的校准步骤来解决这个目标。首先，给定一个预训练的 LLM $m$，我们使用量化器 $Q(·)$ 为精度集 $P$ 中的每个精度 $p$ 获取一组可变精度量化模型 $\{m_p\}$（第1行）。接下来，我们执行阶段感知精度分配步骤（第2行）。具体地，通过在校准数据集上评估不同预填充和解码位宽组合的算法性能，我们确定满足质量目标的最低精度对（第2行）。

4.2 渐进式混合精度解码

核心思想。利用第3.2节的洞察，我们引入了渐进式混合精度解码（PMPD）来进一步提高解码效率。如图1右侧所示，PMPD 利用不同 token 在近似容忍度上的可变性，随着我们深入输出序列而逐渐降低权重的算术精度。通过这种方式，我们为通常对响应质量起决定性作用的早期 token 保持较高精度，同时通过使用更窄位宽的权重来更快地生成更深层的 token。

形式化定义。我们用以下参数来定义 PMPD：i) 精度集 $P$ 和 ii) 精度切换调度 $S = \{st(p) | p \in P\}$，它由每个精度的可能切换点组成，定义为 $st(p) \in [1, OL)$，其中 $OL$ 是给定 LLM 的最大输出长度。这里，起始点 $st(p)$ 定义为我们从当前精度切换到较低精度 $p$ 的输出序列的 token 索引。

优化问题。在这种设置下，为了获得性能最高的 PMPD 配置，我们提出了以下优化问题：

$$\min st(p), \quad \forall p \in \mathcal{P} \setminus \{p_{\min}\}$$ $$\begin{aligned} \begin{aligned} \text{s.t.} \quad & q_{\text{ref}} - \epsilon \le q(S) \quad \& \quad 0 \le st(p) < OL, \quad p \in \mathcal{P} \\ & p > q \implies st(p) \le st(q), \quad p, q \in \mathcal{P} \end{aligned} \end{aligned}$$

其中 $q(S)$ 是调度 $S$ 实现的质量。我们的目标函数旨在最小化在精度集中除最低精度 $p_{min}$ 外的每个精度下处理的 token 数量——最低精度将最后使用直到序列结束以获得更高吞吐量——同时满足指定的质量目标（公式(2)中的第一个约束）、不超过 LLM 的最大上下文长度（公式(2)中的第二个约束），并遵循渐进式降低精度的方法（公式(3)中的优先约束）。

4.3 精度切换调度器

调度器作用。在部署时，PMPD 通过一个精度切换调度器来实施（图4）。该调度器配置了 PMPD 参数 $\langle P, S \rangle$，并根据每个输入提示的调度 $S$ 向系统处理器发出精度切换指令。

图4：精度切换调度器。
图4：精度切换调度器。

搜索空间复杂性。给定在阶段感知精度分配步骤中生成的可变精度量化模型集 $\{m_p\}$（第4.1节和第1行），可以对所有精度切换调度的组合评估公式(1)的目标函数。理想情况下，性能最高的调度会在决定精度切换时，逐个提示、逐个步骤地考虑所有 $p \in P$，而不受公式(3)中的优先约束。这种调度可以通过对 $|P|^{OL}$ 种调度进行穷举枚举来获得。然而，对于实际的 $|P|$ 和 $OL$ 值，这种计算很快变得难以处理。此外，即使可以高效地找到这样的调度，在大多数情况下也是不切实际的，因为它需要过于频繁的切换（例如每个解码步骤），从而加剧了内存带宽需求（第3.2节）。

搜索空间缩减。相反，通过我们解码时渐进式降低精度的方法，待检查的调度总数减少到 $\sum_{r=0}^{|P|-1} \frac{OL!}{r!(OL-r)!} \frac{(|P|-1)!}{r!(|P|-1-r)!}$，其中 $r$ 是精度切换的次数。尽管如此，最优调度的搜索时间随 $OL$ 和 $|P|$ 的增加而增加，这在运行时仍然是禁止的。因此，为了进一步将最优 $S$ 的搜索时间减少到 $O(1)$，我们参数化了候选精度切换点的数量，记为 $N$，并在公式(2)的第二个约束中增加了一个范围约束，即 $st(p) \in \{0\} \cup \{ \frac{OL}{N-1}, \frac{2OL}{N-1}, ..., \frac{(N-2)OL}{N-1}, OL\}$。

两种调度器类型。在这种设置下，我们提出了两种类型的调度器：i) 一种与提示无关的静态调度器，其中 $S$ 在部署前确定；以及 ii) 一种与任务无关的学习型调度器，其中 $S$ 为每个提示动态导出。

与提示无关的静态调度器。为了利用高效的精度调度器，我们提出一种静态方法，该方法通过利用校准集的离线优化过程（第4行）来选择调度 $S$。虽然公式(3)中的优先约束和公式(1)中的优化目标仍然成立，但公式(2)中实现的质量被近似为 $q(S) \approx \frac{1}{|V|} \sum_{x \in V} q_x(S)$，其中 $V$ 是验证集，$q_x(S)$ 是调度在样本 $x$ 上测得的质量分数。由于验证集是按任务定义的，静态调度器保持了与提示无关（即对所有提示使用相同的精度切换调度）和任务特定的特性。因此，这种类型的调度器是根据手头任务的容忍度量身定制的，并且具有几乎不引入运行时开销的优点。然而，通过在给定任务内采用统一的调度，它需要一个能代表任务的验证集，这通常不是一个现实的假设，并且它没有利用不同提示在难度和对近似的容忍度上的可变性。

与任务无关的学习型调度器。为了消除对每个任务验证集的需求，我们引入了一种可训练的调度器，该调度器根据每个输入提示调整调度 $S$。对于一个给定的 LLM，我们的学习型调度器在一个通用的、与任务无关的数据集上进行训练，然后应用于各种下游任务，从而分摊了其训练成本。具体来说，我们的目标是根据每个输入提示的特征设计一个调度。然而，训练调度器做出非常细粒度的每步决策，既需要大量的模型容量，又需要过于频繁的调度器调用，这会引入显著的运行时开销，并抵消 PMPD 的加速优势。

学习型调度器实现。因此，我们利用在预填充阶段提取的特征作为调度器的输入。通过这种方法，我们能够为调度器保持一个轻量级架构，并在推理过程中引入最小的开销，因为调度器在解码前一次性生成调度 $S$（第13行）。在第5.3节中，我们研究了不同的输入特征，即来自模型不同层的 KV 缓存和激活值，并发现利用预填充的 KV 缓存在大多数情况下是有效的。给定输入键缓存 $K \in \mathbb{R}^{T \times D_k}$ 和值缓存 $V \in \mathbb{R}^{T \times D_v}$，学习型调度器的输出计算为 $O = \text{MLP}(\text{Softmax}(\frac{q_{\theta}K^T}{\sqrt{D}})V)$，其中 $q_{\theta} \in \mathbb{R}^{D_k}$ 是一个学习到的查询向量。

学习型调度器性能。我们还发现，尽管不需要特定于任务的验证集和带标签的真实值，我们的学习型调度器在算法和硬件性能上都接近静态调度器，在某些情况下甚至超过了它。其架构和训练细节在附录A.1中有详细说明。

A4 实验环境

  • 模型：实验在适用于边缘部署的模型上进行，包括 Vicuna-7B【【7】，Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90%* ChatGPT Quality，2023】、MobileLLaMA-1.4B【【8】，MobileVLM : A Fast, Strong and Open Vision Language Assistant for Mobile Devices，2023，arXiv】、Stable LM Zephyr-3B 和 Phi-1.5【【22】，Textbooks are all you need ii: phi-1.5 technical report，2023】。
  • 数据集与任务：评估了模型在新闻摘要（CNN/DM【【15】，Teaching Machines to Read and Comprehend，2015，NeurIPS】）、对话摘要（Dialogsum【【6】，DialogSum: A real-life scenario dialogue summarization dataset，2021，ACL-IJCNLP】）和翻译（IWSLT 法语-英语【【5】，Overview of the IWSLT 2017 evaluation campaign，2017，IWSLT】）任务上的零样本生成性能。此外，还在 MT-Bench【【7】，Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90%* ChatGPT Quality，2023】上测试了开放式问答，这是一个静态调度器因缺少验证集而不可行的特殊情况。

  • 硬件配置：

    • GPU：Nvidia RTX 4090 和 A40。使用 PyTorch Profiler 测量了 LLM 中线性层的延迟。
    • NPU：使用一个针对 LLM 优化的加速器 FlightLLM【【33】，FlightLLM: Efficient Large Language Model Inference with a Complete Mapping Flow on FPGAs，2024，ISFPGA】的分析性能模型进行仿真测量。

  • 软件配置：

    • 实现：基于 PyTorch【【25】，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】。
    • 量化方法：采用 Any-Precision LLM【【24】，Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs，2024，ICML】的嵌套量化方法，确保使用多重权重精度不产生内存占用开销。
    • 调度器训练：学习型调度器使用 C4 测试数据集的前256个样本作为种子数据集进行训练。

  • 基线：与单精度量化模型以及使用密集稀疏分解（DNS）【【17】，SqueezeLLM: Dense-and-Sparse Quantization，2024，ICML】增强的量化模型进行比较，DNS 的比率被调整以匹配 PMPD 的平均位宽。

A4 实验结果

算法性能

PMPD 与基线（Baselines）比较（Table 1）：
* 在 CNN/DM 和 Dialogsum 数据集上，静态和学习型调度器均实现了无损性能，同时位宽降低高达33%。在某些情况下，PMPD 的性能甚至比高精度基线高出2.4%。
* 在 IWSLT 数据集上，静态调度器性能有最多2.3%的中度下降，但平均位宽仍减少了9-33%。
* 结果表明，解码阶段并非总是需要高精度模型，静态调度器能够准确识别这种情况并实现最大的位宽降低。

PMPD 与 DNS 比较（Table 1）：
* 在所有任务上，静态调度器在相同位宽下始终优于 DNS，在 IWSLT 数据集上的 BLEU 分数高出29%。
* 对于较小的模型，性能增益更明显：在 Vicuna-7B 的 CNN/DM 任务上，DNS 与静态调度器性能相当，但在 MobileLLaMA 上，DNS 的 Rouge-L 分数落后23%。这表明 PMPD 特别适合边缘部署场景。

静态（Static）与学习型（Learned）调度器比较（Table 1）：
* 平均而言，静态调度器实现了稍高的位宽降低（0.66 vs. 0.63 bits），且性能更稳定，方差更小。这可能是因为它利用了任务特定数据，而学习型调度器可能面临训练集与测试集之间的数据分布偏移挑战。

开放式任务（MT-Bench, Table 2）：
* 在静态调度器不可用的 MT-Bench 任务上，学习型调度器展示了其有效性。它实现了0.47的平均位宽降低，BERTScore 没有下降，Rouge-L 仅下降最多1.4点。
* 学习型调度器始终优于 DNS，Rouge-L 分数高出13.9-18.4%，表明 PMPD 在复杂任务中能更好地保持模型保真度。

表1：静态和学习型调度器与低精度（baseline-l）和高精度（baseline-h）基线以及与调度器性能最高的DNS（具有相同平均位宽）的性能比较。对于CNN/DM和Dialogsum，从上到下使用的模型是Vicuna-7B（2/3位）、MobileLLaMA（3/4位）和Phi-1.5（3/4位）。对于IWSLT，模型是Vicuna-7B、MobileLLaMA和Zephyr-3B（3/4位）。调度器和DNS之间的成对获胜者用下划线标出，而最高的总分用粗体表示。
表1：静态和学习型调度器与低精度（baseline-l）和高精度（baseline-h）基线以及与调度器性能最高的DNS（具有相同平均位宽）的性能比较。对于CNN/DM和Dialogsum，从上到下使用的模型是Vicuna-7B（2/3位）、MobileLLaMA（3/4位）和Phi-1.5（3/4位）。对于IWSLT，模型是Vicuna-7B、MobileLLaMA和Zephyr-3B（3/4位）。调度器和DNS之间的成对获胜者用下划线标出，而最高的总分用粗体表示。
表2：PMPD在MT-Bench上的性能，以fp16输出为参考。
表2：PMPD在MT-Bench上的性能，以fp16输出为参考。

硬件性能

NPU 评估（Figure 5）：
* 在模拟的 LLM 优化 NPU 上，PMPD 相较于 fp16 实现了 3.8-8.0 倍的显著加速。
* 对于更大的 Vicuna-7B 模型，由于其对量化误差的容忍度更高，能够采用更激进的2位精度降低，因此加速效果更明显。
* 尽管带宽有限（NPU 为 32 GB/s），PMPD 仍能为移动级 LLM 实现超过 50 tokens/s 的吞吐量，展示了其在移动设备上部署的有效性。

图5：端到端NPU吞吐量。加速比是与fp16模型比较得出的。
图5：端到端NPU吞吐量。加速比是与fp16模型比较得出的。

GPU 评估（Table 3）：
* 在 GPU 上，PMPD 在线性层上实现了 1.40-12.20 倍的加速（与 fp16 相比）。较小的注意力投影层加速比为 1.40-6.58 倍，较大的前馈投影层加速比为 1.84-12.20 倍。
* 在较小层上加速比较低的原因可能是矩阵向量乘法的内存受限特性。
* 一个瓶颈是 CPU 端的 GPU 内核启动过程，导致 GPU 利用率不佳，可通过 CUDA Graph 等技术优化。

表3：PMPD在线性层上相较于fp16模型的GPU加速比。每个模型展示了两个代表性层：注意力投影（Attn Proj.）和MLP投影（MLP Proj.）。PMPD的加速比是根据两种精度下内核延迟的加权平均值计算的，权重为解码步数。更多细节见附录A.5。
表3：PMPD在线性层上相较于fp16模型的GPU加速比。每个模型展示了两个代表性层：注意力投影（Attn Proj.）和MLP投影（MLP Proj.）。PMPD的加速比是根据两种精度下内核延迟的加权平均值计算的，权重为解码步数。更多细节见附录A.5。

消融研究

阶段感知精度分配（Figure 6）：
* 该策略显著提升了性能：CNN/DM 上的 RougeL 平均提高4.3点，Dialogsum 上的 BERTScore 增加3.6，IWSLT 上的 BLEU 提高9.45。
* 对于 Vicuna-7B 的2位变体，在预填充阶段使用更高精度的3位模型极大地改善了其指令遵循和语言理解能力。
* 使用更高精度预填充的延迟开销极小（0.07% 到 1.05%），证实了预填充阶段是计算密集型的，也证明了该策略的效率。

图6：阶段感知精度分配的性能，Vicuna-7B使用3/2位，其他模型使用4/3位。
图6：阶段感知精度分配的性能，Vicuna-7B使用3/2位，其他模型使用4/3位。

学习型调度器设计（Figure 7）：
* 研究了学习型调度器的不同输入选择，包括来自不同 Transformer 块的激活值和 KV 缓存。
* 实验结果显示，使用最后一个注意力块的 KV 缓存作为输入达到了最高的归一化分数，并实现了显著的位宽降低，证明了其在捕捉上下文信息以预测切换点方面的有效性。
* 该调度器也优于两个随机调度器基线，表明它有效地学习了 KV 缓存中的模式。

图7：在MobileLLaMA上，不同学习型调度器变体在CNN/DM、IWSLT和MT-Bench上的归一化得分。
图7：在MobileLLaMA上，不同学习型调度器变体在CNN/DM、IWSLT和MT-Bench上的归一化得分。

A5 结论

本文利用 LLM 解码过程后期阶段对近似表现出更强容忍度的洞察，引入了 PMPD，一种在推理过程中逐步降低模型权重量化精度的新技术。该方法考虑了 LLM 推理的不同阶段（预填充 vs. 解码）和解码序列的深度，以与提示无关的静态方式或与任务无关的可学习方式来调度精度变化。实验表明，PMPD 能够生成高质量的响应，同时显著降低了平均位宽，从而在 GPU 和 NPU 平台上带来了显著的推理速度提升。

A6 附录

A.1 学习型调度器细节

A.1.1 模型架构。为避免运行时开销，我们为学习型调度器采用了一个轻量级的注意力模块。给定输入键缓存 $K \in \mathbb{R}^{T \times D_k}$ 和值缓存 $V \in \mathbb{R}^{T \times D_v}$，学习型调度器的输出计算如下：

$$\begin{aligned} \begin{aligned} O &= Softmax(\frac{q_{\theta} K^{T}}{\sqrt{D_{k}}}) V \\ Logits &= Softmax(MLP(O)) \end{aligned} \end{aligned}$$

其中 MLP 是一个带有一个隐藏层和 ReLU 激活函数的前馈神经网络，$q_{\theta} \in \mathbb{R}^{D_k}$ 是一个可学习的查询向量。换言之，网络的可训练参数包括查询向量和 MLP 的参数。

A.1.2 训练方法。为了准备训练数据集，我们使用 C4 数据集作为种子数据集，并随机截断每个样本以创建一个补全任务。对于每个样本，我们生成 $N$ 个序列，每个序列对应一个高精度步骤，这些步骤选自 $\{0, \frac{OL}{N-1}, \frac{2OL}{N-1}, \dots, \frac{(N-2)OL}{N-1}, OL\}$。真实标签被设置为能够使序列的 Rouge-L 分数匹配或超过 fp16 模型输出的最低高精度步骤。在训练期间，我们使用交叉熵损失函数作为目标函数。

A.2 预填充和解码位宽的错误容忍度

实验结果。如表4所示，我们观察到在预填充阶段使用更高的位宽能显著提高输出质量。另一方面，在解码阶段使用更高的位宽对性能的提升则没有那么大。这表明预填充阶段应该使用更高的位宽。

表4：不同解码和预填充位宽分配的Rouge-L分数。使用了Vicuna-7B和MTBench。fp16模型的输出用作参考答案。
表4：不同解码和预填充位宽分配的Rouge-L分数。使用了Vicuna-7B和MTBench。fp16模型的输出用作参考答案。

A.3 PMPD 在更大型号上的性能

通用性验证。尽管 PMPD 主要针对边缘和移动应用，但我们对其在更大型号上的效果进行了额外实验，使用了 Llama3-8B 和 longchat-16k-13B 模型，并为 PMPD 配置了静态调度器。表5中的结果表明，我们的方法在不同模型规模上具有良好的通用性。

表5：PMPD在大于7B的模型上的性能。
表5：PMPD在大于7B的模型上的性能。

A.4 GPU 端到端延迟加速

加速效果。我们在表6中提供了与 FP16 模型相比的 GPU 端到端延迟加速。PMPD 相较于统一量化基线，持续提供了高达0.31倍的速度提升。

表6：与FP16模型相比的端到端GPU延迟加速。
表6：与FP16模型相比的端到端GPU延迟加速。

A.5 GPU 矩阵乘法延迟评估

详细数据。表7和表8报告了在 RTX 4090 和 A40 GPU 上使用 self CUDA time 指标测量的 GPU 内核延迟。对于 PMPD，延迟是固定位宽延迟的加权平均值，权重对应于在各个数据集上以每个位宽执行的平均解码步数。在静态和学习型调度器之间，我们选择了性能最高的一个。对于 MobileLLaMA 和 Vicuna-7B，位宽是在 CNN/DM、Dialogsum 和 IWSLT 数据集上平均的。对于 Phi-1.5，位宽是在 CNN/DM 和 Dialogsum 数据集上平均的，而对于 Zephyr-3B，位宽来自 IWSLT 数据集。

表7：RTX 4090上的GPU矩阵内核延迟（µs）。
表7：RTX 4090上的GPU矩阵内核延迟（µs）。
表8：A40上的GPU矩阵内核延迟（µs）。
表8：A40上的GPU矩阵内核延迟（µs）。

A.6 NPU 仿真细节

仿真环境。我们开发了一个分析性能模型，以模拟 PMPD 在 FlightLLM【【33】，FlightLLM: Efficient Large Language Model Inference with a Complete Mapping Flow on FPGAs，2024，ISFPGA】上的部署，这是一个为 LLM 推理优化的 NPU 架构。FlightLLM 包含：i) 一个统一的矩阵处理单元，通过乘法累加（MAC）单元的层次结构执行多种矩阵/向量乘法；ii) 一个特殊功能单元，用于处理 softmax 等 LLM 特定操作；以及 iii) 一个为 LLM 计算模式优化的内存层次结构。我们为实验实例化了两种 NPU 配置：分别为部署小型（MobileLLaMa-1.4B, Phi-1.5）和大型（Vicuna-7B）LLM 的 4K 和 16K MAC 单元，时钟频率为 1 GHz（即峰值吞吐量分别为 8 和 16 TOPS），片外内存带宽为 32 GB/s。

A.7 替代的 PTQ 方法

框架兼容性。我们的方法提供了一个即插即用的框架，与任何 PTQ 方法兼容。为了展示其通用性，我们使用 GPTQ（一种统一量化技术）【【12】，GPTQ: Accurate Post-training Compression for Generative Pretrained Transformers，2023，ICLR】评估了静态和学习型调度器的性能。如表9所示，两种调度器都提供了与高精度基线非常接近的性能。

表9：在Dialogsum上使用GPTQ量化的MobileLLaMA的性能。
表9：在Dialogsum上使用GPTQ量化的MobileLLaMA的性能。

A.8 PMPD 的输出示例

成功案例。
* 问题：人生的各个阶段如何塑造我们对时间和死亡的理解？
* 回答：PMPD 模型能够就人文学科主题产生合理的回答，分点论述了童年、成年、中年和晚年对时间和死亡理解的影响。

失败案例。
* 问题：一个数除以10余4。这个数的两倍除以4的余数是多少？
* 回答：模型生成了连贯的文本，但回答与问题不太相关，给出了一个通用的求余数步骤，而不是直接解决问题。