PRESERVE: Prefetching Model Weights and KV-Cache in Distributed LLM Serving

发表时间: 2025-01 · arXiv:2501.08192 (preprint)

原文: https://arxiv.org/abs/2501.08192

作者/机构：Ahmet Caner Yüzügüler, Jiawei Zhuang, Lukas Cavigelli / Huawei Zurich Research Center Switzerland

速读

一句话结论 本文提出了 PRESERVE 图优化框架，通过在分布式大语言模型推理的通信阶段将权重和 KV Cache 从片外内存预取至片上 L2 缓存，在商用 AI 加速器上实现了最高 1.61 倍的端到端推理加速。

要解决什么问题 大语言模型在推理的解码阶段具有自回归特性，每次只能生成一个 token，这导致底层矩阵乘法操作的计算强度极低，整体性能被片外 HBM 的读取带宽严重受限。为了应对庞大的模型权重和 KV Cache，现代推理系统通常采用张量并行将模型切分到多个加速器设备上。然而，这种分布式架构引入了新的卡点：在每一层的计算之间，设备必须通过 Allreduce 等集合通信操作来同步中间激活值。在执行这些通信操作时，计算核心往往处于空闲等待状态，导致设备利用率低下。为了掩盖这一通信开销，现有的主流路线是算子融合（如将 Matmul 与 Allreduce 融合），试图让计算和通信重叠执行。但这种做法存在致命的机制缺陷：首先，由于严格的数据依赖，它只能融合紧挨着的两个操作，无法覆盖整个通信延迟；其次，在典型的 Transformer 架构中，自注意力操作和 Allreduce 之间还夹杂着其他操作，导致这种融合路线根本无法应用于体积庞大的 KV Cache；最后，算子融合需要极高的底层 Kernel 开发工程代价，且在小批量切块时会导致内存访问效率低下。

怎么做的 PRESERVE 框架的核心思路是：在设备执行 Allreduce 通信操作的等待期间，利用并行的硬件流，提前将下一层所需的模型权重和 KV Cache 从片外 HBM 预取到片上 L2 缓存中。这一设计之所以能绕开上述数据依赖卡点，是因为在解码阶段，虽然线性投影层的输入（激活值）需要等待上一层的 Allreduce 结果，但其对应的模型权重是完全只读的；同理，自注意力层中用于计算的绝大部分历史 KV Cache 也是只读的（仅最后一个位置需要更新）。因此，这些庞大的只读数据完全可以与通信操作并行拉取。当 Allreduce 完成、计算核心准备执行前向传播时，所需数据已经驻留在带宽比 HBM 高出数倍的 L2 缓存中，从而打破了内存带宽瓶颈。为了在不修改用户代码的前提下实现这一机制，PRESERVE 被设计为一个图优化框架。它接收由深度学习框架编译出的计算图中间表示，并在图编译器层面自动插入预取算子。其关键设计包含一个基于广度优先搜索的算子插入算法：框架会遍历计算图，找到所有的通信节点，并向下搜索类型为 MatMul 或 SelfAttention 的子节点。为了防止无节制预取导致 L2 缓存被污染（即有用数据被提前驱逐），PRESERVE 在编译期会严格追踪已预取数据的内存占用。定义目标硬件的 L2 缓存容量为 $C$，节点 $n$ 所需的内存大小为 $M(n)$，框架维护的关键不变量为当前累积的预取数据量 $S_{prefetch}$ 必须严格小于缓存容量：

$$S_{prefetch} = \sum_{n \in \text{prefetched}} M(n) < C$$


只有当该条件满足时，框架才会在与通信操作平行的执行流中插入预取指令；一旦超出 $C$，则停止当前通信节点下的预取插入。在运行时，设备调度器会将预取算子分发到独立流中，并通过事件指令与主计算流进行层间同步，从而透明地实现通信与内存读取的重叠。

效果如何 实验在配备 8 张 Ascend 910B NPU（单卡 192MB L2 缓存、64GB HBM）的 Huawei Atlas 800T A2 服务器上搭建，测试了 Llama3-8B/70B、Qwen2-7B/72B 和 Phi-3 等开源模型。对比基线包含两种：一是完全串行执行的常规推理，二是代表算子融合路线的强基线方法（调用 npu_mm_all_reduce_base 实现的 Fused GEMM + Allreduce）。在常规设置下（Batch Size 为 4，上下文 16K），PRESERVE 相比串行基线在所有模型上均实现了端到端提速，最高在 4 卡并行的 Llama3-8B 和 Phi3-small 上测得 1.61 倍的加速比。在与算子融合路线的直接对比中，PRESERVE 在小批量（Batch Size 小于 512）场景下展现出显著优势，最高比 Fused GEMM + Allreduce 快 1.19 倍，因为它避免了小数据块切分带来的访存低效，且成功覆盖了 KV Cache。此外，作者还通过性能模型进行了硬件设计空间探索，结果表明，当引入预取机制后，AI 加速器最优的 L2 缓存大小从 8MB 激增至 104MB，采用该最优配置能在抵消芯片面积成本后，额外获得 1.25 倍的吞吐量密度（性能/成本）提升。该方法的局限性在于对 L2 缓存容量的绝对依赖：当 Batch Size 达到或超过 512，或者序列长度达到极端的 32K/64K 时，急剧膨胀的 KV Cache 会突破 L2 容量上限，导致预取失效，此时其性能会回落至不如算子融合基线的水平（最低降至 0.93 倍），因此该方法更适合对延迟敏感的在线推理而非离线大吞吐场景。

主要贡献

当前大型语言模型（LLMs）通常部署在包含大量设备（如GPU/NPU）的集群上进行推理服务。然而，这些设备之间的通信会产生显著的开销，这不仅增加了推理的延迟和成本，还限制了系统的可扩展性。目前解决这一问题的方法主要是将通信与计算重叠（例如融合矩阵乘法和Allreduce操作），但由于这些操作之间存在数据依赖性，这些方法存在严重的局限性，无法掩盖整个通信延迟，也无法应用于KV缓存。

为了解决上述问题，本文提出了以下核心贡献：
* 提出PRESERVE框架：这是一种新颖的模型权重和KV缓存预取框架。该框架能够在集合通信（Collective Communication）操作期间，将模型权重和KV缓存从片外高带宽内存（HBM）预取到AI加速器的片上缓存（L2 Cache）中，从而利用内存读取操作来掩盖通信延迟。
* 无缝的图优化方案：提出了一种图优化机制，能够自动且最优地将预取操作插入到LLM推理的计算图中。这使得开发者无需修改任何用户代码即可获得性能提升，同时框架会在编译时管理流同步和缓存状态，防止缓存污染。
* 显著的端到端性能提升：在商用AI加速器上进行的大量实验表明，针对开源的最先进LLMs，PRESERVE能够实现高达 $1.6\times$ 的端到端推理加速。
* 最优硬件配置的设计空间探索：通过设计空间探索（Design Space Exploration, DSE）确定了采用该预取方法时AI加速器的最优硬件配置。结果表明，当考虑到预取机制时，最佳的L2缓存大小从8 MB增加到了 $104 \mathrm{MB}$，通过选择最优的L2缓存大小，可以使性能成本比进一步提升 $1.25\times$。

概念对比图：A) 常规推理，B) 采用GEMM+Allreduce融合内核的推理，C) PRESERVE（本工作）
概念对比图：A) 常规推理，B) 采用GEMM+Allreduce融合内核的推理，C) PRESERVE（本工作）

背景知识

LLM推理工作负载 如今大多数最先进的生成式AI大模型（例如GPT-4【33, GPT-4 Technical Report + 2023 + arXiv + https://arxiv.org/abs/2303.08774】和Llama2 【41, Llama 2: Open Foundation and Fine-Tuned Chat Models + 2023 + arXiv + https://arxiv.org/abs/2307.09288】）都采用了仅解码器（decoder-only）的Transformer架构 【44, Attention is All you Need + 2017 + NeurIPS + https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html】。这些模型由堆叠的解码器层组成，每层包含一个自注意力（self-attention）层和一个多层感知机（MLP）层，并且都跟随着归一化层。自注意力层可进一步分解为三个输入投影层、位置嵌入、自注意力内核（例 如Flash attention【8, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness + 2022 + NeurIPS + http://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html】）以及一个输出投影层。同样，MLP层可分解为三个投影层和一个激活层。Transformer架构中很大一部分内存占用和计算工作量来自于投影层和自注意力内核，这使它们成为LLM加速的重点。使用LLM的生成过程分为预填充（prefill）和解码（decode）两个阶段。在预填充阶段，LLM处理输入提示并生成第一个输出Token；在解码阶段，由于自回归特性，LLM基于过去的Token逐个生成Token，直到生成特殊的序列结束Token或达到最大Token数限制。因为在大多数生成式AI应用中输入提示的长度很容易超过数百或数千个Token，典型预填充阶段中线性层和自注意力内核的矩阵乘法操作极大地受益于数据重用，这增加了它们的运算强度（以FLOPS/bytes为单位）并最大化了计算资源的利用率 。

LLM解码受限于内存带宽 与预填充阶段相反，解码阶段自注意力层中的代数运算大多是矩阵与一个非常窄的矩阵（通常宽度不超过8）相乘，因此提供的运算强度为 $16 \mathrm{Op}/$word，而大多数加速器的屋顶线（roofline）在 $>100 \mathrm{Op}/$word。同样，解码阶段投影层中的数据重用仅限于批次（batch）大小，这也通常导致较低的运算强度。因此，在各种应用场景中（例如移动平台或有延迟约束的在线服务），由于传入查询率低或内存容量有限，批次大小可能被限制在一个很小的数字，解码阶段会遭受内存带宽瓶颈的困扰。结果就是，具有LLM的自回归生成式AI应用通常被认为是一个受内存带宽限制的过程。为了消除解码阶段自回归步骤中的冗余计算，计算出的键（key）和值（value）张量被存储在内存中以便在后续迭代中重用，这种机制通常被称为KV缓存（KV-cache）。在具有长上下文长度的生成任务中（例如书籍摘要或多模态模型），KV缓存的大小很容易变得与模型权重相当甚至超过模型权重。此外，对于长上下文长度，从片外内存读取KV缓存可能比读取模型权重花费更长的时间。因此，提高处理和读取KV缓存的性能对LLM的整体性能至关重要。基于此，本工作不仅考虑模型权重的预取，还考虑了KV缓存的预取。

多设备推理 如今许多最先进的LLM模型大小达到数千亿个参数。即使使用降低精度的格式（如bfloat16, float8, int8），这些模型的内存占用也超过了单个现代AI加速器设备的HBM内存容量。此外，单个设备的吞吐量通常受到HBM带宽的限制，高端AI加速器的HBM带宽在几TB/s的数量级。由于当今最先进的模型在模型大小和KV缓存上超过数百GB，单个设备不足以实现确保响应式用户体验所需的单Token延迟目标（通常设定为每Token $100 \mathrm{ms}$）。因此，LLM推理通常是分布式的，并在多个紧密互连的设备上并行执行。在分布式LLM推理中，提出了各种并行化策略以最小化设备间的通信开销。最常用的并行化策略之一称为张量并行（tensor-parallelism）。在这种并行化技术中，LLM中的权重和KV缓存被划分并分布在各个设备上。在推理期间，每个设备使用其权重和KV缓存的划分在本地执行计算，并通过通信原语 allreduce 对它们的结果进行求和。为了最小化通信开销，自注意力层通常在注意力头维度上进行划分，其中每个注意力头的计算相互独立，直到自注意力层的输出投影层【37, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism + 2019 + arXiv + http://arxiv.org/abs/1909.08053】。同样，MLP层也在其中间维度（对应于输入和输出线性投影的列和行）进行划分，使得大部分计算可以在本地执行。作为这种划分方案的结果，加速器每个自注意力层和MLP层只需要执行一 次 allreduce 调用。尽管分布式推理每层只需要一次 allreduce 调用，但根据数据和集群大小，它们的执行可能需要相当长的时间。需要在加速器之间通信的数据大小随着LLM架构的嵌入（embedding）大小和输入提示的批次大小线性增长。此外，网络流量也随着参与张量并行执行的加速器数量线性增长。虽然现代LLM推理服务器配备了设备之间的高带宽互连结构（如NVLink和HCCS），但扩展到单个服务器之外仍然需要通过较慢的网络（如PCIe或InfiniBand）执行 allreduce 调用。因此，allreduce 调用的执行时间可能占据总时间的一大部分，并成为分布式LLM推理系统可扩展性的限制因素。总结来说，当今LLM有两个限制其推理性能的显著计算特征：第一，底层操作表现出低运算强度，使得LLM推理的单设备性能受限于内存带宽；第二，由于其巨大的内存占用和延迟约束，LLM必须分布在多个设备上，这带来了显著的通信开销。这两个计算挑战限制了AI推理系统的性能和效率，增加了每Token的成本，并阻碍了它们的可扩展性和响应性，特别是在实时应用中。

各LLM在不同设备数量下Attention和MLP层的内存占用图
各LLM在不同设备数量下Attention和MLP层的内存占用图

方法细节

常规执行模式 在常规执行中，注意力层的查询（query）、键（key）和值（value）线性投影与前面的 allreduce 操作存在数据依赖关系，因为它们需要等待上一层激活值（activations）的规约结果。同样，MLP层的门控（gate）和向上（up）投影也依赖于前面的 allreduce 操作。由于这种依赖性，现代机器学习框架（如Pytorch）会在单个流中顺序执行这些操作。不幸的是，这导致设备在等待 allreduce 操作完成时处于空闲状态，造成了资源利用率低下和较长的延迟。

采用提出方法的执行模式 虽然线性投影需要因为激活值的数据依赖性而等待 allreduce 操作，但它们的权重是只读的，因此它们与前面的操作没有任何依赖关系。同样，在解码阶段，自注意力层中 $Q\bar{K}^T$ 和其他计算所使用的 $K$ 和 $V$ 缓存也是只读的（除了最后一个条目需要根据前面线性投影计算出的键和值进行更新）。因此，可以在执行这些层的前向传播之前预取权重和KV缓存。通过提出的方法，设备在等待 allreduce 操作的同时，开始在并行流中将权重和KV缓存从片外HBM内存预取到片上缓冲区（如L2缓存）。当 allreduce 操作完成且设备准备好执行下一层时，权重和KV缓存已经被获取到了L2缓存中，而L2缓存通常提供比片外HBM内存高 $5-10\times$ 的带宽。结果就是，线性投影和自注意力的前向传播执行得快得多，从而提高了资源利用率并加速了解码阶段。

PRESERVE框架概述图
PRESERVE框架概述图

L2缓存容量需求 提出的方法要求AI加速器将后续层的预取权重和KV缓存存储在其片上内存（通常是L2缓存）中。如果加速器的L2缓存大小不足以存储这些权重和KV缓存，预取的数据将被驱逐（evicted），从而降低提出方法的有效性。因此，AI加速器应具有足够的L2缓存大小来存储 allreduce 操作之间所有层的权重和KV缓存。为了评估提出方法在流行LLM上的可行性，作者分析了各个LLM层的内存需求，并将其与现代AI加速器的L2缓存大小进行了比较（如图2所示）。随着设备数量的增加（前提是设备数量不超过注意力头的数量），权重和KV缓存被划分为更小的块，从而减少了每台设备的内存大小。图2表明，高达16个设备的张量并行足以将即使是最大模型的Attention和MLP层的内存需求降低到可以适应现有商用AI加速器L2缓存的水平。现代LLM推理系统通常包含更多的设备，例如Deepspeed-inference可扩展到256个GPU【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale + 2022 + SC22 + doi:10.1109/SC41404.2022.00051】，Cloud TPUv5e推理实例包含多达256个芯片【43, Expanding our AI-optimized infrastructure portfolio: Introducing Cloud TPU v5e and announcing A3 GA + 2023 + Google Cloud Blog + https://cloud.google.com/blog/products/compute/announcing-cloud-tpu-v5e-and-a3-gpus-in-ga】。因此，考虑到当前LLM推理系统的规模和商用AI加速器的L2缓存容量，提出的方法适用于大多数广泛使用的开源LLM 。

框架集成 提出的方法如何向用户暴露对其有效性和适用性至关重要。提出方法的一种简单实现是开发一个预取操作作为自定义算子，并将插入预取算子的任务交给程序员。然而，这种方法有三个缺点：首先，程序员需要显式地将预取算子与通信操作并行插入，这增加了程序员的负担；其次，程序员需要对底层硬件平台的内存层次结构有很好的了解，才能发挥提出方法的全部潜力；第三，以此方式编写的程序将特定于某些硬件平台，无法在不降低性能的情况下轻松迁移到不同的平台。因此，提出的方法不向程序员暴露预取，而是自动在ML应用的计算图中插入所需的算子。当今许多硬件供应商提供图优化框架（例如CUDA Graphs【29, Getting Started with CUDA Graphs + 2019 + NVIDIA Blog + https://developer.nvidia.com/blog/cuda-graphs/】 、CANN Graph Engine【17, CANN Community Edition + 2024 + Huawei + https://www.hiascend.com/en/software/cann/community】等），这些框架接收用硬件无关的高级编程语言（如Pytorch）编写的程序，并执行编译器级别的优化（如算子插入、融合、移除）。此外，此类图优化框架通常具有整个图的完整视图，这使得能够优化以最佳方式使用内存资源。为此，作者开发了PRESERVE框架，它将权重和KV缓存的预取算子插入到LLM推理的计算图中，同时自动管理流同步和缓存优化，以最小化缓存污染和冗余的内存带宽使用 。

PRESERVE框架工作流 图3展示了提出的PRESERVE框架的概述。在主机（host）侧，框架接收用户代码（使用Pytorch或Tensorflow等支持的ML框架编写的LLM模型推理实现）。然后，用户代码使用ML框架的内置编译器（例如TorchDynamo）或特定硬件的编译器（例如Ascend ATC【16, ATC tool + 2024 + Huawei + https://www.hiascend.com/document/detail/en/canncommercial/700/inferapplicationdev/atctool/atlasatc_16_0005.html】）编译成中间表示（IR）。接着，IR被传递给图优化器（例 如Ascend Graph Engine (GE)【26, DaVinci: A Scalable Architecture for Neural Network Computing + 2019 + Hot Chips + doi:10.1109/HOTCHIPS.2019.8875654】），图优化器执行图级别的优化，并将合适的计算和通信操作的预取算子插入到图中。最后，优化的图使用供应商特定的算子库（例如Ascend CANN【17, CANN Community Edition + 2024 + Huawei + https://www.hiascend.com/en/software/cann/community】）编译成离线模型（可执行二进制文件）。在接收到请求时，PRESERVE加载离线模型，将输入数据和模型权重复制到设备内存，并使用任务队列将编译后的算子发送到设备。当设备中的运行时调度器接收到编译后的算子时，它会根据算子被编译的执行顺序和流来分发算子。预取算子在与主流并行的流中执行，并使用事件（event）指令在层之间进行同步，以促进通信和预取之间的重叠 。

算子插入算法 算法1描述了用于权重和KV缓存预取的算子插入算法。该算法将计算图 $g$ 和目标硬件平台的L2缓存容量 $C$ 作为输入。接着，算法查找并遍历计算图中的所有通信操作 $O$。对于每个通信操作，算法以广度优先（BFS）顺序访问子节点 $n$，直到到达图的末尾或下一个通信操作（通过设置 reachedEOG = true 标志）。对于类型为 MatMul 或 SelfAttention 的每个子节点，算法估计预取所需的内存大小 n.memSize()，并计算当前节点和通信操作之间所有前面的预取算子的总内存大小 cacheSum。如果总内存大小小于L2缓存容量 $C$，则在与通信操作并行的流中插入当前节点的权重或KV缓存的预取算子 $P$。如果总内存大小超过了L2缓存容量，为了避免缓存驱逐，预取算子将不会被插入，算法跳出当前循环并继续处理下一个通信操作。通过这种方式，PRESERVE在运行时允许AI加速器在等待通信操作时开始读取权重和KV缓存，从而掩盖通信延迟并提高LLM推理性能。

实验设置

  • 模型与数据集配置：选用了社区广泛采用的开源LLMs，包括Llama3-8b、Llama3-70b、Qwen2-7B、Qwen2-72B、Phi-3-small和Phi-3-medium。批次大小（batch size）在1到64之间变化，序列长度在 $2\mathrm{k}$ 到 $32\mathrm{k}$ 之间变化。假设采用静态批处理，序列长度相等，预填充（prefill）和解码（decode）长度分别占序列总长度的2/3和1/3。
  • 软件配置：所有基准测试均在Pytorch中实现，使用 torch-npu 后端，基于 torchair 库中提供的参考实现，并使用TorchDynamo进行编译。并行执行通过 torch.distributed 库提供的通信原语实现。提出的方法使用CANN的Graph Engine (GE) 特性实现。激活值和权重均采用int8量化格式。
  • 硬件配置：实验在一台华为Atlas 800T A2服务器上进行，该服务器配备4个Kunpeng 920 48核CPU和8个Ascend 910B NPU。每个NPU拥有 $192\mathrm{MB}$ 的L2缓存，配备24个DaVinci AI核心和64 GB HBM内存，提供 $800 \mathrm{TOPS}/s$ 的理论int8算力和1.6TB/s的片外内存带宽。服务器中的Ascend 910B NPU通过HCCS互连结构以全网状（full-mesh）拓扑紧密相连。

实验结果

端到端执行时间 Table 1总结了在不同基准测试和集群大小（NPU数量）下，基线实现和提出方法的端到端LLM推理执行时间。结果显示，提出的方法在所有基准测试和集群大小中都加速了执行。加速比通常随着NPU数量的增加而增加，这是因为随着NPU数量的增加，每个设备的计算时间减少，而设备间的通信时间增加。因此，掩盖通信操作的延迟对执行时间有更大的影响，从而提高了加速比。作者观察到，除了Qwen2-7B（最大加速比出现在2个NPU时），最大加速比大多出现在NPU数量等于4时。这是因为当这些模型被划分到8个设备时，每个设备的KV头（KV-heads）数量等于1。当每个设备的KV头数量等于1时，从HBM读取KV缓存的内存操作是连续的而不是跨步的（strided）；因此，HBM带宽得到了更好的利用，内存操作变得更快，给提出的预取方法留下的改进空间较小。因此，当每个设备的本地KV头数量等于2时，观察到了最高的加速比。在Llama3-8B和Phi3-small模型上获得了最高加速比（$\sim 1.6\times$），高于它们的大型对应版本（Llama3-70B为 $1.36\times$，Phi3-medium为 $1.43\times$）。这是由于小模型中的通信计算比更高，因此将内存读取与通信重叠会导致更高的加速比。总体而言，实验证明该方法显著改善了各种最先进模型在不同NPU数量下的端到端执行时间，加速比从 $1.09\times$ 到 $1.61\times$ 不等。

批次大小与序列长度的影响 由于KV缓存及其对计算特征和内存需求的影响，提出方法的有效性对批次大小和最大序列长度很敏感。图4展示了在4个NPU上，批次大小从1到256，序列长度从1k到 $64\mathrm{k}$ 变化时，三个最大模型的加速比。结果表明，在给定的范围内，该方法实现了高达 $1.82\times$ 的加速。最大加速比出现在长序列长度和小批次大小的情况下。加速比通常随着批次大小的增加而降低，原因有二：首先，随着批次大小增加，MLP层中的计算变得不再那么受内存限制，这降低了方法在MLP层中的有效性；其次，KV缓存的L2需求随批次大小增加。因此，对于大批次大小，KV缓存无法装入L2，导致自注意力算子没有加速。同时，加速比通常随着序列长度的增加而增加，直到达到某个阈值。这是因为自注意力层的内存访问模式是不规则和跨步的，因此在自注意力层中获得的加速比高于MLP层。增加序列长度使自注意力层在总执行时间中占据更大比例。然而，KV缓存的内存需求也随序列长度增加，当超过阈值导致KV缓存不再适合L2时，该方法对自注意力层失效，加速比显著下降。尽管如此，这种下降仅在非常大的序列长度（如批次大小为4时序列长度为32k）下发生，覆盖了广泛的实际应用范围。

不同批次大小和最大序列长度下所获加速比的热力图
不同批次大小和最大序列长度下所获加速比的热力图

与计算-通信重叠融合内核的对比 为了展示提出方法相对于将GEMM和Allreduce操作重叠的融合内核（先前方法）的性能优势，作者将PRESERVE与一个基线进行了比较。在基线中，每个Attention和MLP层的最后一个线性层（分别为 $W_{out}$ 和 $W_{down}$）与随后的Allreduce操作融合（使用 torch_npu 库的 npu_mm_all_reduce_base 算子）。图5显示，PRESERVE在小批次大小下优于基线，而基线在大批次大小下表现更好。PRESERVE在小批次大小下具有优势的原因在于：首先，PRESERVE允许将任何权重和KV缓存的HBM读取与通信重叠，而融合算子只能融合连续的操作，完全忽略了KV缓存；其次，融合算子通常将matmul操作分成多个块并通过流水线执行，但小数据块通常导致内存访问效率低下。因此，PRESERVE在小批次大小下优于融合基线，加速比高达 $1.19\times$。相反，由于计算强度的提高和KV缓存大小的增长，PRESERVE的有效性随着批次大小的增加而逐渐减弱。在批次大小等于或大于512时，基线变得比PRESERVE更快（图5中小于1的加速比，最小值为 $0.93\times$）。因此，PRESERVE更适合批次大小小于512的场景（如在线和资源受限的推理），而融合算子仍然是批次大小等于或大于512（如离线推理）的首选方法。

与融合GEMM+Allreduce算子基线的对比图
与融合GEMM+Allreduce算子基线的对比图

补充细节

系统设计空间探索（DSE）
为了调查硬件设计参数如何影响PRESERVE框架的有效性，作者开发了AI加速器针对Transformer架构的性能和成本模型。计算吞吐量被建模为加速器flops容量的线性函数。内存读写吞吐量被假定为片外内存或片上L2缓存带宽的线性函数。通信延迟被建模为恒定初始延迟加上数据大小除以链路带宽的总和。假设集群中的加速器以环形拓扑互连。使用了7nm工艺节点的数据，核心和L2 SRAM的面积参数取自Lin等人【27, 7.1 A 3.4-to-13.3TOPS/W 3.6TOPS Dual-Core Deep-Learning Accelerator for Versatile AI Applications in 7nm 5G Smartphone SoC + 2020 + ISSCC】的设计。

  • L2缓存大小的性能影响：图6展示了归一化到最小L2容量的各模型推理延迟。随着L2缓存大小的增加，由于可以将更多数据预取到L2，延迟减少了 $20\%$ 到 $36\%$（平均 $25\%$）。大多数模型的最大加速比在 $136\mathrm{MB}$ 和 $144\mathrm{MB}$ 的L2缓存大小处达到，超过此值后没有进一步的加速。
  • 吞吐量-面积权衡：分配更大的L2缓存会增加硅面积，提高成本。为此，作者设计了“吞吐量密度”指标（Token/s除以总裸片面积 $\mathrm{mm}^2$）。图7显示，不启用预取的基线加速器设计的吞吐量密度在8 MB L2缓存时达到峰值。相比之下，使用预取的加速器的吞吐量密度在L2缓存大小高达 $104\mathrm{MB}$ 时持续增加，因为更大的L2缓存掩盖了更多的通信开销。即使以增加L2缓存大小为代价，预取更多数据也能将吞吐量密度从11.8提升至 $14.8\mathrm{token}/s/\mathrm{mm}^2$，对应 $1.25\times$ 的改进。
  • 网络带宽影响：图8显示了提出方法的加速比与网络带宽的关系。增加带宽会降低allreduce延迟，从而能够掩盖更大比例的通信开销。对于包含128个设备的集群，将设备到设备链路带宽从200 Gbps增加到1000 Gbps，平均加速比从 $1.18\times$ 提升至 $1.27\times$。

各种LLM的推理延迟与L2缓存大小的关系图
吞吐量密度与L2缓存大小的关系图（启用与不启用预取）
预取带来的加速比与设备间网络带宽的关系图

相关工作对比
先前的研究提出了各种优化技术。例如Flash Attention【8】和PagedAttention【22】通过内核优化和内存管理缓解内存瓶颈，但单设备性能仍受限。分布式推理通过张量并行【37】等减少了单设备负载，但引入了通信瓶颈。通过融合内核将计算与通信重叠的方法【5, 12, 34, 36, 45】有局限性：划分小块降低计算效率、需要大量底层内核开发工程，且只能与紧接的下一个计算操作重叠（无法掩盖KV缓存读取）。其他软件预取技术（如PrefetchML【7】、ZeRO-Inference【4】）主要是将数据从CPU/NVMe预取到GPU内存，没有利用片上L2缓存来最大化数据局部性并掩盖设备间通信。

局限性
提出方法的有效性受限于AI加速器的片上内存能够容纳多少模型权重和KV缓存。因此，任何增加每层每设备模型权重和KV缓存大小的因素（如批次大小、序列长度、精度）都可能影响加速比。然而，在线推理（如聊天机器人）的紧凑延迟约束通常对批次大小设置了上限；对于长序列，通常采用序列并行来减少每设备的KV缓存；更积极的量化方法（如AWQ【28】、KVQuant【13】）进一步降低了内存需求。因此该方法仍适用于广泛的应用场景。

结论

本文介绍了PRESERVE，这是一个新颖的模型权重和KV缓存预取框架，能够将HBM读取与集合通信操作重叠，从而掩盖通信延迟。PRESERVE使用图优化算法自动将预取命令与计算图中的通信集合操作并行插入，同时防止缓存污染，且无需修改用户代码。在商用AI加速器上的实验表明，该方法在最先进的LLM上实现了高达 $1.6\times$ 的端到端加速。此外，设计空间探索表明，当考虑预取机制时，AI加速器的最优L2缓存大小从8 MB转移到了 $104 \mathrm{MB}$，在最优L2缓存大小下，PRESERVE使性能成本比进一步提高了 $1.25\times$。研究得出两个关键结论：将模型权重和KV缓存从HBM预取到L2缓存是掩盖分布式LLM推理中通信开销的有效方法；在硬件设计时考虑预取机制能显著提高AI加速器的性能成本比。

附录

将设计空间探索扩展至L2带宽和计算单元
由于模型权重和KV缓存被预取到L2缓存中，预取数据被消耗的速度对提出方法的有效性也至关重要。在大多数加速器中，数据首先从片外内存获取到L2缓存，然后在操作执行期间获取到计算核心。由于提出的方法在操作执行之前将数据预取到L2缓存，因此在计算期间访问L2的延迟会暴露出来，从而阻碍性能提升。图9绘制了针对不同立方体吞吐量（cube throughput）和L2总线带宽，预取方法相对于基线加速器的加速比。随着立方体吞吐量的增加，计算密集型操作的执行时间减少，加速比也随之增加。此外，增加L2带宽可将提出方法的加速比提高多达 $10\%$。这是因为计算期间暴露的L2访问延迟变得更短，提高了方法的有效性。因此，L2总线的访问延迟和带宽是采用该预取方案的加速器的重要设计考量。

不同L2带宽下，平均加速比与立方体吞吐量的关系图
不同L2带宽下，平均加速比与立方体吞吐量的关系图

多节点横向扩展（Scale-out）
作者使用加速器性能模型将分析扩展到单个服务器之外，以评估提出方法在横向扩展系统上的性能增益。理论上，当allreduce延迟和预取延迟相等时，会出现最大加速比。图10显示了在不同张量并行设备数量下，预取方法相对于基线加速器的加速比，以及Attention和FFN层的allreduce和预取延迟。观察到加速比随着设备数量的增加而增加，直到32，超过32后开始下降。在设备数量较少（8和16）时，FFN层的预取延迟明显长于allreduce延迟，限制了加速比。随着设备数量增加，借助于张量并行，每个设备的FFN层权重大小减小，FFN的预取延迟也随之减少，并在集群大小为32时与allreduce延迟相匹配，此时观察到最高加速比。将设备数量增加到超过32的集群大小会导致预取延迟短于allreduce延迟，在这种情况下，后者无法被完全掩盖，导致加速比下降。

不同设备数量下，加速比以及Attention和FFN层的allreduce与预取延迟关系图
不同设备数量下，加速比以及Attention和FFN层的allreduce与预取延迟关系图

参考文献引用汇总：
* 【4, DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale + 2022 + SC22 + doi:10.1109/SC41404.2022.00051】
* 【5, FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion + 2024 + arXiv + doi:10.48550/ARXIV.2406.06858】
* 【7, PrefetchML: a framework for prefetching and caching models + 2016 + MODELS + http://dl.acm.org/citation.cfm?id=2976775 】
* 【8, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness + 2022 + NeurIPS + http://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html 】
* 【12, Overlapping Communication and Computation with High Level Communication Routines + 2008 + CCGrid + doi:10.1109/CCGRID.2008.15】
* 【13, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization + 2024 + NeurIPS + http://papers.nips.cc/paper_files/paper/2024/hash/028fcbcf85435d39a40c4d61b42c99a4-Abstract-Conference.html 】
* 【16, ATC tool + 2024 + Huawei + https://www.hiascend.com/document/detail/en/canncommercial/700/inferapplicationdev/atctool/atlasatc_16_0005.html 】
* 【17, CANN Community Edition + 2024 + Huawei + https://www.hiascend.com/en/software/cann/community 】
* 【22, Efficient Memory Management for Large Language Model Serving with PagedAttention + 2023 + SOSP + doi:10.1145/3600006.3613165】
* 【26, DaVinci: A Scalable Architecture for Neural Network Computing + 2019 + Hot Chips + doi:10.1109/HOTCHIPS.2019.8875654】
* 【27, 7.1 A 3.4-to-13.3TOPS/W 3.6TOPS Dual-Core Deep-Learning Accelerator for Versatile AI Applications in 7nm 5G Smartphone SoC + 2020 + ISSCC】
* 【28, AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration + 2024 + MLSys + https://proceedings.mlsys.org/paper_files/paper/2024/hash/42a452cbafa9dd64e9ba4aa95cc1ef21-Abstract-Conference.html 】
* 【29, Getting Started with CUDA Graphs + 2019 + NVIDIA Blog + https://developer.nvidia.com/blog/cuda-graphs/ 】
* 【33, GPT-4 Technical Report + 2023 + arXiv + https://arxiv.org/abs/2303.08774 】
* 【34, Optimizing Distributed ML Communication with Fused Computation-Collective Operations + 2024 + SC + doi:10.1109/SC41406.2024.00094】
* 【36, Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms + 2021 + ISCA + doi:10.1109/ISCA52012.2021.00049】
* 【37, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism + 2019 + arXiv + http://arxiv.org/abs/1909.08053 】
* 【41, Llama 2: Open Foundation and Fine-Tuned Chat Models + 2023 + arXiv + https://arxiv.org/abs/2307.09288 】
* 【43, Expanding our AI-optimized infrastructure portfolio: Introducing Cloud TPU v5e and announcing A3 GA + 2023 + Google Cloud Blog + https://cloud.google.com/blog/products/compute/announcing-cloud-tpu-v5e-and-a3-gpus-in-ga 】
* 【44, Attention is All you Need + 2017 + NeurIPS + https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html 】
* 【45, Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models + 2023 + ASPLOS + doi:10.1145/3567955.3567959】