AI 大模型推理优化旨在解决高成本与低效率问题，核心指标为吞吐量与时延。关键技术包括显存管理（KV Cache、Paged Attention）、算法优化（Flash Attention、GQA）、计算加速（算子融合、高性能算子）、服务调度（Continuous Batching、异步 Tokenize）、分布式并行（模型并行、流水线并行）及低比特量化。通过空间换时间策略减少重复计算，利用分页机制提升显存利用率，结合量化降低访存开销，最终实现推理性能与资源消耗的最佳平衡。
自从 OpenAI 发布了 ChatGPT，全球范围内对基于 Transformer 结构的大语言模型 (LLM) 的技术关注达到了前所未有的高度。这些模型凭借其强大的理解和生成能力，正在深刻地改变我们对人工智能的认知和应用。然而，大语言模型在推理应用方面的成本过高，极大地阻碍了技术的广泛应用。
因此，优化大语言模型的推理性能成为了业界研究的焦点。大语言模型推理不仅面临着计算资源的巨大需求，还面临着计算效率的挑战。通过优化推理性能，我们不仅可以降低硬件成本，还能提高模型的实时响应速度。这将使模型能够更快速地执行自然语言理解、翻译、文本生成等任务，从而改善用户体验，加速科学研究，推动各行业应用的发展。
LLM 推理服务的关键在于两大核心指标：吞吐量和时延。
吞吐量关注的是系统成本，高吞吐量代表系统能够高效处理更多请求，提高系统利用率。而时延则直接影响用户体验，快速响应是提升满意度的关键。然而，这两个指标往往相互影响，需要我们在实践中精心权衡。
例如，为了提升吞吐量，一个常见的策略是增大 batchsize，即将用户请求从串行转为并行处理。然而，这也会在一定程度上增加每个用户的等待时间，因为系统现在要同时处理多个请求。
为了优化 LLM 的推理性能，我们的目标是同时提高吞吐量和降低时延。这涉及到七个关键部分，下面将详细展开描述。
大模型推理性能优化的利器之一就是 KV Cache 技术。这项技术在不牺牲任何计算精度的前提下，巧妙地运用空间换时间的策略，显著提升了推理性能。目前，业界主流的 LLM 推理框架都默认支持并启用了这一功能。
Transformer 模型因其自回归推理的特性而独树一帜，即每次推理仅预测输出一个 token。当前轮输出的 token 与历史输入的 tokens 拼接后，成为下一轮的输入 tokens，如此循环往复。在这个过程中，相邻两轮的输入仅相差一个 token，导致了大量的重复计算。而 KV Cache 技术正是为了解决这个问题而生，它通过将可复用的键值向量结果保存下来，巧妙避免了这些重复计算，从而大幅提升了推理效率。
具体来讲，KV Cache 技术在每次自回归推理过程中发挥了巨大作用。当 Transformer 模型的每一层 Attention 模块执行时，它会将 Q、K、V 的计算结果保存到一个预先分配好的数据结构（我们称之为 KV Cache）中。这样，当下一次自回归推理开始时，我们只需将新的 Q 与已有的 KV Cache 中的 K、V 拼接起来，即可供后续计算使用。
值得注意的是，由于 KV Cache 缓存了每一轮已计算完毕的键值向量，因此会额外增加显存开销。以 LLaMA-7B 模型为例，每个 token 对应的 KV Cache 空间大小可通过下面的公式来计算，这样我们就能更精准地掌握和预测显存使用情况。
cache_per_tokens = 2 * n_layer * n_head * d_head * dtype_byte_size
公式中的第一个因子 2 代表 Key/Value 两个向量，它们在每一层中都需要被存储。这里的 n_layer 表示 Transformer 层的数量，n_head 代表 KV head 的个数（在多头注意力模型中，这个值等于注意力头数；而在多查询注意力模型中，这个值为 1）。而 d_head 则代表每个 KV head 的维度，dtype_byte_size 则表示每存放一个数据所需的字节数。
模型推理所需的 KV Cache 总量计算公式如下：
其中 Context_length 表示输入和输出序列长度之和。因此，KV Cache 的大小与 batch size 和序列长度呈线性关系。
total_cache = batchSize * Context_length * cache_per_token
KV Cache 的引入巧妙地将推理过程划分为两个阶段，进而为后续的优化策略提供了更多可能性。
当谈到 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 技术，可有效避免显存溢出问题。
Transformer 类模型的核心自注意力机制在计算和存储上，其复杂度随序列长度呈二次增长，这成为了性能瓶颈。虽然之前的优化工作主要关注减少 FLOPs 以提升计算速度，但显存 IO 访问的优化却被忽视了。
事实上，GPU SRAM 的读写速度高达 19TB/s，而 GPU HBM 的读写速度却与之相差甚远，存储容量也相差巨大。Flash Attention 的设计充分考虑了这一点，它通过从 HBM 读取数据块到 SRAM 进行计算，然后再写回 HBM，从而有效避免了频繁从 HBM 中读取或写入注意力矩阵。
FlashAttention 的主要目标包括：
论文中提出的具体算法，通过精细的设计和高效的实现，成功达到了上述目标，为 Transformer 类模型的优化提供了新的思路。
GQA 将查询头巧妙地划分为 G 组，每组共享单一的键头和值头。
在将多头检查点转换为 GQA 检查点时，采用平均池化的策略，将组内的所有原始头进行合并，构建出每个组的键和值头。
中间组数的选择使得 GQA 在模型质量上插值于 MQA 和 MHA 之间，既保证了模型质量高于 MQA，又保证了推理速度优于 MHA。
从 MHA 过渡到 MQA，键值缓存的大小从 H 键和值头减少到单个键和值头，这意味着需要加载的数据量减少了 H 倍。然而，随着模型大小的增加，头部的数量通常也会增加，这使得多查询注意力在内存带宽和容量方面表现更佳。
但 GQA 的设计允许我们随着模型大小的增加，仍然保持相同的带宽和容量比例下降，从而实现了高效且灵活的注意力机制。
算子融合是深度学习模型推理的秘密武器，它通过减少计算过程中的访存次数和 Kernel 启动耗时，来显著提升模型推理性能。对于 LLM 推理来说，这一技术同样不可或缺。
以 HuggingFace Transformers 库中的 LLaMA-7B 模型为例，通过深入分析模型推理时的算子执行分布，我们发现该模型拥有 30 个类型共计 2436 个算子，其中 aten::slice 算子就高达 388 次。这么多小算子频繁执行，无疑会拖慢 GPU 的步伐，影响推理速度。
为了应对这一挑战，业界已经进行了许多尝试。以 DeepSpeed Inference 为例，算子融合主要聚焦在以下四个方面：
然而，算子融合并非易事，通常需要定制化实现算子的 CUDA kernel，这对 GPU 编程能力提出了较高要求。幸运的是，随着编译器技术的飞速发展，OpenAI Triton、TVM 等优秀框架已经能够实现算子融合的自动化或半自动化，让这一技术更加易于应用。
为了降低 LLM 推理的时延，我们专注于编写针对运行热点函数的高性能算子。以下是两个关键优化点：
实现这些高性能算子对 GPU 编程能力有着较高要求，并且算法实现中的某些超参数与特定问题规模密切相关。因此，编译器相关的自动调优技术也成为了我们研究的重点，以确保算子在各种场景下都能发挥出最佳性能。
为了提升服务性能，引入了两大关键优化策略：Continuous Batching 和异步 Tokenize / Detokenize。
| 问题分类 | 现象 | 解决方法 | 实现原理 | 特点 |
|---|---|---|---|---|
| 问题一 | 同批次序列推理时，存在'气泡'，导致 GPU 资源利用率低 | Continuous Batching | 由 batch 粒度的调度细化为 step 级别的调度 | 在时间轴方向动态插入新序列 |
| 问题二 | Tokenize / Detokenize 过程在 CPU 上执行，期间 GPU 处于空闲状态 | 异步 Tokenize / Detokenize | 多线程异步 | 流水线 overlap 实现降低时延 |
大语言模型的输入和输出长度可变，给服务带来了挑战。传统方法在同批次序列推理过程中，存在'气泡'现象，导致 GPU 资源浪费。而 Continuous Batching 技术通过动态调整 Running 队列中的序列，消除了这种浪费，实现了 GPU 资源的最大化利用。
此外，还引入了 Dynamic Batching 的概念。在自回归迭代生成每个 token 后，调度器会根据当前剩余显存量，动态调整 Running 队列的长度。这允许我们在 batch 维度上动态插入新序列，以充分利用显存空间。
这两项优化技术已经在 HuggingFace 的 Text-Generation-Interface (TGI)、vLLM、OpenPPL-LLM 等多个框架中得到实现，并显著提升了推理吞吐量。
由于大语言模型参数量庞大，单一计算设备可能难以承载，因此分布式并行技术显得至关重要。在 LLM 推理中，模型并行和流水线并行已成为得力助手。模型并行巧妙地将权重参数分散到多个计算设备上，实现高效的分布式计算。
各种并行术语小课堂：
在上图中，第一行展示了 Column Parallel，即将权重数据纵向分割到多个 GPU 中，结果需在列方向拼接；第二行则是 Row Parallel，权重数据横向分割，结果需通过 AllReduce 规约。
业界翘楚 Megatron-LM 提出了针对 Self-Attention 和 MLP 的简洁高效模型并行方案：
除了上述的并行方式，LLM 模型中的 Input Embedding 采用 Row Parallel，Output Embedding 采用 Column Parallel；而 Dropout、Layer Norm 等操作则保持完整，未进行并行拆分。
| Layers | Model Parallel Method |
|---|---|
| Input Embeding | Row Parallel |
| Self-Attention | Column Parallel + Row Parallel |
| MLP | Column Parallel + Row Parallel |
| Output Embeding | Column Parallel |
以 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 越大，通信效率越高，模型并行的收益也就越明显。
在提升 LLM 模型推理性能上，吞吐量和时延是两个关键指标。吞吐量受限于显存容量，而低比特量化技术正是解决显存占用和访存量的有效手段。通过降低显存占用量，我们可以运行更大的 batchsize，从而提升吞吐量；同时，由于 LLM 推理具有 Memory-bound 特性，降低访存量将在吞吐量和时延上带来双重收益。低比特量化的关键在于其节省的显存量和访存量，以及量化计算带来的加速远大于反量化带来的额外开销。
| 被量化的对象 | 量化方法 | 特点 |
|---|---|---|
| 权重量化 | LLM.int8()，GPTQ | 显存占用减半，但由于计算结果需要反量化，时延基本无收益 |
| 权重和激活同时量化 | Smooth Quant | 显存占用减半，时延有收益，精度几乎匹配 FP16 |
| KV Cache 量化 | INT8 或 FP8 量化 | 方法简单，吞吐量收益明显 |
| 基于硬件特点的量化：英伟达 Hopper 架构下的 FP8 | 直接利用 TensorCore FP8 计算指令 | 不需要额外的量化/反量化操作，时延收益明显 |
表中的四类量化方法各有特点，业界在低比特量化方向的研究日新月异，我们期待能够探索出一个适用于大语言模型的量化方法，它能在保证精度的同时，以较高的压缩率压缩模型，并加速端到端推理。
在 LLM 推理领域，除了引入传统优化技术外，业界也在不断探索新的解决方案。从大模型自回归解码的特点出发，一些新技术正在尝试通过调整推理过程和引入新的模型结构来进一步提升性能。
例如，投机采样（Speculative decoding）技术，它针对 LLM 推理的串行解码特点，引入了一个近似模型来执行串行解码，而原始模型则执行并行评估采样。通过近似模型和原始模型的配合，该技术能够在保证精度一致性的同时，降低大模型串行解码的次数，从而显著降低推理时延。这是一项令人兴奋的技术，有望为 LLM 推理性能带来质的飞跃。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online