AI 大模型推理优化技术详解

2.1.2 Paged Attention

当谈到 LLM 推理服务的吞吐量时，显存限制往往成为主要瓶颈。研究团队洞察到，由于缺乏精细的显存管理方法，现有的系统浪费了高达 60% 至 80% 的显存，而这大部分浪费都来自 KV Cache。因此，有效管理 KV Cache 变得至关重要。

在 Paged Attention 之前，主流 LLM 推理框架在 KV Cache 管理上存在不足。HuggingFace Transformers 库中，KV Cache 随着执行动态申请显存空间，这不仅导致高时延开销，还引发显存碎片化问题。而 FasterTransformer 则采用预先分配长显存空间的方法，若用户实际使用的上下文长度低于预设值，则会造成显存浪费。

Paged Attention 将传统操作系统的内存管理思想引入 LLM 领域，创新地实现了一个高效的显存管理器。它通过精细化的显存管理，实现了在物理非连续的显存空间中以极低成本存储、读取、新增和删除键值向量。

具体而言，Paged Attention 将每个序列的 KV Cache 分为多个块，每个块包含固定数量的 token 键值对。在实际推理前，系统会预跑一次推理计算，根据用户设置的 max_num_batch_tokens 和 gpu_memory_util 参数，计算出 KV Cache 可用的最大空间，并预先申请缓存空间。这里，max_num_batch_tokens 代表硬件显存一次最多能容纳的 token 总量，gpu_memory_util 为模型推理的最大显存占用比例，block_size 为块大小（默认设为 16）。

在实际推理过程中，Paged Attention 维护一个逻辑块到物理块的映射表。多个逻辑块可以对应一个物理块，通过引用计数来管理物理块的使用情况。当引用计数大于 1 时，物理块被使用；当引用计数为 0 时，物理块被释放。这种设计巧妙地将地址不连续的物理块串联起来，实现了统一管理。

Paged Attention 技术将操作系统的分页内存管理应用于 KV Cache 管理，极大地提高了显存利用效率。通过 token 块粒度的显存管理，系统能够精确计算出剩余显存可容纳的 token 块数量，配合 Dynamic Batching 技术，可有效避免显存溢出问题。

2.2 算法

2.2.1 Flash Attention

Transformer 类模型的核心自注意力机制在计算和存储上，其复杂度随序列长度呈二次增长，这成为了性能瓶颈。虽然之前的优化工作主要关注减少 FLOPs 以提升计算速度，但显存 IO 访问的优化却被忽视了。

事实上，GPU SRAM 的读写速度高达 19TB/s，而 GPU HBM 的读写速度却与之相差甚远，存储容量也相差巨大。Flash Attention 的设计充分考虑了这一点，它通过从 HBM 读取数据块到 SRAM 进行计算，然后再写回 HBM，从而有效避免了频繁从 HBM 中读取或写入注意力矩阵。

FlashAttention 的主要目标包括：

• 目标一：在无需访问整个输入数据的情况下，实现 softmax 函数的缩减计算；
• 目标二：在后向传播过程中，避免存储中间注意力矩阵，从而节省显存空间。

论文中提出的具体算法，通过精细的设计和高效的实现，成功达到了上述目标，为 Transformer 类模型的优化提供了新的思路。

2.2.2 GQA

GQA 将查询头巧妙地划分为 G 组，每组共享单一的键头和值头。

• GQA-G: 代表 G 组分组查询，将查询头分为多个小组，每组拥有共享的键和值头。
• GQA-1: 具有单个组，因此相当于 MQA（Multi-Query Attention），即单个键和值头。
• GQA-H: 组数等于头数，这实际上等同于传统的 MHA（Multi-Head Attention）。

在将多头检查点转换为 GQA 检查点时，采用平均池化的策略，将组内的所有原始头进行合并，构建出每个组的键和值头。

中间组数的选择使得 GQA 在模型质量上插值于 MQA 和 MHA 之间，既保证了模型质量高于 MQA，又保证了推理速度优于 MHA。

从 MHA 过渡到 MQA，键值缓存的大小从 H 键和值头减少到单个键和值头，这意味着需要加载的数据量减少了 H 倍。然而，随着模型大小的增加，头部的数量通常也会增加，这使得多查询注意力在内存带宽和容量方面表现更佳。

但 GQA 的设计允许我们随着模型大小的增加，仍然保持相同的带宽和容量比例下降，从而实现了高效且灵活的注意力机制。

2.3 计算

2.3.1 算子融合

算子融合是深度学习模型推理的秘密武器，它通过减少计算过程中的访存次数和 Kernel 启动耗时，来显著提升模型推理性能。对于 LLM 推理来说，这一技术同样不可或缺。

以 HuggingFace Transformers 库中的 LLaMA-7B 模型为例，通过深入分析模型推理时的算子执行分布，我们发现该模型拥有 30 个类型共计 2436 个算子，其中 aten::slice 算子就高达 388 次。这么多小算子频繁执行，无疑会拖慢 GPU 的步伐，影响推理速度。

为了应对这一挑战，业界已经进行了许多尝试。以 DeepSpeed Inference 为例，算子融合主要聚焦在以下四个方面：

• 归一化层和 QKV 横向融合：将原本需要三次计算的 Query/Key/Value 操作，合并为一个高效算子，并与前面的归一化算子融为一体。
• 自注意力计算融合：将自注意力计算中的多个算子'一网打尽'，形成一个强大的融合算子。FlashAttention 就是这一领域的佼佼者。
• 残差连接、归一化层、全连接层和激活层融合：在 MLP 中，将这些相互关联的算子'打包'成一个，简化计算流程。
• 偏置加法和残差连接融合：让这两个'小伙伴'携手共进，减少计算开销。

然而，算子融合并非易事，通常需要定制化实现算子的 CUDA kernel，这对 GPU 编程能力提出了较高要求。幸运的是，随着编译器技术的飞速发展，OpenAI Triton、TVM 等优秀框架已经能够实现算子融合的自动化或半自动化，让这一技术更加易于应用。

2.3.2 高性能算子

为了降低 LLM 推理的时延，我们专注于编写针对运行热点函数的高性能算子。以下是两个关键优化点：

• GEMM 操作优化：在 LLM 推理的预填充阶段，Self-Attention 和 MLP 层中的 GEMM 操作占据了推理时延的 80% 以上。英伟达已经推出了多种优化方案，如 cuBLAS、CUDA、CUTLASS 等，来针对这一问题进行 GPU 级别的优化。例如，在 FasterTransformer 框架中，我们大量采用了基于 CUTLASS 编写的 GEMM 内核函数。同时，由于 Self-Attention 中存在特殊的 GEMM+Softmax+GEMM 结构，我们还会结合算子融合技术来进行联合优化。
• GEMV 操作优化：当进入 LLM 推理的解码阶段时，运行热点函数会从 GEMM 转变为 GEMV。由于 GEMV 的计算强度相对较低，我们的优化重点主要放在降低访存开销上。

实现这些高性能算子对 GPU 编程能力有着较高要求，并且算法实现中的某些超参数与特定问题规模密切相关。因此，编译器相关的自动调优技术也成为了我们研究的重点，以确保算子在各种场景下都能发挥出最佳性能。

2.4 服务

为了提升服务性能，引入了两大关键优化策略：Continuous Batching 和异步 Tokenize / Detokenize。

• Continuous Batching：该技术通过提高可并发的 batchsize 来显著提升吞吐量。它采用 step 级别的任务调度，将传统 batch 粒度的任务细化，以消除'气泡'现象。我们维护两个队列——Running 和 Waiting，队列中的序列状态可以在这两者之间灵活转换。每当一个序列推理完成后，它会从 Running 队列中移除，并从 Waiting 队列中取出新的序列进行推理。这种方式确保了 GPU 资源的最大化利用，有效提升了推理效率。
• 异步 Tokenize / Detokenize：这一技术通过多线程并行处理 Tokenize 和 Detokenize 操作，与模型推理过程时间交叠，从而显著降低时延。它允许我们在处理用户请求时，同时进行文本的分词和重组，提高了整个服务流程的效率。

大语言模型的输入和输出长度可变，给服务带来了挑战。传统方法在同批次序列推理过程中，存在'气泡'现象，导致 GPU 资源浪费。而 Continuous Batching 技术通过动态调整 Running 队列中的序列，消除了这种浪费，实现了 GPU 资源的最大化利用。

此外，还引入了 Dynamic Batching 的概念。在自回归迭代生成每个 token 后，调度器会根据当前剩余显存量，动态调整 Running 队列的长度。这允许我们在 batch 维度上动态插入新序列，以充分利用显存空间。

这两项优化技术已经在 HuggingFace 的 Text-Generation-Interface (TGI)、vLLM、OpenPPL-LLM 等多个框架中得到实现，并显著提升了推理吞吐量。

2.5 分布式

由于大语言模型参数量庞大，单一计算设备可能难以承载，因此分布式并行技术显得至关重要。在 LLM 推理中，模型并行和流水线并行已成为得力助手。模型并行巧妙地将权重参数分散到多个计算设备上，实现高效的分布式计算。

各种并行术语小课堂：

• 流水线并行（PP）：将模型的不同层分配到各个 GPU，通过分割 mini-batch 实现数据的流水线处理，从而提升 GPU 的计算与通信效率。
• 数据并行（DP）：遵循 Single Program Multiple Data 原则，任务被分配到多个进程（设备）上，每个进程处理不同的数据，但维护相同的模型参数和计算任务。
• 张量并行（TP）：模型参数被分割成不同部分，并分配到多个设备上进行计算，最后通过通信聚合结果，确保数学一致性。
• 模型并行（MP）：有两种含义，一是层间模型并行，类似于流水线并行；二是层内模型并行，类似于张量并行。

在上图中，第一行展示了 Column Parallel，即将权重数据纵向分割到多个 GPU 中，结果需在列方向拼接；第二行则是 Row Parallel，权重数据横向分割，结果需通过 AllReduce 规约。

业界翘楚 Megatron-LM 提出了针对 Self-Attention 和 MLP 的简洁高效模型并行方案：

• MLP: 首个全连接层采用 Column Parallel，第二个则采用 Row Parallel，整个 MLP 仅需在 Row Parallel 后执行一次 AllReduce。
• Self-Attention：在计算 Query、Key 和 Value 向量时采用 Column Parallel，而在空间映射时采用 Row Parallel，整个过程也仅需一次 AllReduce。

除了上述的并行方式，LLM 模型中的 Input Embedding 采用 Row Parallel，Output Embedding 采用 Column Parallel；而 Dropout、Layer Norm 等操作则保持完整，未进行并行拆分。

以 LLaMA-34B 模型为例，我们进行了通信量分析。该模型包含 48 个 Transformer layers，隐藏层大小为 8192。在每次单 batch 推理中，共需进行 2*48 次 Broadcast 和 248 次 AllReduce 操作，每次通信数据量均为 16KB。当考虑到多 batch 推理，假设 batchsize 为 64 时，每次通信数据量也仅为 1MB。

在 A100-PCIE-40GB 机器上的 NCCL AllReduce 带宽测试显示，尽管理论带宽高达 32-64GB/s，但实际推理场景下的通信数据量主要集中在 1MB 以下，对应的实际带宽约为 1-10GB/s。因此，模型参数量和 batchsize 越大，通信效率越高，模型并行的收益也就越明显。

2.6 低比特量化

在提升 LLM 模型推理性能上，吞吐量和时延是两个关键指标。吞吐量受限于显存容量，而低比特量化技术正是解决显存占用和访存量的有效手段。通过降低显存占用量，我们可以运行更大的 batchsize，从而提升吞吐量；同时，由于 LLM 推理具有 Memory-bound 特性，降低访存量将在吞吐量和时延上带来双重收益。低比特量化的关键在于其节省的显存量和访存量，以及量化计算带来的加速远大于反量化带来的额外开销。

表中的四类量化方法各有特点，业界在低比特量化方向的研究日新月异，我们期待能够探索出一个适用于大语言模型的量化方法，它能在保证精度的同时，以较高的压缩率压缩模型，并加速端到端推理。

2.7 其他新技术

在 LLM 推理领域，除了引入传统优化技术外，业界也在不断探索新的解决方案。从大模型自回归解码的特点出发，一些新技术正在尝试通过调整推理过程和引入新的模型结构来进一步提升性能。

例如，投机采样（Speculative decoding）技术，它针对 LLM 推理的串行解码特点，引入了一个近似模型来执行串行解码，而原始模型则执行并行评估采样。通过近似模型和原始模型的配合，该技术能够在保证精度一致性的同时，降低大模型串行解码的次数，从而显著降低推理时延。这是一项令人兴奋的技术，有望为 LLM 推理性能带来质的飞跃。