LLM推理如何优化内存瓶颈和计算效率？（附KV缓存与TensorRT-LLM方案）

通过堆叠 Transformer 层构建的大型模型，在广泛的自然语言任务上展现出更高的准确性、更强的少样本学习能力，甚至接近人类的涌现能力。然而，这些基础模型的训练成本高昂，并且在推理阶段（一项持续性成本）可能对内存和计算资源要求极高。目前最流行的大型语言模型参数量可达数百亿甚至上千亿。根据具体用例，模型可能需要处理长输入（或上下文），这也会增加开销。例如，检索增强生成（RAG）流水线需要将大量信息放入模型输入中，这极大地增加了 LLM 需要处理的工作量。

本文将探讨 LLM 推理中最紧迫的挑战以及一些实用的解决方案。读者应对 Transformer 架构和注意力机制有基本的了解。

开发者也可以使用开源的 NVIDIA 和社区模型（例如 Nemotron 3 系列开放模型）以及运行在开源 TensorRT-LLM 库（可在 GitHub 上获取）来探索这些推理优化技术。这使得使用生产级代码而非抽象示例来实验现实世界的推理权衡成为可能。

理解 LLM 推理

大多数流行的仅解码器 LLM（例如 GPT-3）都是基于因果建模目标进行预训练的，本质上作为下一个词预测器。这些 LLM 将一系列词元作为输入，并以自回归方式生成后续词元，直到满足停止条件（例如，生成词元数量的限制或停止词列表），或者直到生成一个特殊的 <end> 词元标记生成结束。这个过程涉及两个阶段：预填充阶段和解码阶段。

请注意，词元 是模型处理的语言基本单位。一个词元大约相当于四个英文字符。所有自然语言输入在输入模型之前都会被转换为词元。

预填充阶段：处理输入

在预填充阶段，LLM 处理输入词元以计算中间状态（键和值），这些状态用于生成“第一个”新词元。每个新词元都依赖于之前的所有词元，但由于输入的完整范围是已知的，从高层次看，这是一个高度并行化的矩阵-矩阵运算。它有效地饱和了 GPU 的利用率。

解码阶段：生成输出

在解码阶段，LLM 以自回归方式逐个生成输出词元，直到满足停止条件。每个顺序输出的词元都需要知道之前所有迭代的输出状态（键和值）。这类似于矩阵-向量运算，与预填充阶段相比，它未能充分利用 GPU 的计算能力。数据（权重、键、值、激活）从内存传输到 GPU 的速度主导了延迟，而不是计算实际发生的速度。换句话说，这是一个内存受限的操作。

本文中介绍的许多推理挑战和相应解决方案都涉及优化这个解码阶段：高效的注意力模块、有效管理键和值等。

不同的 LLM 可能使用不同的分词器，因此，比较它们之间的输出词元可能并不直接。在比较推理吞吐量时，即使两个 LLM 每秒输出的词元数相似，如果它们使用不同的分词器，也可能不等效。这是因为对应的词元可能代表不同数量的字符。

批处理

提高 GPU 利用率和有效吞吐量的最简单方法是通过批处理。由于多个请求使用相同的模型，权重的内存成本被分摊。更大的批次被一次性传输到 GPU 进行处理，将利用更多可用的计算资源。

然而，批处理大小只能增加到某个极限，超过这个极限可能会导致内存溢出。为了更好地理解为什么会发生这种情况，需要查看键值缓存和 LLM 内存需求。

传统的批处理（也称为静态批处理）是次优的。这是因为对于批次中的每个请求，LLM 可能生成不同数量的完成词元，因此它们的执行时间也不同。结果，批次中的所有请求必须等待最长的请求完成，如果生成长度差异很大，这种情况会加剧。有一些方法可以缓解这个问题，例如动态批处理。

例如，像 NVIDIA TensorRT-LLM 这样的开源运行时具有针对流行开放模型（例如 Llama 和 Nemotron）的动态批处理和相关的调度优化。这不需要编写自定义调度器或 CUDA 内核。

键值缓存

解码阶段的一个常见优化是 KV 缓存。解码阶段在每个时间步生成一个词元，但每个词元都依赖于所有先前词元（包括在预填充阶段计算的输入词元的 KV 张量，以及直到当前时间步为止计算的任何新 KV 张量）的键和值张量。

为了避免在每个时间步为所有词元重新计算所有这些张量，可以将它们缓存在 GPU 内存中。每次迭代计算新元素时，只需将它们添加到运行中的缓存中，以供下一次迭代使用。在某些实现中，模型的每一层都有一个 KV 缓存。

图 1. 键值缓存机制示意图

LLM 内存需求

实际上，GPU LLM 内存需求的两个主要贡献者是模型权重和 KV 缓存。

• 模型权重：内存被模型参数占用。例如，一个拥有 70 亿参数的模型（如 Llama 2 7B），以 16 位精度（FP16 或 BF16）加载，将占用大约 7B * sizeof(FP16) ≈ 14 GB 的内存。
• KV 缓存：内存被用于缓存自注意力张量以避免冗余计算。

通过批处理，批次中每个请求的 KV 缓存仍然必须单独分配，并且可能占用大量内存。下面的公式描述了 KV 缓存的大小，适用于当今大多数常见的 LLM 架构。

每个词元的 KV 缓存大小（字节）= 2 * (层数) * (头数 * 头维度) * 精度字节数

第一个因子 2 代表 K 和 V 矩阵。通常，(头数 * 头维度) 的值与 Transformer 的隐藏大小（或模型维度，d_model）相同。这些模型属性通常可以在模型卡片或相关的配置文件中找到。

这个内存大小是针对输入序列中的每个词元，跨越输入批次所需的。假设使用半精度，KV 缓存的总大小由以下公式给出。

KV 缓存总大小（字节）= (批次大小) * (序列长度) * 2 * (层数) * (隐藏大小) * sizeof(FP16)

例如，对于一个使用 16 位精度的 Llama 2 7B 模型，批次大小为 1，KV 缓存的大小将为 1 * 4096 * 2 * 32 * 4096 * 2 字节，约等于 2 GB。

高效管理这个 KV 缓存是一项具有挑战性的任务。随着批次大小和序列长度线性增长，内存需求会迅速扩大。因此，它限制了可以服务的吞吐量，并对长上下文输入提出了挑战。这是本文介绍的几种优化背后的动机。

通过模型并行化扩展 LLM

减少模型权重在每个设备上的内存占用的一种方法是将模型分布在多个 GPU 上。分散内存和计算占用空间使得能够运行更大的模型或更大的输入批次。模型并行化对于训练或推理需要比单个设备可用内存更多的模型是必要的，并且可以使训练时间和推理指标（延迟或吞吐量）适合某些用例。根据模型权重的分割方式，有几种并行化模型的方法。

请注意，数据并行性也是经常与下面列出的其他技术在同一上下文中提到的技术。在这种技术中，模型的权重被复制到多个设备上，输入的（全局）批次大小被分片到每个设备上成为微批次。它通过处理更大的批次来减少总体执行时间。然而，这是一种训练时的优化，在推理时不太相关。

请注意，任何模型并行技术——包括流水线并行和张量并行——都可以在开源框架中使用，例如 NVIDIA Megatron-LM 和 NVIDIA NeMo 框架，它们为广泛的开放模型提供训练和推理工作流的支持。

流水线并行

流水线并行涉及将模型（垂直）分片成块，其中每个块包含在单独设备上执行的层子集。图 2a 是四路流水线并行的示意图，其中模型被顺序分区，所有层的四分之一子集在每个设备上执行。一个设备上的一组操作的输出被传递到下一个设备，该设备继续执行后续块。(F_n) 和 (B_n) 分别表示设备 n 上的前向和后向传递。每个设备上存储模型权重的内存需求实际上减少了四分之一。

这种方法的主要限制在于，由于处理的顺序性质，在等待前一层输出（激活、梯度）时，某些设备或层可能保持空闲。这导致前向和后向传递中的效率低下或“流水线气泡”。在图 2b 中，白色空白区域是朴素流水线并行中的大型流水线气泡，设备处于空闲状态且利用率不足。

微批次可以在一定程度上缓解这个问题，如图 2c 所示。输入的全局批次大小被分割成子批次，这些子批次被逐个处理，梯度在最后累积。请注意，(F_{n,m}) 和 (B_{n,m}) 分别表示设备 (n) 上微批次 (m) 的前向和后向传递。这种方法缩小了流水线气泡的大小，但并没有完全消除它们。

图 2. 四路流水线并行示意图。(a) 模型按层分区为 4 部分，每个子集在单独设备上执行。(b) 朴素流水线并行导致大的流水线气泡和 GPU 利用率不足。(c) 微批次减少了流水线气泡的大小，提高了 GPU 利用率。

常见问题（FAQ）

LLM推理中KV缓存具体如何提升效率？

KV缓存通过存储预填充阶段计算的键值状态，避免解码阶段重复计算，将内存受限的矩阵-向量运算转化为高效的内存访问，显著减少延迟并提升GPU利用率。

批处理策略对LLM推理性能有什么实际影响？

批处理能同时处理多个请求，将多个独立的矩阵-向量运算合并为更高效的矩阵-矩阵运算，提高GPU计算资源利用率，从而显著提升整体吞吐量并降低单位计算成本。

TensorRT-LLM如何帮助解决模型并行化问题？

TensorRT-LLM提供生产级框架，支持流水线并行和张量并行等模型并行化技术，可将超大模型拆分到多个GPU上运行，有效解决单卡内存瓶颈问题，实现大规模LLM的高效推理部署。