最近，SGLang 引起了广泛关注，出现了许多 “SGLang 吊打 vLLM 和 TRT-LLM” 的言论。不得不说，SGLang 确实是一项非常出色的工作。与此同时，vLLM 的性能问题和 TRT-LLM 的易用性问题也广受诟病，但是在实际应用中，我们仍然需要保持理性。比如，已经使用了 LMDeploy 或 TRT-LLM，是否要在当前阶段切换到 SGLang；SGLang 在对应的场景是否一定有这么大的提升？
不过，本文中并非要介绍 SGLang，而是旨在探讨 vLLM 的基石——PagedAttention 的一些问题，以及现有的一些改进工作，比如 vAttention 和 vTensor 等。此外，我们还会简单介绍一些与 Attention 密切相关的 MHA/GQA 的实现考量。
需要注意的是，网上已经有许多关于 vAttention 和 vTensor 的论文解读，本文不会详细介绍这些论文，甚至可能忽略一些实现的细节，这里只是通过这些论文引出一些需要关注的细节。
如下图所示为标准的 LLM Decoding 阶段的 Multi-Head Attention（MHA）计算，其中的 D 表示 hidden size，H 表示 Head 个数，L 表示当前是在序列的第 L 个 Token。可以看出：
红色和蓝色部分：因为是 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 的一个重要原因。
早期通常只有比较大的模型才会采用 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。
使用 MHA 时，Q、K 和 V 的 Attention 计算可以使用 CUDA Core 也可以使用 Tensor Core。由于 Tensor Core 要求矩阵的 Shape 是 8 的整数倍，如果不满足就只能 Padding：
PS：对于 GQA 而言，理论上也可以期望 GPU 的 L2 Cache 能够缓存到共享的 Key 和 Value Cache，从而缓解 IO Bound 问题，然而实际上无法人为控制，不一定能达到理想的效果。
一般所说的 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 版本：
CUDA 的 Low-level VMM API 从 CUDA 10.2 版本正式引入，用于更高效地管理GPU虚拟内存。其提供如下一些主要的 API（部分），使得开发者能够构建更高效的动态数据结构，更好地控制应用程序中的 GPU 内存使用：
这里不再赘述，具体可以参考：Introducing Low-Level GPU Virtual Memory Management | NVIDIA Technical Blog。
需要注意的是：NVIDIA 官方提供的 Low-level API 能分配的最小 Physical Chunk 为 2MB。
PagedAttention 是一种受操作系统中虚拟内存和分页技术启发的注意力算法。在此基础上，作者构建了 vLLM 推理框架，可实现：
传统的 LLM 系统通常会为每个 Request 预先分配显存，如下图 Figure 3 所示，假设模型支持的最大训练长度为 2048，则最多会为每个 Request 预先分配 2048 个 Token 的显存空间。这样如果实际执行中的输入和输出包含 10 个 Token，则将有 2038 个 Token 空间的浪费（内存碎片），随着服务支持的 Batch Size 增加，这一问题会愈加明显：
为了解决上述问题，作者参考操作系统中内存碎片和 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 可以远小于模型支持的序列长度，以此来明显缓解内存碎片化。
除了减少内存碎片外，这种方式的另外一个好处是可以更好的进行 Prefix Caching。如下图所示，如果一个 Request 要进行并行采样，或者两个 Request 有公共的前缀（比如 System Prompt），那么实际上只用计算、存储一份 KV Cache 即可，即可减少计算，也可减少存储。当然，一个 Block 中只会来自一个 Request（Sample），因此都是按照整个 Block 的粒度共享，如果 Block Size 过大，则可能影响可以共享的长度。比如，如果 Block Size 为 100，两个序列在第 99 个 Token 时出现不同，则无法实现共享。
在实际的 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 的计算开销和实现的复杂度。
PagedAttention 中的动态 Block Mapping 会影响涉及存储 KV Cache 的 GPU 操作的性能。与传统连续存储的方式相比，PagedAttention 中的 GPU Kernel 会引入访问 Block Table、执行额外的分支，以及处理可变序列长度的额外开销。如下图所示，作者与高度优化的 FasterTransformer 实现相比，Attention Kernel 的延迟会高出 20-26%。因为只会影响 Attention Kernel，不会影响其他算子，比如 Linear Kernel，作者认为其还可以接受（PS：这里用的 Context Length 非常短，随着序列变长，差距可能会越来越大）：
Block 的大小可能会对 PagedAttention 的性能产生重大影响：
如下图所示，作者使用 ShareGPT 和 Alpaca 数据评估了不同 Block Size 下的 Latency（越低越好），可以看出，当 Block Size 为 16 和 32 时表现最好，因此作者在 vLLM 中将默认的 Block Size 设置为 16：
这里的 Block Size 应该指的是 Token 数？那么不同的配置是否会有不同的最优值，比如 LLaMA3 8B、70B、405B 的最优配置是否相同？除此之外，使用了 GQA 的模型是否又会有不同的 Block Size？
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%）。
如下图 Figure 2 所示为 3 种分配机制的区别：
如下图 Figure 7 所示为 GMLake 的方案概览，它提供了与现有的 Caching Allocator 接口相同的 GMLake Allocator，但是在内部集成了虚拟内存拼接机制（Virtual Memory Stiching，VMS），其主要是依赖 CUDA 的 Low-level VM 管理 API实现。GMLake 主要是包含 3 个组件：
原始的 Low-level VMM API 也非常耗时，因此减少其使用量对于实现高效率 GMLake 至关重要。作者从 Caching Allocator 中汲取灵感，设计了具有 Caching 功能的虚拟内存池（VMP），从而显著减少物理内存分配的次数。如下图 Figure 8 所示，作者设计了两种内存池：
LLM 推理与训练不同，训练中通常会将输入序列拼接到模型支持的最大序列长度，比如 4096 Token，可以充分提高效率，并且是等价的。这种方式对于内存分配相对比较友好，并且通常是一些大块的内存分配。而在 LLM 推理场景，输入、输出的序列长度可能差异很大，尤其是 Decoding 阶段是一个一个 Token 生成，导致分配的最小粒度变为单个 Token，就会涉及很多小块内存分配，也就需要更精细化的内存管理。
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）内存，从而支持更多内存密集型工作负载。
如下图 Figure 1 所示为本文提出的 vTensor 与原始的 KV Cache 机制以及 Paged KV Cache 的区别，和 vAttention 类似，也是基于 Low-level 的 vMM API，有 3 个好处：
如下图 Figure 5 所示为具体的 vTensor 的实现，vTensor 指针由 vTensor Manager（VTM）生成。
基于 vTensor，作者进一步开发了 FlexInfer，它是一个 CPU 和 GPU 异构框架，将大多数内存操作解耦并卸载到 CPU，并通过重叠 GPU 计算来隐藏它们。与之前在 GPU 上内存操作（比如 PagedAttention）相比，CPU 更擅长与内存相关的操作。当请求发送到 FlexInfer 时，它会根据请求的配置（例如 Batch Size 和 Seq Length）解耦内存和计算操作。
如下图 Figure 7 所示，作者对比了不同场景下相应 Attention Kernel 的性能，其中 FlexInfer attn 表示使用 vTensor 的 FlashAttention，Paged flash attn 表示使用 Paged 的 FlashAttention，Flash attn 表示原始的 FlashAttention：
MQA（KV Head 1）时，加速最明显，可以最充分发挥 Tensor Core 的算力。
MHA（KV Head 32）时，基本没有加速，此时都无法充分发挥 Tensor Core 算力。
GQA（KV Head 4 和 8）时，有一定加速，可以部分发挥 Tensor Core 算力。
如下图 Figure 8 所示，作者进一步对比了 Prefilling 阶段 Kernel 的性能（对应的 Sequence Length 固定为 16K），可以看出，在不同 Batch Size 和 Prefix/Prompt 比例下，PagedAttention 的性能都是最差的，其他几种方式性能相当。当然，将 Paged KV Cache 适配到 FlashAttention 的代价也很高，而 FlexInfer attn 的实现代价小得多。
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。
原生的 CUDA Low-Level VMM API 最小只能分配 2MB 的 Physical Chunk，作者认为其依然会导致存在一部分的内存碎片（PS：后文会介绍），因此决定修改 CUDA Low-level API，以允许分配更细粒度的 Physical Chunk，比如 64KB、128KB 和 256KB。并提供了一些新的 API：
如上只是一部分代码修改，实际的修改有 200-300 行，具体可以参考 。需要指出的是，这个修改是针对 545.23.06 版本，其他版本需要进一步适配；此外，因为涉及了底层的 nvidia-uvm Driver 的修改，因此需要替换系统已有 Driver 并且重新启动 Server，相应的代价也比较高。
如下图 Table 3 所示，作者也进一步对相关 API 进行了封装，并测试了不同分配大小的时延：
具体的思路和 vTensor 类似，不再赘述，这里阶段介绍一个 vAttention 中动态内存管理的示例：
）中提供了对 vAttention 的支持，但是目前还没有合入，也没有实现所有功能。比如，当前还只支持 2MB 的 Physical Chunk，可能会存在比较大的显存碎片（下面会介绍）。
在 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 完全可以接受：
在进行 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 空间：
从以上的统计数据可以看出，每个请求最大的 Memory 浪费与 Chunk Size 成正比，以 Batch Size 100，TP=1 为例：