最近，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 为例：