LLM 推理的 Attention 计算和 KV Cache 优化：PagedAttention、vAttention 等

一、背景

最近，SGLang 引起了广泛关注，出现了许多 “SGLang 吊打 vLLM 和 TRT-LLM” 的言论。不得不说，SGLang 确实是一项非常出色的工作。与此同时，vLLM 的性能问题和 TRT-LLM 的易用性问题也广受诟病，但是在实际应用中，我们仍然需要保持理性。比如，已经使用了 LMDeploy 或 TRT-LLM，是否要在当前阶段切换到 SGLang；SGLang 在对应的场景是否一定有这么大的提升？

不过，本文中并非要介绍 SGLang，而是旨在探讨 vLLM 的基石——PagedAttention 的一些问题，以及现有的一些改进工作，比如 vAttention 和 vTensor 等。此外，我们还会简单介绍一些与 Attention 密切相关的 MHA/GQA 的实现考量。

需要注意的是，网上已经有许多关于 vAttention 和 vTensor 的论文解读，本文不会详细介绍这些论文，甚至可能忽略一些实现的细节，这里只是通过这些论文引出一些需要关注的细节。

二、引言

2.1 MHA Attention 计算

如下图所示为标准的 LLM Decoding 阶段的 Multi-Head Attention（MHA）计算，其中的 D 表示 hidden size，H 表示 Head 个数，L 表示当前是在序列的第 L 个 Token。可以看出：

• 当Batch Size 为 1时，图中红色、绿色、蓝色处的矩阵乘法全部为矩阵乘向量，是明显的 Memory Bound，算术强度不到 1。
• 当Batch Size 大于 1时（比如 Continuous Batching）：

红色和蓝色部分：因为是 Weight 乘以 Activation，所以不同的 Request 之间可以共享 Weight。这里变成矩阵乘矩阵，并且 Batch Size 越大，算术强度越大，也就越趋近于 Compute Bound（FFN 层也类似）。

绿色部分：这里 Q、K 和 V 的 Attention 计算，是 Activation 乘以 Activation，所以不同的 Request 之间没有任何相关性。即使 Batching，这里也是Batched 矩阵乘向量，并且因为序列长度可能不同，这里不同 Request 的矩阵乘向量是不规则的。也就是说，这里算术强度始终不到 1，是明显的 Memory Bound。

从上可以看出，通过 Continuous Batching 可以很好的将 Memory Bound 问题转变为 Compute Bound，但 Q、K 和 V 的 Attention 计算的算术强度却始终小于 1。根据 Amdahl 法则，如果系统中有一部分无法优化，即使把其他部分优化到可以忽略，不可优化的部分也会决定整个系统的性能上限。不幸的是，Sequence Length 越长，这里的计算量就越不可忽略。

根据模型配置信息可以估算出模型中 Q、K 和 V 的 Attention 计算与其他矩阵计算的比例大约为 (L+D)/(12*D)（PS：准确值需要根据具体的模型参数计算）。也就是说，当序列长度 L 等于 12 倍的 hidden size 时，两部分的计算量相当，即使其他矩阵计算优化到 0，加速比也只有 2x。比如 LLaMA 2 7B 的 hidden size 为 4K，当序列长度达到 44K 时，两部分的计算量相当，要优化的重点也会很不一样，这也是很多长序列相关工作会在 Attention 部分采用稀疏 Attention 的一个重要原因。

2.2 GQA Attention 计算

早期通常只有比较大的模型才会采用 GQA（[2305.13245] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints），比如 LLaMA -2 70B，而 LLaMA-2 7B/13B 都没有采用 GQA。然而，LLaMA-3 8B 中也用上了 GQA，甚至其他更小的模型也在将 MHA 替换为 GQA。

• 使用 GQA 有个非常大的好处：在推理阶段可以显著降低 KV Cache 的大小，比如,相比 32 个 KV Head 的 MHA，32 个 Query Head，8 个 KV Head 的 GQA 的 KV Cache 大小可以降低到 MHA 的 8/32=1/4，这也为更大的 Batch Size 提供了空间，可以进一步提升吞吐。
• 除此之外，还有一个比较大的好处：可以明显提升 Q、K 和 V 的 Attention 计算的算术强度。此时虽然不同的 Request 之间同样不能共享，但是同一个 Request 中的不同 Head 可以共享，比如 4 个 Query Head 共享 1 个 KV Head，则算术强度就会接近于 4，也可以更充分发挥 Tensor Core 的算力。

使用 MHA 时，Q、K 和 V 的 Attention 计算可以使用 CUDA Core 也可以使用 Tensor Core。由于 Tensor Core 要求矩阵的 Shape 是 8 的整数倍，如果不满足就只能 Padding：

• 对于MHA而言，其是矩阵乘向量，则有7/8 的计算是冗余的。
• 对于GQA而言，如果 4 个 Query Head 共享 1 个 KV Head，则 Attention 计算有 4/8 的计算是冗余的，如果8 个 Query Head 共享 1 个 KV Head，则没有计算的冗余。很多框架已经做了相关优化，比如 LMDeploy，TRT-LLM 的 XQA 等。
• 此外，PagedAttention 的 KV Cache 是非连续存储的，导致即使使用 GQA 也无法利用 Tensor Core。

PS：对于 GQA 而言，理论上也可以期望 GPU 的 L2 Cache 能够缓存到共享的 Key 和 Value Cache，从而缓解 IO Bound 问题，然而实际上无法人为控制，不一定能达到理想的效果。

2.3 NVIDIA 驱动

一般所说的 NVIDIA Driver 包含以下两个部分，NVIDIA 从 2022 年 5 月正式开源的驱动程序也只是下述的 GPU Kernel-Mode Driver：NVIDIA Linux open GPU kernel module source，没有开源 CUDA User-Mode Driver。可以使用 modinfo nvidia 或 cat /proc/driver/nvidia/version 查看当前系统安装的 GPU Kernel-Mode Driver 版本：

• GPUKernel-Mode Driver，部分如下所示：

• nvidia.ko：提供对 NVIDIA 硬件的 Low-level 访问，供其他组件使用。
• nvidia-uvm.ko：为 CUDA Driver 提供统一虚拟内存（Unified Virtual Memory，UVM）功能。
• nvidia-drm.ko：向 Linux 内核的 DRM 子系统注册 DRM 驱动程序。
• nvidia-peermem.ko：提供 Mellanox InfiniBand 基于 HCA 的直接 P2P 读写 NVIDIA GPU 显存的功能。

• CUDAUser-Mode Driver，部分如下所示：
• libcuda.so：为 CUDA 应用程序提供运行时支持。
• libnvidia-ml.so：提供监控和管理 API。
• libnvidia-nvvm.so：由 CUDA 驱动程序加载，用于进行 JIT 链接时优化。
• libnvidia-ptxjitcompiler.so：将 PTX 编译为 GPU 机器码，并由 CUDA 驱动程序使用。

2.4 Low-level VMM API

CUDA 的 Low-level VMM API 从 CUDA 10.2 版本正式引入，用于更高效地管理GPU虚拟内存。其提供如下一些主要的 API（部分），使得开发者能够构建更高效的动态数据结构，更好地控制应用程序中的 GPU 内存使用：

• cuMemCreate：创建物理内存句柄
• cuMemAddressReserve：保留虚拟地址范围
• cuMemMap：将物理内存句柄映射到虚拟地址范围
• cuMemSetAccess：设置分配的内存访问权限

这里不再赘述，具体可以参考：Introducing Low-Level GPU Virtual Memory Management | NVIDIA Technical Blog。

需要注意的是：NVIDIA 官方提供的 Low-level API 能分配的最小 Physical Chunk 为 2MB。

三、PagedAttention

3.1 摘要

PagedAttention 是一种受操作系统中虚拟内存和分页技术启发的注意力算法。在此基础上，作者构建了 vLLM 推理框架，可实现：

• 近似实现 KV 的零浪费。
• 在请求内部和请求之间灵活共享 KV Cache，以进一步减少内存使用。

对应的 Paper 为：[2309.06180] Efficient Memory Management for Large Language Model Serving with PagedAttention

对应的 vLLM 的代码库：GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs

3.2 问题

传统的 LLM 系统通常会为每个 Request 预先分配显存，如下图 Figure 3 所示，假设模型支持的最大训练长度为 2048，则最多会为每个 Request 预先分配 2048 个 Token 的显存空间。这样如果实际执行中的输入和输出包含 10 个 Token，则将有 2038 个 Token 空间的浪费（内存碎片），随着服务支持的 Batch Size 增加，这一问题会愈加明显：

3.3 方案

3.3.1 内存管理

为了解决上述问题，作者参考操作系统中内存碎片和 Sharing 的解决方案：带分页的虚拟内存，并提出 PagedAttention。如下图 Figure 6 所示，PagedAttention 将 Request 的 KV Cache 划分为多个 Block，每个 Block 可以包含固定数量的 Token 的 Key 和 Value，在 Logical KV Blocks 中 Block 是连续的，但是在 Physical KV Blocks 中 Block 是不连续的（对应预先分配的大块物理内存）。因此，可以像操作系统的虚拟内存一样，以更灵活的方式管理 KV Cache，当然，也就需要 Block Table 来管理这种映射关系。这样的好处是可以将整个 Request 对应的 KV Cache 划分为相同大小的 Block，Block 可以远小于模型支持的序列长度，以此来明显缓解内存碎片化。

3.3.2 Prefix Caching

除了减少内存碎片外，这种方式的另外一个好处是可以更好的进行 Prefix Caching。如下图所示，如果一个 Request 要进行并行采样，或者两个 Request 有公共的前缀（比如 System Prompt），那么实际上只用计算、存储一份 KV Cache 即可，即可减少计算，也可减少存储。当然，一个 Block 中只会来自一个 Request（Sample），因此都是按照整个 Block 的粒度共享，如果 Block Size 过大，则可能影响可以共享的长度。比如，如果 Block Size 为 100，两个序列在第 99 个 Token 时出现不同，则无法实现共享。

3.3.3 Attention 计算

在实际的 Attention Kernel 计算时，会按照 Block 的粒度执行，以 Query Token qi 为例（对应 “forth”），第一次会计算 qi 对应的向量与 Block 0 中（对应 “Four score and seven”）Kj 的乘积，来计算 Attention Score Aij，第二次计算 Block 1，以此类推：

也就是说：PagedAttention 允许不同的 Block 在物理内存中以不连续的方式存储，可以充分增加内存管理的灵活性。当然，这也就导致实现的 CUDA Kernel 是和内存管理耦合的，这也就增加了 Kernel 的计算开销和实现的复杂度。

3.4 消融实验

3.4.1 Block Mapping 对性能的影响

PagedAttention 中的动态 Block Mapping 会影响涉及存储 KV Cache 的 GPU 操作的性能。与传统连续存储的方式相比，PagedAttention 中的 GPU Kernel 会引入访问 Block Table、执行额外的分支，以及处理可变序列长度的额外开销。如下图所示，作者与高度优化的 FasterTransformer 实现相比，Attention Kernel 的延迟会高出 20-26%。因为只会影响 Attention Kernel，不会影响其他算子，比如 Linear Kernel，作者认为其还可以接受（PS：这里用的 Context Length 非常短，随着序列变长，差距可能会越来越大）：

3.4.2 Block Size 对性能的影响

Block 的大小可能会对 PagedAttention 的性能产生重大影响：

• 如果Block 太小，PagedAttention 可能无法充分利用 GPU 的并行性来读取和处理 KV Cache。
• 如果Block 过大，则内存碎片会增加，Prefix Cache 共享的可能性会降低。

如下图所示，作者使用 ShareGPT 和 Alpaca 数据评估了不同 Block Size 下的 Latency（越低越好），可以看出，当 Block Size 为 16 和 32 时表现最好，因此作者在 vLLM 中将默认的 Block Size 设置为 16：

3.5 问题

这里的 Block Size 应该指的是 Token 数？那么不同的配置是否会有不同的最优值，比如 LLaMA3 8B、70B、405B 的最优配置是否相同？除此之外，使用了 GQA 的模型是否又会有不同的 Block Size？

四、GMLake

4.1 摘要

GMLake 是蚂蚁和上海交通大学的工作，作者指出，GPU 原生的内存分配器开销很大，常见的 DNN 框架（比如 Pytorch 和 Tensorflow）会采用 Caching 机制，通过使用 GPU 原生的分配器分配大块内存，然后通过分割机制来实现快速分配和释放。然而，Caching 机制会因为重计算、Offload、分布式训练以及低秩微调等方式而引入频繁和不规则的内存分配和释放，导致存在严重的内存碎片问题。

为了解决上述问题，作者提出了一种基于 Low-level 的 GPU 虚拟内存管理的内存分配框架，称为 GMLake（GPU Memory Lake）。使用 GMLake，可以融合或连接非连续的内存块，并通过虚拟内存地址进行映射。在 A100 80GB GPU 上，通过这种方式可以显著减少 GPU 内存使用量（平均减少 9.2 GB，最多 25 GB）和内存碎片（平均减少 15%，最多 33%）。

对应的论文为：[2401.08156] GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching

4.2 背景

如下图 Figure 2 所示为 3 种分配机制的区别：

• 原生的 GPU 内存分配：每次都调用 cudaMalloc 分配，调用 cudaFree 释放。
• Caching 分配：Pytorch 和 Tensorflow 采用的 BFC 算法，预先分配大块内存，然后通过查找、拆分、合并等方式实现最大化利用，减少 cudaMalloc 和 cudaFree 的巨大开销。作者实验表明，Caching 分配的吞吐大概是原生 GPU 内存分配的 10 倍。
• Virtual Memory：利用 GPU 的 Low-level 虚拟内存管理方案分配。

4.3 方案

4.3.1 概览

如下图 Figure 7 所示为 GMLake 的方案概览，它提供了与现有的 Caching Allocator 接口相同的 GMLake Allocator，但是在内部集成了虚拟内存拼接机制（Virtual Memory Stiching，VMS），其主要是依赖 CUDA 的 Low-level VM 管理 API实现。GMLake 主要是包含 3 个组件：

• Virtual memory API：用于指示 GPU 使用虚拟内存地址分配和释放内存的 Low-level API。
• Virtual memory pool：作为基础数据结构，用于缓存虚拟内存，提高效率。
• GMLake allocator：包括管理 VM 池所必须的所有函数、算法和策略。

4.3.2 Stitched Memory Pool & Primitive Memory Pool

原始的 Low-level VMM API 也非常耗时，因此减少其使用量对于实现高效率 GMLake 至关重要。作者从 Caching Allocator 中汲取灵感，设计了具有 Caching 功能的虚拟内存池（VMP），从而显著减少物理内存分配的次数。如下图 Figure 8 所示，作者设计了两种内存池：

• Primitive Memory Pool（pPool）：pPool 使用有序集合来存储 pBlock，对于每个 pBlock，pPool 首先构造一个结构来记录指向 pBlock 的指针，其包含基础属性，比如 pBlock 的激活状态。随后，将新分配的 pBlock 插入到集合中，所有 pBlock 按照块大小降序排列。pBlock 作为原始块，代表 High-level Tensor 可访问的最小单元（Low-level 能分配的最小 Physical Chunk 为 2MB，一个 pBlock 可以包含多个 Physical Block），它作为基本数据结构，可以被多个 sBlock 拼接或指向。
• Stitched Memory Pool（sPool）：sPool 也被组织成一个有序集合，类似于 pPool。它的元素包含拼接的 block 结构，该结构组合了多个 pBlock。比如下图中的 sBlock 3 包含一个组合在一起的 pBlock 3 和 pBlock 5。同时，pBlock 3 也可以指向 sBlock 2。

4.4 问题

LLM 推理与训练不同，训练中通常会将输入序列拼接到模型支持的最大序列长度，比如 4096 Token，可以充分提高效率，并且是等价的。这种方式对于内存分配相对比较友好，并且通常是一些大块的内存分配。而在 LLM 推理场景，输入、输出的序列长度可能差异很大，尤其是 Decoding 阶段是一个一个 Token 生成，导致分配的最小粒度变为单个 Token，就会涉及很多小块内存分配，也就需要更精细化的内存管理。

五、vTensor

5.1 摘要

vTensor 同样由蚂蚁和上海交大发表，和 GMLake 大部分作者相同，可以认为是将 GMLake 由 DNN Training 扩展到 LLM Inference 场景。其比 vAttention 晚发表几个月，和 vAttention 的思路非常类似，不过两个工作都是最近才开源的。

很自然，vTensor 也是一种基于 GPU 虚拟内存管理（VMM）的 LLM 推理 Tensor 结构。vTensor 通过将计算与内存碎片整理解耦并提供动态可扩展性来解决现有的限制（主要指 PagedAttention）。其采用 CPU-GPU 异构方案，确保高效、无碎片的内存管理（PS：近似？），同时可以适应不同 LLM 架构的各种计算 Kernel。

实验表明，vTensor 在不同的模型中实现了平均 1.86x 加速，在多轮聊天场景中最高可达 2.42x。此外，vTensor 在 Kernel 评估中，可以实现平均 2.12x 和 3.15x 加速，与 SGLang Triton prefix-prefilling Kernel 和 vLLM 的 PagedAttention Kernel 相比，在 A100 上可以释放 71.25%（57GB）内存，从而支持更多内存密集型工作负载。

对应的论文为：[2407.15309] vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving

对应的代码库为：​​https://github.com/intelligent-machine-learning/glake/tree/master/GLakeServe​​

5.2 方案

5.2.1 vTensor

如下图 Figure 1 所示为本文提出的 vTensor 与原始的 KV Cache 机制以及 Paged KV Cache 的区别，和 vAttention 类似，也是基于 Low-level 的 vMM API，有 3 个好处：

• 可以实现内存管理与 CUDA Kernel 的解构，更加通用，Kernel 实现更加简单。
• 可以实现内存的动态扩展，减少浪费。
• Attention Kernel 可以使用更强算力的Tensor Core（虚拟地址连续），而 PagedAttention 只能使用 CUDA Core。

​

如下图 Figure 5 所示为具体的 vTensor 的实现，vTensor 指针由 vTensor Manager（VTM）生成。

• 当向vTensor Scheduler（VTS）发送创建请求时，VTS 将创建分配策略，然后让 vTensor Operation（VTO）分配虚拟内存和相应的 Physical Chunk（PC）。
• vTensor 指针 *A 指向 GPUVirtual Address（VA），该地址必须是连续的，才能与标准 CUDA 分配的 Tensor 兼容。
• Virtual Memory Management（VMM）维护 VTM 注册的 Physical Chunk 的完整映射信息，VMM 允许 GPU Tensor Core 通过 Virtual Address 访问 GPU 内存中所需的数据。
• Physical Chunk在 GPU 内存中分配，但是Physical Chunk Handle（PH）和 Virtual Address 可以由 CPU 访问和管理。Physical Handle 和 Virtual Memory 仅具有 Physical Chunk 的索引信息，而不包含 Device-Host 传输。
• Physical Chunk同样采用了 2MB 的 Physical Chunk，其对应的 Handle 只有几个字节，由 vTensor Pool（VTP）记录并存储在 CPU 内存中。
• 作者也开发了一系列基于 vTensor 的方法来操作 KV Cache 的碎片整理，这是 vTensor 和 FlexInfer 设计的基础。

5.2.2 FlexInfer

基于 vTensor，作者进一步开发了 FlexInfer，它是一个 CPU 和 GPU 异构框架，将大多数内存操作解耦并卸载到 CPU，并通过重叠 GPU 计算来隐藏它们。与之前在 GPU 上内存操作（比如 PagedAttention）相比，CPU 更擅长与内存相关的操作。当请求发送到 FlexInfer 时，它会根据请求的配置（例如 Batch Size 和 Seq Length）解耦内存和计算操作。

• 在计算方面，FlexInfer Scheduler 采用原始的高度优化的 Kernel 在 GPU Tensor Core 上运行 LLM，在不受算术强度限制的情况下保持计算灵活性和效率。
• 在内存方面，FlexInfer 还提供了高度定制的调度方案，以在启动、Prefill、Decoding 以及终止阶段与计算重叠，隐藏内存分配和释放，可以显著降低 LLM Service System 的内存操作开销。

5.3 消融实验

5.3.1 Decoding Kernel 评估

如下图 Figure 7 所示，作者对比了不同场景下相应 Attention Kernel 的性能，其中 FlexInfer attn 表示使用 vTensor 的 FlashAttention，Paged flash attn 表示使用 Paged 的 FlashAttention，Flash attn 表示原始的 FlashAttention：

• Batch Size：随着 Batch Size 增加，FlexInfer attn 和 Flash attn 始终保持最低 Latency，FlexInfer attn Latency 平均比 Paged attn 低 42%。
• Sequence Length：以 Batch Size 16 为例，随着 Sequence Length 增加，FlexInfer attn 和 Flash attn 依然保持最低 Latency，在 1K 时加速最明显（但是在 1K 序列长度时，Attention 在整个计算中的占比也比较小）。
• KV Head：以 Batch Size 16 和 Sequence Length 16K 为例，其中 KV Head 为 1 对应 MQA，KV Head 32 对应 MHA（Yi-6B、Yi-9B），KV Head 4 和 8 对应 GQA。可以看出

MQA（KV Head 1）时，加速最明显，可以最充分发挥 Tensor Core 的算力。

MHA（KV Head 32）时，基本没有加速，此时都无法充分发挥 Tensor Core 算力。

GQA（KV Head 4 和 8）时，有一定加速，可以部分发挥 Tensor Core 算力。

5.3.2 Prefix-prefilling Kernel 评估

如下图 Figure 8 所示，作者进一步对比了 Prefilling 阶段 Kernel 的性能（对应的 Sequence Length 固定为 16K），可以看出，在不同 Batch Size 和 Prefix/Prompt 比例下，PagedAttention 的性能都是最差的，其他几种方式性能相当。当然，将 Paged KV Cache 适配到 FlashAttention 的代价也很高，而 FlexInfer attn 的实现代价小得多。

5.4 问题

六、vAttention

6.1 摘要

vAttention 是微软的工作，其发表在 GMLake 和 vTensor 之间，思路和 vTensor 非常接近。同样是使用 Low-level 的 VMM API 实现，也同样是为了减少 KV Cache 内存碎片，降低 Attention Kernel 的开发成本。与 vTensor 不同的是，作者还进一步修改了 GPU 内核 Driver，以提供更细粒度的 Physical Chunk 的分配，比如 64KB、128KB 和 256KB，而不局限于 2MB。

结果表明，vAttention 可以为各种 Attention Kernel 实现无缝的兼容，并提供动态内存管理能力。vAttention 生成 Token 的速度比 vLLM 快 1.97x，同时处理 Prompt 的速度比 FlashAttention 和 FlashInfer 的 PagedAttention 快 3.92x 和 1.45x。

对应的论文为：[2405.04437] vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

6.2 方法

6.2.1 Low-level VMM API 适配

原生的 CUDA Low-Level VMM API 最小只能分配 2MB 的 Physical Chunk，作者认为其依然会导致存在一部分的内存碎片（PS：后文会介绍），因此决定修改 CUDA Low-level API，以允许分配更细粒度的 Physical Chunk，比如 64KB、128KB 和 256KB。并提供了一些新的 API：

如上只是一部分代码修改，实际的修改有 200-300 行，具体可以参考 ​​https://github.com/microsoft/vattention/tree/main/nvidia-vattn-uvm-driver​​。需要指出的是，这个修改是针对 545.23.06 版本，其他版本需要进一步适配；此外，因为涉及了底层的 nvidia-uvm Driver 的修改，因此需要替换系统已有 Driver 并且重新启动 Server，相应的代价也比较高。

如下图 Table 3 所示，作者也进一步对相关 API 进行了封装，并测试了不同分配大小的时延：

6.2.2 vAttention

具体的思路和 vTensor 类似，不再赘述，这里阶段介绍一个 vAttention 中动态内存管理的示例：

• a：两个请求 R1 和 R2 的 virtural tensor，没有使用 Physical Memory。
• b：R1 分配了一个 Physical Page。
• c：R1 分配了两个 Physical Page，R2 分配了一个 Physical Page。
• d：R1 处理完，但是不会立即释放对应的 Physical Page；R2 分配了两个 Physical Page。
• e：新的请求 R3 分配了两个 Physical Page，会利用之前 R2 分配的 Physical Page。

其实 vLLM 也已经在 PR（​​https://github.com/vllm-project/vllm/pull/6102​​）中提供了对 vAttention 的支持，但是目前还没有合入，也没有实现所有功能。比如，当前还只支持 2MB 的 Physical Chunk，可能会存在比较大的显存碎片（下面会介绍）。

6.3 消融实验

6.3.1 Physical Chunk 大小对 Prefill 的影响

在 Prefill 阶段的 Token 数比较多，每层的 K 或 V Cache 可能远超 2M，此时使用过小的 Chunk 有可能会引入比较多的开销。如下图 Figure 11 所示，使用 64KB Chunk 最多会导致 Prefill Latency 增加 15%，不过通过计算和分配的 Overlap 可以隐藏掉这一部分开销（对应 vAttention）。

6.3.2 Physical Chunk 大小对分配带宽的影响

如下图 Table 7 所示，使用更小的 Physical Chunk Size 会导致每秒分配的显存大小降低，当 Physical Chunk Size 为 64KB 时，每秒分配的显存量只有 2MB 时的 1/5 左右：

如果 64KB 时的分配速度无法满足实际需要的分配速度，那么就可能成为瓶颈。作者也进行了相应的测试，如下图所示，当 Batch Size 增加到 320 时（绿线 LLama3-8B - 2 A100，实线 Yi-6B - 1A100，蓝线 Yi-34B - 2 A100），需要的最大分配带宽也只有 600MB/s，远小于 64KB 对应的 7.59GB/s，也证明 64KB 的 Physical Chunk 完全可以接受：

6.3.3 Memory 碎片

在进行 Attention 计算时，各个 Request 之间是没有任何关系的，即使是 Continuous Batching，也是 N 个矩阵乘向量（MHA）或者 N 个矩阵乘矩阵（GQA、MQA）。因此，不管是 PagedAttention 还是 vTensor 或者 vAttention，其每个 Chunk 中都只会存储同一个 Request 的 Token，此时，在每个 Request 的最后一个 Chunk 中就可能存来空闲未被使用的空间，Chunk 的 Size 越大，理论浪费的空间就越大。

如下图 Table 8 所示，作者以 Yi-6B、LLaMA-3-8B 和 Yi-34B 为例，统计了不同 Chunk 可以容纳的 Token 数，以及浪费的最大 Memory 空间：

• 模型结构：

Yi-6B 的 Hidden Size 为 4096，总共 32 个 Q-Head，4 个 KV-Head，32 个 Layer。

Yi-34B 的 Hidden Size 为 7168，总共 56 个 Q-Head，8 个 KV-Head， 60 个 Layer。

LLaMA-3-8B 的 Hidden Size 为 4096，总共 32 个 Q-Head，8 个 KV-Head， 32 个 Layer。

• 每一层一个 Token 的 Key 或 Value Cache 的大小为（FP16 存储，如果采用 TP，则一般会按照 Head 切分，所以每个 GPU 上 1 个 Token 相应的存储占用的空间减少）：
• Yi-6B：4096(dim)/32(head)*4(head)*2(Byte)=1KB
• Yi-34B：7168(dim)/56(head)*8(head)*2(Byte)=2KB
• LLaMA-3-8B：4086(dim)/32(head)*8(head)*2(Byte)=2KB
• 每一层一个 Chunk 可以存储的 Key 或 Value Cache 的 Token 数为（TP1）：
• Yi-6B 64KB：64KB/1KB=64
• Yi-34B 256KB：256KB/2KB=128
• 每一层最多浪费 Key 和 Value 两个 Chunk，则理论上一个 Request 浪费的最大空间为：
• Yi-6B 2MB：2MB*2(K/V)*32(Layer)=128MB
• LLaMA-3-8B 128KB：128KB*2*32(Layer)=8MB

从以上的统计数据可以看出，每个请求最大的 Memory 浪费与 Chunk Size 成正比，以 Batch Size 100，TP=1 为例：

• 2MB 的 Chunk Size 最大浪费 12.5GB-23.4GB（128MB-240MB * 100）。
• 64KB 的 Chunk Size 最大浪费 400MB-750MB（4MB-7.5MB * 100）。

七、参考链接

1. ​​https://arxiv.org/abs/2305.13245​​
2. ​​https://github.com/NVIDIA/open-gpu-kernel-modules​​
3. ​​https://developer.nvidia.com/blog/introducing-low-level-gpu-virtual-memory-management/​​
4. ​​https://arxiv.org/abs/2309.06180​​
5. ​​https://github.com/vllm-project/vllm​​
6. ​​https://arxiv.org/abs/2401.08156​​
7. ​​https://arxiv.org/abs/2407.15309​​
8. ​​https://github.com/intelligent-machine-learning/glake/tree/master/GLakeServe​​
9. ​​https://arxiv.org/abs/2405.04437​​
10. ​​https://github.com/microsoft/vattention/tree/main/nvidia-vattn-uvm-driver​​
11. ​​https://github.com/vllm-project/vllm/pull/6102​​

本文转载自 ​​AI闲谈​​，作者： AI闲谈