一、背景

最近，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閑談