点击查看 动态视频演示

图 1：高通 Cloud AI 100 Ultra 执行生成任务 (1) FP16 权重基线 (2) 使用 MX6 加速 (3) MX6 和 SpD 加速

推测采样 (SpS)，也称为推测解码(SpS)，和通过 MXFP6 微缩放格式进行权重压缩，是两种高级技术，可以显著提高大语言模型 (LLM) 推理 AI 工作负载的LLM解码速度。这两种技术都可用于高通数据中心人工智能加速器的LLM加速。要显著提高推理性能，请先了解 Amazon AWS EC2 和 Cirrascale 云服务。

本文探讨了两种高级技术在两个大语言模型上的应用：CodeGen 1-7B 和 Llama 2-7B-Chat-FT，展示了加速 AI 处理和效率的潜力。我们将为您揭开这一进步的细节，您也可以先浏览 高通云人工智能 Github 上的推测解码功能文档。

语言模型（LM）和自回归生成

本节介绍基于传统自回归解码的LLM解码的基础知识，及其固有的多令牌生成顺序特性。

文本补全是 LM 的常见任务：给定提示文本，LM 会生成合理的最终文本。

LM 以令牌为文本原子单位粒度处理文本。令牌相当于单词中的子词。LM 可以处理的令牌集称为词汇表，用 Σ 表示。

语言模型将一系列令牌作为输入，并在词汇表中生成概率分布函数（pdf）。从此概率分布函数中抽样在输入后给出了一个合理的令牌。在实现过程中，Σ 上的 pdf 表示为长度为 |Σ|的向量，包含元素[0,1]，总和为 1。

自回归 （AR） 生成

通常，我们希望生成多个后续令牌，而不仅仅是一个令牌。给定 m 个令牌 u1,...,um 的提示语，生成 n 个令牌 v1,...,vn，但是需要 n 次调用 LM（作为仅解码器转换器模型实现），如下所示：

图 2 自回归 （AR）生成概述

图 2 中，观察上述完成生成过程的以下内容：

Logits 和翘曲：LM 生成一个词汇表大小的向量：logits。通过 softmax 对其进行归一化，在 Σ 上得到一个 pdf。从此 pdf 中采样一般称为多项式采样。

这样的 pdf 也可以被翘曲，将概率质量重新分布于 Σ 的小子集。

多项式采样常用的特殊形式是贪婪采样，由top-1翘曲产生。它将所有概率质量置于具有最高 logit 值的令牌，而将 0 置于其他令牌。从此类分布中采样将始终产生具有最高概率/logit 值的令牌。

多次调用：LM 每次被调用 n 次 - 每个生成的令牌调用一次。每次调用都使用模型的所有参数。

自回归：在下一次迭代中，LM 的输出与其输入相结合。

因果关系：如果不生成所有先前的令牌，就不能生成v_i，因此，生成是一个连续的过程——vi 和 vj （i≠j） 的生成不能并行处理。

重复计算：LM 调用是独立的，并且 LM 不会在调用之间维护内部状态。这将导致同一输入被多次处理，从而生成相同的键值向量集。例如，提示语本身被处理 n 次，生成的令牌 vi 被处理 n-i 次。

使用 KV 缓存（KV$）减少计算，加速自回归生成

避免重复计算的常用方法是在调用之间维护 LM 的内部状态：键值缓存（KV$）。此模型 LM_with_KV$ 的功能与上述 LM 相同，但操作上存在以下差异：

• 当首次调用提示语 u1,...,um 时，将创建 KV$ 并写入其中，即 m 提示语的键值向量。
• 为后续令牌生成而调用时：
  • 仅将上一次迭代生成的单个令牌作为输入，而非不断累积的令牌输入。
  • 计算单个输入令牌的键查询值向量，并将键值附加到 KV$
  • 仅处理 LM 所有层的单个令牌，但使用 KV$ 中所有键值向量计算单个令牌的因果注意力（causal attention）。

使用 LM_with_KV$() 的文本补全可以重写如下：

值得强调的是，具有适当 KV$ 内容的 LM_with_KV$() 和 LM() 是相同的函数——两者均计算可用于预测 vi 之后下一个令牌的词汇表上相同条件概率分布。

LM(u1..um v1…vi ) ≡ LM_with_KV$(vi) ≡ p( .|u1,u2,…,um,v1,…vi)

LM 补全生成性能配置文件

LM_with_KV$() 执行阶段将提示语作为输入并生成第一个令牌，称为预填充阶段。以自回归方式生成后续令牌（一次生成一个令牌）的执行阶段称为解码阶段。

这些阶段使用完全不同的性能配置文件。预填充阶段只需要对 LM 调用一次，需要从 DRAM 获取一次模型的所有参数，并重用 m 次，处理提示语中的所有 m 个令牌。如果m值足够大，则预填充阶段的计算受限高于内存带宽受限。

另一方面，解码阶段对 DRAM 带宽的限制则高于计算限制，因为需要获取模型的所有参数，处理一个输入令牌。此外，加速器计算资源的利用率非常低，因为一次只处理一个令牌。换句话说，解码阶段可以在不影响延迟的情况下执行额外的计算。

对于一阶近似值，在解码阶段，LM 可以生成的令牌数量可以使用以下公式计算：DRAM-Read-Bandwidth（字节/秒）除以模型的容量（以字节为单位）。

推测解码

AR 解码无法并行处理，因为其因果性只能连续处理：令牌v i+1 只能在令牌 vi 生成后生成。下面我们来看看 SpD 是如何克服这个障碍的。

检查 K (>1) 令牌序列是否是提示语的合理延续比生成该 K 令牌延续更快——因为前者只需要调用一次 LM，‌SpD 正是利用了这样的特点。

在 SpD 中，合理补全由另一个容量较低的 LM（称为草稿语言模型 （DLM））生成。在这种情况下，与此相对的原始 LM 称为目标 LM （TLM）。

要检查v_1,...,v_K是否是u‌1‌…u‌m 的合理补全，需要检查以下 K 条件：

v1...vi 是u‌1‌…u‌m的合理补全（i=1,..,K）

由于检查上述 K 条件内在并行性，可以通过调用一次 LM 实现——而生成 K 个令牌需要 K 次顺序调用。

为使整体推测解码方案更快，需要适当选择以下参数：

• 与 TLM 相比，DLM 的相对容量需要足够小，以自回归方式生成快速推测，同时相对容量又需要足够大，以生成 TLM 合理补全。
• 推测长度 （K） 需要足够小，确保 TLM 单次调用检查补全和 DLM 生成时间不会过分消耗计算资源。

更正式地说，给定 u‌1‌…u‌m 的提示语和可能的补全 v1...vK，可以有效地评估以下条件 pdf：

pLM(.|u1…um)

pLM(.|u1…um,v1)

pLM(.|u1…um,v1,…,vi-1) for i=2..K

以上表示 LM_with_KV$() 的一次调用，输入 K 个令牌 v1,...,vK 和 KV$，其中包含 u‌1‌…u‌m‌ 键值向量。此调用也称为对 K 个令牌 v1,...,vK 的序列进行评分。

总之，SpD 依赖于这样一个事实，即对于较小的 K 值，对 K 令牌进行评分产生的延迟与自回归生成单个令牌相比，没有区别，或者略高，因为计算单元在 AR 解码阶段未得到充分利用。

推测解码的组件

SpD 中有三个功能组件，可以根据给定 m 个令牌 u‌1‌,…u‌m 提示语生成 n 个后续令牌。

TLM: 我们用于生成补全的 LM 称为目标或真 LM（TLM）。

DLM: 一个较小的 LM 用于猜测 K 令牌的补全，称为草稿 LM （DLM），一次自回归生成一个令牌。

• • 与 n 相比，K 是一个较小的数字。由于 DLM 很小，因此增加的延迟也应该是最小的。
  • DLM 生成的序列是 TLM 推测补全。

MRS: 用于决定 TLM 是否接受推测补全或其前缀的概率接受方案，即修改拒绝采样 （MRS）。

• • 这是解码过程的核心，保证了推测采样的正确性。请注意，从 DLM 的角度来看，推测补全可能是合理补全，但对于 TLM 来说可能并非如此。因此，TLM 可能只接受推测的前缀，甚至不接受任何前缀。接受决定本质上基于概率的。作为 MRS 核心的概率推理是确保接受 u‌1‌,..u‌m‌t‌1‌,..t‌i‌‌  的概率确实匹配 TLM 分配的概率p‌TLM‌(u‌1‌,..u‌m‌t‌1‌,..t‌i‌)。

推测解码过程

首先，我们从较简单的情况入手，即 DLM 和 TLM 使用贪婪采样而不是一般的多项式采样。这种情况较为简单的原因是，当 pdf 变成 1-hot pdf 时，所有的概率推理都可以简化为确定性推理。

图 3：DLM 的 K 令牌 AR 生成

贪婪采样

对于给定的 m 个令牌提示语，DLM 自回归生成 K 个令牌 t1,..,tK，一次一个，并调用 K 次（如上图 3 所示）。

TLM 将处理序列 u‌1‌,u‌2‌,..u‌m‌,t‌1‌,..,t‌K，调用一次 TLM 并计算 v‌1‌,…,v‌K+1（如下图 4 所示）。

• • v1 是 TLM 根据  u‌1‌…u‌m 预测的下一个令牌。

图 4 推测时 TLM 单次调用

• • vi 是 TLM 根据 u‌1‌,u‌2‌,..u‌m‌,t‌1‌,..,t‌i-1 预测的下一个令牌（1 < i ≤ K + 1）。

改进的拒绝采样（MRS）算法：该组件从 DLM t1‌,…tK 中获取 K 个推测令牌，通过评分过程从 TLM v1‌,…v(K+1) 获取 K+1 个推测令牌。输出是提示语 u1…um期望的合理补全。从左到右扫描推测令牌，一次扫描一个，查看其是否与 TLM 在评分期间采样的相应令牌匹配。如果是，则被接受为合理补全；如果不是，则该过程将停止，来自 TLM 的令牌作为被拒绝令牌的替代，而其余的推测则被拒绝。如果全部接受，则接受额外的令牌。参见图 7，了解 K=3 时可能的接受情况。图 5 在 DLM 和 TLM 上使用贪婪采样的 MRS

图 6：推测的可能补全

此时，我们只需调用 1 次 TLM 和 K 次调用 DLM，即可为 TLM 生成至少 1 个且最多 K+1 个令牌。

重复上述推测-接受过程，直至生成所需数量的 n 个令牌。

有一个重要的实现细节，我们在这里不作详述，即 DLM 和 TLM 的 KV$ 应该清除被拒绝的令牌。

多项式采样

高通 Cloud AI 100支持 DLM 和 TLM 执行多项式采样的一般用例。在这种情况下，当 TLM 对输入序列进行评分并输出条件 pdf 时，MRS 方案需要以适当的概率做出概率决策，以便从所需的目标分布中对补全采样。有关多项式采样的进一步说明，请参阅此处本文的扩展版本。

总之，SpD 使用自回归的较小容量模型 （DLM），推测多令牌补全。由于 DLM 较小，因此执行速度远快于  TLM。MRS 算法仅调用一次 TLM 并接受推测的前缀，验证推测是否是 TLM 的合理补全。理论上可以保证，DLM 和 TLM 之间的推测-接受握手过程不会损失准确性。

使用微缩放（Mx）格式执行仅权重量化

除推测采样之外，使用微缩放（Mx）格式的仅权重量化还可以在 LLM 解码上实现 ~2 倍的加速。

2023 年，AMD、Arm、英特尔、Meta、Microsoft、NVIDIA 和Qualcomm 成立了微缩放格式（MX）联盟，创建用于 AI 训练和推理的下一代 6 位和 4 位数据类型并形成标准。高通 Cloud AI 100 已将 MxFP6（6 位数据类型）作为权重压缩格式。当用户选择此编译选项时，编译器将在离线编译阶段自动将 FP32 或 FP16 中的权重压缩为 MxFP6 格式。以这种形式压缩可节省 61% 的权重，从而减轻 DRAM 容量压力。

在推理运行时，高通云人工智能 100 高通 Cloud AI 100使用矢量引擎和优化的解压缩内核，在软件中执行动态解压缩。解压缩可以与权重获取和计算并行执行，因此开销大多是不明显的。解压缩后，和以前一样以 FP16 精度执行计算。由于 LLM 仍然受到 DRAM 的约束，因此计算不会成为瓶颈，所以‌可以使用FP16。FP16 还允许保留更高精度的激活，从而克服量化造成的精度损失。

实验结果

本节详细介绍了在高通 Cloud AI 100上使用贪婪 SpD 和 MX6 权重压缩实现的性能加速，基线为无推测解码，并使用 FP16 的参数格式。

3 种不同的高通 Cloud AI 100版本实现的加速：

• 标准（75W TDP），
• Pro（75W TDP）和
• Ultra (150W TDP)

为确保多样性，我们报告了两个网络 CodeGen 1-7B mono 和 Llama 2-7B-chat-FT 的加速，均包含不同的上下文长度。CodeGen 1 是用于程序合成的模型系列。mono 子系列经过微调，可根据自然语言的规范生成 python 程序。Llama 2-7B-chat-FT 模型是高通根据 Meta 发布的 Llama 2-7B 微调的模型。

图 1 从博客的开头，显示了同时开始的使用特定提示语三次运行的 Llama 2 7B-chat-FT。左列为模型在基线条件下的执行情况，参数格式为 FP16；中间一列是 MX6 压缩下加速的模型；右列则是使用 MX6 压缩和 SpD 加速的模型。

从最右列生成和补全令牌的速度可以看到明显的加速。值得注意的是，MX6 + SpD 列中突出显示为蓝色的文本是已被接受的 DLM 的推测补全。

图 7 记录了图 1 中观察到的每秒生成令牌数（t/s）指标的加速。

图 7：云 AI 100 Ultra 加速 Llama 2 7B

使用高通 Cloud AI 100的 Standard 和 Pro 版本可以实现类似的加速，如图 8 所示。

图 8：高通 Cloud AI 100Standard 和 Pro 的加速范围

配置

本节讲述用于测量此加速的配置。

	CodeGen 1-7B	Llama 2-7B 聊天微调
TLM	~~7B 参数格式：MXFP6	~7B 参数格式：MXFP6
DLM	~350M 参数格式 FP16	115M 参数格式：FP16
最大提示语长度	32	128
最大生成长度	224	896
推测长度 （K）	7	7
批量规模	1	1
各种提示语	Hello World、BFS、拓扑排序、Tarjan SCC、线性回归	获取有关天气、旅行指南、餐厅建议、汽车控制、音乐搜索、航班搜索的信息

结语和下一步行动

大语言模型可以通过推理解码进行加速，从而在不损失准确性的情况下提高性能效率。使用 MXFP6 进行权重压缩是另一种 2 倍提升 LLM 速度的技术。通过协同作用，这两个功能共为 LLM 解码加速 ~4 倍。这两种技术都可以在高通 Cloud AI 100推理加速器上使用，可用于 Amazon AWS 和 Cirrascale 云服务. 请务必查看 高通云人工智能Github 存储库中为推测解码提供的模型方案  ，以简化您的 LLM 工作流程！

如果您对本文中的常见问题解答、附录和参考资料感兴趣，请浏览此页面。