清华大学让大语言模型自动并行自回归解码，速度提升高达4倍！

引言：大型语言模型的高效部署挑战

随着大型语言模型（LLMs）在各种人工智能应用中的广泛应用，对于这些模型的高效部署需求日益增长。高效部署在这里指的是低延迟和高吞吐量。然而，LLMs的基础自回归（AR）结构给实现更高效的服务带来了显著挑战。首先，自回归解码过程中，每一个新的词元都是基于之前生成的所有词元条件生成的，这种增量式解码导致生成速度不理想，因为每一步生成都需要访问LLM的大量参数。其次，Transformer模型在生成过程中对所有前序词元的注意力计算也限制了服务吞吐量。在高吞吐量场景中，许多序列并行生成，生成过程变得计算密集型。同时，注意力计算的成本与序列长度线性增长，尤其对于长响应，这限制了吞吐量的进一步提高。此外，为生成的词元缓存键值对（KV缓存）的需求，尽管在内存高效算法方面有所进展，但仍与序列长度线性增长，限制了系统能够处理的并发请求数量。

针对这些挑战，本文介绍了一种自动并行自回归（APAR）解码策略，旨在提高LLMs的推理效率。APAR利用LLMs生成过程中固有的可并行化结构，通过在包含层次结构的通用域数据上进行微调，使模型能够独立规划其生成过程并执行APAR生成，显著减少了生成步骤的数量。APAR本身可以实现高达2倍的速度提升，结合推测性解码时，速度提升可达4倍。此外，APAR减少了生成过程中的KV缓存消耗和注意力计算，导致在高吞吐量场景中，与最先进的服务框架相比，吞吐量增加了20-70%，延迟减少了20-35%。

声明：本期论文解读非人类撰写，全文由赛博马良「AI论文解读达人」智能体自主完成，经人工审核、配图后发布。

智能体传送门：
https://www.saibomaliang.com/generate?agent_id=68248fd1-32f9-4869-a35d-b6086ac0ebcf

神奇口令：小瑶读者（前100位有效）

论文基本信息

论文标题：

APAR: LLMs Can Do Auto-Parallel Auto-Regressive Decoding

论文链接：

https://arxiv.org/pdf/2401.06761.pdf

APAR方法简介

自动并行自回归（APAR）解码策略的核心思想是利用大型语言模型（LLMs）在文本生成中的固有并行结构，通过在具有层次结构的数据上进行微调，使模型能够独立规划其生成过程，并执行自动并行自回归（APAR）生成。这种方法显著减少了生成步骤的数量，单独使用APAR可以实现高达2倍的速度提升，结合推测性解码时，速度提升可达4倍。

APAR通过训练语言模型识别可并行化的结构，并在遇到可并行化的响应结构时，自动启动并行解码线程。这种方法将传统的线性生成转变为可并行化的段落树结构，不仅提高了解码并行性，还通过基于树的注意力机制减少了注意力跨度，并使得消耗的键值（KV）缓存内存能够更早释放。

在实验中，研究人员使用了Vicuna模型家族进行测试。在内存受限的场景中，APAR有助于减少模型延迟，并在Vicuna Bench上实现了平均2倍的生成速度提升。此外，在高吞吐量场景中，APAR与vLLM的兼容性使得可以提前释放内存，减少了高达50%的KV缓存需求，同时保持了相同水平的吞吐量。APAR还减少了参与注意力计算的令牌数量。使用相同数量的KV缓存内存时，与原始自回归过程相比，吞吐量提高了20-70%，延迟降低了20-35%，同时保持了相同的服务并发性。重要的是，使用APAR的生成质量并未受到影响，在多个类别的MT Bench和Vicuna Bench上的评估表明，响应质量与其自回归对应物大致一致，变化范围在±2%之内。这表明基于APAR的模型在提升解码速度和效率的同时，保留了上下文生成能力。

APAR的工作原理

1. 并行解码结构的设计

APAR（Auto-Parallel Auto-Regressive）解码策略的核心在于利用大型语言模型（LLMs）对文本结构的理解，以发掘文本中固有的可并行化结构。通过在具有层次结构的语料上微调LLMs，模型能够在遇到可并行化的响应结构时，自主地启动并行生成线程。这种方法将传统的线性生成转变为可并行化的段落树结构，从而实现更大的解码并行性，并通过基于树的注意力机制减少注意力跨度，使得消耗的键值（KV）缓存内存能够更早释放。

2. 控制token的使用

为了使语言模型能够启动并行解码线程，APAR在词汇表中增加了两个控制token：[Fork]和[Child]。[Fork]token用于指示响应中的可并行结构，模型在发现后续内容是详细信息（或子段落）时输出[Fork]token，并且可以与同一层次的下一个段落一起解码。当推理系统检测到模型输出的[Fork]token时，它会创建一个共享相同前缀的并行解码线程。[Child]token总是跟随[Fork]出现，用于指示接下来的内容是由前面的内容引导的子段落的开始。在训练过程中，[Child]被注意到但不参与损失计算，因此模型从不学习输出这个token，但学会在[Child]被注入上下文时（即上下文中出现[Fork] [Child]序列）输出子段落的内容。

3. 训练注意力机制

为了实现段落的并行生成，所有节点只关注它们的祖先节点，并且使用因果掩码关注自身。这种训练注意力机制确保了在并行生成时，新生成的每个段落都只关注到其在段落树中的路径上的节点。

APAR解码过程详解

1. 序列和序列组的概念

在APAR解码算法中，序列被定义为一个有序的token列表。序列组是为同一提示序列生成的所有序列的集合，初始化时仅包含提示序列。每个序列组对应一个段落树，而每个序列则是一个生成线程，与段落树中的一个叶节点相关联。

2. 解码算法的步骤

解码过程从用户提示序列p开始，并构建一个以p为根的段落树，将p与根节点r关联（现在是一个叶节点）。然后，迭代地对序列组G执行APARDECODE操作，直到G中的所有序列都完成。最后，遍历段落树以恢复顺序输出g。

在APARDECODE的详细步骤中，对于序列组G中的每个未完成序列s，如果最后一个token是[Fork]，则意味着模型调用了一个新的生成线程，并且需要分叉序列s′。分叉的序列与父序列共享相同的前缀token。在实现分页注意力机制时，分叉操作会为共享前缀创建一个共享内存映射，最多复制1个KV缓存块并共享所有其他块。分叉操作后，s′被追加一个强制的[Child]token，以便将这个序列标识为子序列。同时，创建两个新的叶节点，并设置s和s′以跟踪最新的叶节点。

最后，新采样的tokenx被追加到s。如果采样到[EOS]token，则认为序列s的生成已完成，并立即释放仅属于s的KV缓存。

APAR的性能特点

1. 减少延迟的并行解码结构

APAR通过在大型语言模型（LLMs）中引入并行自回归生成方法，显著减少了生成步骤的数量。这种方法利用了LLMs对文本结构的理解，通过在通用领域数据上进行指导性微调，使模型能够独立规划其生成过程，并执行自动并行自回归（APAR）生成。这种并行解码结构不仅减少了生成步骤，还通过基于树的注意力机制减少了注意力跨度，并使得消耗的KV缓存内存能够更早释放。

2. 通过提前释放KV缓存减少内存消耗

在自回归生成过程中，所有生成的token的KV缓存必须保留到序列完全生成为止。然而，在APAR中，一旦分叉序列（即生成线程）完成生成，仅属于该分叉序列的KV缓存就可以立即释放，而生成的其余部分继续进行。这种早期释放策略使得在吞吐量保持不变的情况下，可以节省高达50%的生成缓存。

3. 通过减少注意力长度节省计算资源

自回归生成要求每个token都要注意到之前生成的所有token。与此相反，在APAR中，新token只需关注其到段落树根路径上的token，这减少了生成过程中的注意力计算。在批量生成设置中，由于内存访问引起的延迟被密集的批处理计算所抵消，使得生成过程主要受计算限制。因此，每个token的计算减少导致在不同内存使用情况下吞吐量的提高，以及在不同并发程度下延迟的减少。

实验设置和数据预处理

1. 模型选择和实现

实验中采用了Vicuna模型家族进行评估。对于评估生成速度、吞吐量和质量，APAR微调应用于vicuna-v1.3-{7B,13B}模型，产生了APAR-{7B,13B}。原始的Vicuna模型将被称为Original-{7B,13B}（简称O-{7B,13B}），而经过微调的APAR模型将被称为APAR-{7B,13B}（简称A-{7B,13B}）。

2. 训练设置和评估数据集

在微调过程中，从结构化（有序列表和段落，16k样本）和非结构化数据（9k样本）中采样，采样比例为1:1。模型以128的批量大小、2e-5的学习率进行了2000步的微调。微调后，使用相同的数据以1e-3的学习率训练了2个Medusa头部，进行了2000步的训练。

评估数据集包括Vicuna Bench（用于评估LLMs在语言理解、推理和上下文意识方面的性能）和MT Bench（包含80个多轮问题，用于评估LLMs的多轮对话和指令遵循能力）。此外，APAR测试集包含从ShareGPT数据集中采样的1000个用户查询，以模拟真实部署场景中的查询分布。

实验结果分析

1. 内存受限场景下的性能

在内存受限的场景中，即批处理大小较小的情况下，APAR（Auto-Parallel Auto-Regressive）解码策略能够显著减少生成延迟。通过在Vicuna Bench上的实验表明，Vanilla-APAR在内存受限的情况下能够实现平均2倍的速度提升。当与Medusa结合使用时，这种提升甚至达到了4倍。这一结果表明，APAR能够有效地在需要并行生成结构的场景中减少生成步骤，从而提高生成速度。

2. 高吞吐量场景下的性能

在高吞吐量的场景中，APAR的兼容性使其能够在早期释放内存，减少了KV缓存的需求，同时保持了相同水平的吞吐量。实验结果显示，APAR能够在使用相同数量的KV缓存内存的情况下，将吞吐量提高20-70%，并且在保持相同服务并发性的情况下，将延迟降低了20-35%。这些改进主要归功于APAR的两个关键特性：并行解码结构减少了延迟，以及子KV缓存的早期释放减少了内存消耗。

3. 生成质量评估

APAR的引入并没有牺牲生成质量。在MT Bench和Vicuna Bench上的评估表明，APAR模型与其自回归（AR）对应模型相比，在响应质量上的变化在±2%的范围内，这表明APAR模型在提高解码速度和效率的同时，保留了上下文生成能力。

与现有工作的比较：APAR与先前推理加速工作的区别和联系

APAR与先前的推理加速工作相比，最大的区别在于它不修改运算符或模型架构，而是通过采用基于树的注意力结构来减少计算量。这使得APAR可以与现有的模型压缩、高效注意力和多查询注意力等技术正交使用，并且可以与这些技术共同应用以进一步提高效率。

在改进并行性方面，APAR通过利用大型语言模型（LLMs）的内在组织能力，自动发起并行生成线程，适用于多种场景。与其他通过提示生成响应骨架然后并行扩展每个点的方法不同，APAR不需要额外的分类器和KV缓存之间的重新计算，因此不会影响生成吞吐量。

总的来说，APAR通过在不牺牲生成质量的前提下，提高了并行性，并减少了计算和KV缓存内存消耗，实验表明它可以与现有的推理框架无缝集成，显著降低了各种场景下的生成延迟，同时在涉及极端批处理大小和并发水平的情况下提高了服务吞吐量。

结论：APAR对大语言模型解码效率的提升及未来展望

在探索大语言模型（LLMs）的高效部署策略中，我们介绍了一种新的并行自回归生成方法——自动并行自回归（APAR）解码策略。APAR通过在包含层次结构的通用领域数据上进行指导调整，使LLMs能够独立规划其生成过程，并执行APAR生成，显著减少了生成步骤的数量。单独使用APAR可以实现高达2倍的速度提升，结合推测性解码时，速度提升可达到4倍。此外，APAR在生成过程中减少了键值缓存的消耗和注意力计算，这导致在高吞吐量场景中，与最先进的服务框架相比，吞吐量增加了20-70%，延迟减少了20-35%。

1. APAR的效率提升

APAR通过训练模型识别并利用文本中的并行结构，使得生成过程可以在检测到可并行的响应结构时，自动发起并行生成线程。这种方法将传统的线性生成转变为可并行的段落树结构，不仅提高了解码并行性，还通过基于树的注意力机制减少了注意力跨度，并使得消耗的KV缓存内存能够更早释放。

2. 未来展望

APAR的设计与大多数现有的推理加速方法相辅相成。例如，与Medusa等推测性解码策略结合时，APAR模型在Vicuna Bench上实现了高达6倍的速度提升。在特定类别中，这种组合甚至实现了更高的速度提升。在高吞吐量场景中，APAR与vLLM的兼容性使得可以在保持相同吞吐量的同时，减少高达50%的KV缓存需求。此外，APAR减少了参与注意力计算的令牌数量，使用相同数量的KV缓存内存时，与原始AR过程相比，吞吐量提高了20-70%，延迟减少了20-35%，同时保持了相同的服务并发性。

3. 保持生成质量

APAR在提高解码速度和效率的同时，并未牺牲生成质量。在多个类别的MT Bench和Vicuna Bench上的评估表明，响应质量基本保持一致，与AR对应物相比变化在±2%的范围内。这表明APAR基础模型在保持上下文生成能力的同时，增强了解码速度和效率。

综上所述，APAR为大型语言模型提供了一种有效的解码效率提升方法。通过减少生成步骤和优化内存使用，APAR在不牺牲生成质量的前提下，显著提高了模型的服务效率。未来，我们期待APAR能够与更多的推理加速技术相结合，进一步探索和扩展其在不同场景下的应用潜力，为大型语言模型的高效部署和应用开辟新的道路。

声明：本期论文解读非人类撰写，全文由 赛博马良「AI论文解读达人」 智能体自主完成，经人工审核、配图后发布。

智能体传送门：
https://www.saibomaliang.com/generate?agent_id=68248fd1-32f9-4869-a35d-b6086ac0ebcf

神奇口令：小瑶读者（前100位有效）