大型语言模型，特别是那些拥有数十亿参数 (parameter)的模型，不仅在计算方面，而且在内存容量和带宽方面都带来了巨大的挑战。尽管量化 (quantization)和剪枝等技术可以减小静态模型大小，但在推理 (inference)过程中管理动态内存占用，特别是激活和主要的键值（KV）缓存，对于在硬件上的高效部署非常重要。内存不足会导致失败，而低效的内存访问模式会造成明显的延迟瓶颈，其对推理时间的影响常常超过纯粹的计算时间。专门用于减轻LLM推理过程中这些内存压力的先进技术将被讨论。
推理 (inference)过程中的内存占用来源
了解内存分配的位置是优化的第一步。在推理过程中，主要的内存消耗者包括：
- 模型参数 (parameter)（权重 (weight)）： 这是模型的静态大小。虽然庞大（对于基础模型通常为几十到数百千兆字节），但权重通常只加载一次并保持不变。量化 (quantization)等技术直接处理这一部分。
- 激活： 这些是网络正向传播期间计算的中间结果（例如，注意力层、前馈网络的输出）。它们的大小取决于批次大小、序列长度、隐藏维度和模型深度。对于长序列或大批次，激活会消耗大量内存，可能超过权重大小。
- 键值（KV）缓存： 专门用于自回归 (autoregressive)生成，KV缓存存储了来自先前token的自注意力 (self-attention)层计算出的键和值。这避免了为每个新生成的token进行冗余计算。然而，缓存大小随生成的token数量线性增长（如果未优化，则随上下文 (context)长度二次增长）。对于长对话或文档摘要任务，KV缓存可能成为单一最大的内存消耗者，轻松超过模型权重。
- 工作区缓冲区： 深度学习 (deep learning)框架和库（如cuDNN、cuBLAS）为中间计算、算法选择或优化核执行分配的临时内存。它们的大小可能因操作和库实现而异。
不同情境下的内存使用示意。请注意，在长上下文情况下，即使是量化模型，KV缓存也占据了主导地位。
激活重计算（检查点）
激活重计算（或在训练语境中的梯度检查点）是一种主要在训练中知名的技术，它以计算量换取内存。在正向传播过程中，它不是存储所有计算梯度所需的中间激活，而是仅保存一部分激活（在策略性选择的“检查点”层）。检查点之间的激活在反向传播 (backpropagation)过程中按需重新计算。
虽然在推理 (inference)过程中通常不需要梯度，但如果激活内存本身成为限制因素，该核心思路可以进行调整，特别是在极长序列或大批次情况下，即使一次正向传播也可能超出可用内存。在这种情况下，中间激活可以被丢弃，并在正向传播的后续阶段或由后续处理阶段需要时重新计算。
- 权衡： 显著减少峰值激活内存。
- 代价： 引入大量计算开销（重新运行网络部分的正向传播），增加延迟。
- 推理适用性： 由于延迟惩罚，通常是推理的最后手段。在训练场景或高度专业化的推理流程中更为常见，这些场景中将模型适配到内存是绝对优先事项。
内存卸载
卸载是指在内存层次结构的不同层级之间移动数据，以释放更快但更有限的内存，例如GPU高带宽内存（HBM）。
-
CPU卸载： 不需要立即使用的张量（通常是模型权重 (weight)，但也可能是激活或KV缓存条目）通过PCIe总线从GPU HBM传输到主机系统的主内存（CPU RAM）。当再次需要时，它们会被传回。
- 优点： 显著增加模型可用的有效内存容量。使得原本无法适配HBM的模型得以运行。
- 缺点： PCIe带宽比HBM带宽低几个数量级，导致传输缓慢，如果管理不当，可能会引入明显的延迟空泡。需要复杂的调度来将计算与数据传输重叠。
- 使用场景： 为超大型模型顺序加载层，在长时间生成暂停期间卸载激活或部分KV缓存。像Accelerate（Hugging Face）这样的框架提供了自动权重卸载的实用工具。
-
磁盘/NVMe卸载： CPU卸载的延伸，数据被进一步下移到更慢的存储介质，如NVMe固态硬盘甚至传统硬盘。
- 优点： 大幅增加表观容量。
- 缺点： 由于存储访问时间，引入比CPU卸载更高的延迟。仅适用于对延迟不敏感的任务或专用设置。
- 使用场景： 研究环境，拥有超大型模型但RAM/HBM严重受限的系统。
有效的卸载严重依赖于预测接下来需要哪些数据，并异步预取它们，以将传输延迟隐藏在正在进行的计算之后。
高效的KV缓存管理
如前所述，KV缓存通常是长序列推理 (inference)过程中的主要内存瓶颈。因此，优化其管理具有很高的效果。
-
KV缓存量化 (quantization)： 专门对缓存中存储的键和值应用量化（例如，8位整数（INT8）、4位浮点数（FP4）或其他低位格式）。
- 优势： 直接将缓存的内存占用减少2倍（FP16到INT8）或更多。
- 考量： 需要仔细校准或微调 (fine-tuning)以最大程度地减少精度下降。KV缓存量化的影响有时可能比权重 (weight)量化更显著。对低精度格式的硬件支持是有利的。
-
注意力变体（减少缓存大小）： 架构修改可以从本质上减少KV缓存大小。
- 多查询注意力（MQA）： 使用所有查询头共享的单个键和值头。显著减小K和V张量的大小，从而减小缓存。
- 分组查询注意力（GQA）： 标准多头注意力 (multi-head attention)（MHA）和MQA之间的一种折衷。使用多个K/V头，但少于查询头的数量。在MQA的内存节省和MHA的潜在质量优势之间提供了平衡。
- 滑动窗口注意力： 只关注固定窗口内最近的token（例如Mistral, Mixtral）。这自然限制了所需的KV缓存大小，因为窗口外的键和值可以被丢弃。
-
PagedAttention： 由vLLM项目推广的一种复杂的内存管理技术。它从操作系统中的虚拟内存和分页中获得启发。
- 原理： PagedAttention不是将每个序列的KV缓存存储在连续的内存块中（这会导致碎片化和内存浪费），而是以固定大小、非连续的块（称为“页面”）分配缓存。一个块表将逻辑token位置映射到这些物理页面。
- 优势：
- 减少碎片： 显著减少内部和外部内存碎片，提高内存利用率（更接近理论容量）。
- 高效共享： 支持写时复制机制。例如，当多个请求共享一个共同的提示（如在束搜索或并行采样中），它们的初始KV缓存页面可以在内存中共享，避免冗余存储和计算。
- 灵活分配： 更优雅地处理可变序列长度。
PagedAttention的示意图。序列中的逻辑token位置通过块表映射到可能非连续的物理内存块（页面），这些块存储了KV缓存数据。
统一内存管理与池化
除了卸载或PagedAttention等特定技术，高效的底层内存管理也很重要。
- 自定义分配器： 用自定义池分配器替换默认的CUDA分配器（
cudaMalloc/cudaFree）可以减少开销。预先分配大内存池并在这些池中管理分配/释放，可最大程度减少对CUDA驱动程序的昂贵调用并减少碎片。
- 统一内存系统： 一些推理 (inference)框架（例如vLLM、TensorRT）实现了复杂的内存管理器，它们统一管理权重 (weight)、激活、KV缓存和工作区，通常结合了基于启发式方法或性能分析的池化和智能放置策略。
评估内存管理策略
选择和调整内存管理技术需要权衡复杂的考量：
- 内存节省： 主要目标。以减少的峰值内存使用量（GB）衡量。
- 延迟影响： 重计算增加计算延迟。卸载增加数据传输延迟。量化 (quantization)可能会为数据打包/解包增加少量开销。这必须进行衡量（例如，每个输出token的时间，总生成时间）。
- 吞吐量 (throughput)影响： 该策略如何影响可以并发处理的请求数量或系统的整体每秒token数。例如，PagedAttention通常通过更好的内存利用率增加批次大小来提高吞吐量。
- 实现复杂性： 某些技术（如复杂的卸载或自定义内存管理器）实现和正确调试起来很复杂。
- 准确性： 特别与KV缓存量化相关。确保所选技术不会根据下游任务指标不可接受地降低模型输出质量。
有效的内存管理并非单一技术，而是通常结合了根据特定模型、硬件限制（HBM大小、PCIe速度）和应用需求（延迟敏感度、序列长度）量身定制的方法。在实际负载下分析内存使用和推理 (inference)延迟，对于找出瓶颈并量化这些优化策略的好处非常重要。像vLLM、TensorRT和TGI这样的框架通常集成了这些技术中的几种，开箱即用地提供高性能推理，但了解其背后的机制有助于更好的配置和故障排除。
