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