一、背景

最近，SGLang 引起了廣泛關(guān)注，出現(xiàn)了許多 “SGLang 吊打 vLLM 和 TRT-LLM” 的言論。不得不說，SGLang 確實(shí)是一項(xiàng)非常出色的工作。與此同時(shí)，vLLM 的性能問題和 TRT-LLM 的易用性問題也廣受詬病，但是在實(shí)際應(yīng)用中，我們?nèi)匀恍枰３掷硇?。比如，已?jīng)使用了 LMDeploy 或 TRT-LLM，是否要在當(dāng)前階段切換到 SGLang；SGLang 在對應(yīng)的場景是否一定有這么大的提升？

不過，本文中并非要介紹 SGLang，而是旨在探討 vLLM 的基石——PagedAttention 的一些問題，以及現(xiàn)有的一些改進(jìn)工作，比如 vAttention 和 vTensor 等。此外，我們還會簡單介紹一些與 Attention 密切相關(guān)的 MHA/GQA 的實(shí)現(xiàn)考量。

需要注意的是，網(wǎng)上已經(jīng)有許多關(guān)于 vAttention 和 vTensor 的論文解讀，本文不會詳細(xì)介紹這些論文，甚至可能忽略一些實(shí)現(xiàn)的細(xì)節(jié)，這里只是通過這些論文引出一些需要關(guān)注的細(xì)節(jié)。

二、引言

2.1 MHA Attention 計(jì)算

如下圖所示為標(biāo)準(zhǔn)的 LLM Decoding 階段的 Multi-Head Attention（MHA）計(jì)算，其中的 D 表示 hidden size，H 表示 Head 個(gè)數(shù)，L 表示當(dāng)前是在序列的第 L 個(gè) Token。可以看出：

• 當(dāng)Batch Size 為 1時(shí)，圖中紅色、綠色、藍(lán)色處的矩陣乘法全部為矩陣乘向量，是明顯的 Memory Bound，算術(shù)強(qiáng)度不到 1。
• 當(dāng)Batch Size 大于 1時(shí)（比如 Continuous Batching）：

紅色和藍(lán)色部分：因?yàn)槭?Weight 乘以 Activation，所以不同的 Request 之間可以共享 Weight。這里變成矩陣乘矩陣，并且 Batch Size 越大，算術(shù)強(qiáng)度越大，也就越趨近于 Compute Bound（FFN 層也類似）。

綠色部分：這里 Q、K 和 V 的 Attention 計(jì)算，是 Activation 乘以 Activation，所以不同的 Request 之間沒有任何相關(guān)性。即使 Batching，這里也是Batched 矩陣乘向量，并且因?yàn)樾蛄虚L度可能不同，這里不同 Request 的矩陣乘向量是不規(guī)則的。也就是說，這里算術(shù)強(qiáng)度始終不到 1，是明顯的 Memory Bound。

從上可以看出，通過 Continuous Batching 可以很好的將 Memory Bound 問題轉(zhuǎn)變?yōu)?Compute Bound，但 Q、K 和 V 的 Attention 計(jì)算的算術(shù)強(qiáng)度卻始終小于 1。根據(jù) Amdahl 法則，如果系統(tǒng)中有一部分無法優(yōu)化，即使把其他部分優(yōu)化到可以忽略，不可優(yōu)化的部分也會決定整個(gè)系統(tǒng)的性能上限。不幸的是，Sequence Length 越長，這里的計(jì)算量就越不可忽略。

根據(jù)模型配置信息可以估算出模型中 Q、K 和 V 的 Attention 計(jì)算與其他矩陣計(jì)算的比例大約為 (L+D)/(12*D)（PS：準(zhǔn)確值需要根據(jù)具體的模型參數(shù)計(jì)算）。也就是說，當(dāng)序列長度 L 等于 12 倍的 hidden size 時(shí)，兩部分的計(jì)算量相當(dāng)，即使其他矩陣計(jì)算優(yōu)化到 0，加速比也只有 2x。比如 LLaMA 2 7B 的 hidden size 為 4K，當(dāng)序列長度達(dá)到 44K 時(shí)，兩部分的計(jì)算量相當(dāng)，要優(yōu)化的重點(diǎn)也會很不一樣，這也是很多長序列相關(guān)工作會在 Attention 部分采用稀疏 Attention 的一個(gè)重要原因。

2.2 GQA Attention 計(jì)算

早期通常只有比較大的模型才會采用 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 有個(gè)非常大的好處：在推理階段可以顯著降低 KV Cache 的大小，比如,相比 32 個(gè) KV Head 的 MHA，32 個(gè) Query Head，8 個(gè) KV Head 的 GQA 的 KV Cache 大小可以降低到 MHA 的 8/32=1/4，這也為更大的 Batch Size 提供了空間，可以進(jìn)一步提升吞吐。
• 除此之外，還有一個(gè)比較大的好處：可以明顯提升 Q、K 和 V 的 Attention 計(jì)算的算術(shù)強(qiáng)度。此時(shí)雖然不同的 Request 之間同樣不能共享，但是同一個(gè) Request 中的不同 Head 可以共享，比如 4 個(gè) Query Head 共享 1 個(gè) KV Head，則算術(shù)強(qiáng)度就會接近于 4，也可以更充分發(fā)揮 Tensor Core 的算力。

使用 MHA 時(shí)，Q、K 和 V 的 Attention 計(jì)算可以使用 CUDA Core 也可以使用 Tensor Core。由于 Tensor Core 要求矩陣的 Shape 是 8 的整數(shù)倍，如果不滿足就只能 Padding：

• 對于MHA而言，其是矩陣乘向量，則有7/8 的計(jì)算是冗余的。
• 對于GQA而言，如果 4 個(gè) Query Head 共享 1 個(gè) KV Head，則 Attention 計(jì)算有 4/8 的計(jì)算是冗余的，如果8 個(gè) Query Head 共享 1 個(gè) KV Head，則沒有計(jì)算的冗余。很多框架已經(jīng)做了相關(guān)優(yōu)化，比如 LMDeploy，TRT-LLM 的 XQA 等。
• 此外，PagedAttention 的 KV Cache 是非連續(xù)存儲的，導(dǎo)致即使使用 GQA 也無法利用 Tensor Core。

PS：對于 GQA 而言，理論上也可以期望 GPU 的 L2 Cache 能夠緩存到共享的 Key 和 Value Cache，從而緩解 IO Bound 問題，然而實(shí)際上無法人為控制，不一定能達(dá)到理想的效果。

2.3 NVIDIA 驅(qū)動

一般所說的 NVIDIA Driver 包含以下兩個(gè)部分，NVIDIA 從 2022 年 5 月正式開源的驅(qū)動程序也只是下述的 GPU Kernel-Mode Driver：NVIDIA Linux open GPU kernel module source，沒有開源 CUDA User-Mode Driver?？梢允褂?modinfo nvidia 或 cat /proc/driver/nvidia/version 查看當(dāng)前系統(tǒng)安裝的 GPU Kernel-Mode Driver 版本：

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

• nvidia.ko：提供對 NVIDIA 硬件的 Low-level 訪問，供其他組件使用。
• nvidia-uvm.ko：為 CUDA Driver 提供統(tǒng)一虛擬內(nèi)存（Unified Virtual Memory，UVM）功能。
• nvidia-drm.ko：向 Linux 內(nèi)核的 DRM 子系統(tǒng)注冊 DRM 驅(qū)動程序。
• nvidia-peermem.ko：提供 Mellanox InfiniBand 基于 HCA 的直接 P2P 讀寫 NVIDIA GPU 顯存的功能。

• CUDAUser-Mode Driver，部分如下所示：
• libcuda.so：為 CUDA 應(yīng)用程序提供運(yùn)行時(shí)支持。
• libnvidia-ml.so：提供監(jiān)控和管理 API。
• libnvidia-nvvm.so：由 CUDA 驅(qū)動程序加載，用于進(jìn)行 JIT 鏈接時(shí)優(yōu)化。
• libnvidia-ptxjitcompiler.so：將 PTX 編譯為 GPU 機(jī)器碼，并由 CUDA 驅(qū)動程序使用。

2.4 Low-level VMM API

CUDA 的 Low-level VMM API 從 CUDA 10.2 版本正式引入，用于更高效地管理GPU虛擬內(nèi)存。其提供如下一些主要的 API（部分），使得開發(fā)者能夠構(gòu)建更高效的動態(tài)數(shù)據(jù)結(jié)構(gòu)，更好地控制應(yīng)用程序中的 GPU 內(nèi)存使用：

• cuMemCreate：創(chuàng)建物理內(nèi)存句柄
• cuMemAddressReserve：保留虛擬地址范圍
• cuMemMap：將物理內(nèi)存句柄映射到虛擬地址范圍
• cuMemSetAccess：設(shè)置分配的內(nèi)存訪問權(quán)限

這里不再贅述，具體可以參考：Introducing Low-Level GPU Virtual Memory Management | NVIDIA Technical Blog。

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

三、PagedAttention

3.1 摘要

PagedAttention 是一種受操作系統(tǒng)中虛擬內(nèi)存和分頁技術(shù)啟發(fā)的注意力算法。在此基礎(chǔ)上，作者構(gòu)建了 vLLM 推理框架，可實(shí)現(xiàn)：

• 近似實(shí)現(xiàn) KV 的零浪費(fèi)。
• 在請求內(nèi)部和請求之間靈活共享 KV Cache，以進(jìn)一步減少內(nèi)存使用。

對應(yīng)的 Paper 為：[2309.06180] Efficient Memory Management for Large Language Model Serving with PagedAttention

對應(yīng)的 vLLM 的代碼庫：GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs

3.2 問題

傳統(tǒng)的 LLM 系統(tǒng)通常會為每個(gè) Request 預(yù)先分配顯存，如下圖 Figure 3 所示，假設(shè)模型支持的最大訓(xùn)練長度為 2048，則最多會為每個(gè) Request 預(yù)先分配 2048 個(gè) Token 的顯存空間。這樣如果實(shí)際執(zhí)行中的輸入和輸出包含 10 個(gè) Token，則將有 2038 個(gè) Token 空間的浪費(fèi)（內(nèi)存碎片），隨著服務(wù)支持的 Batch Size 增加，這一問題會愈加明顯：

3.3 方案

3.3.1 內(nèi)存管理

為了解決上述問題，作者參考操作系統(tǒng)中內(nèi)存碎片和 Sharing 的解決方案：帶分頁的虛擬內(nèi)存，并提出 PagedAttention。如下圖 Figure 6 所示，PagedAttention 將 Request 的 KV Cache 劃分為多個(gè) Block，每個(gè) Block 可以包含固定數(shù)量的 Token 的 Key 和 Value，在 Logical KV Blocks 中 Block 是連續(xù)的，但是在 Physical KV Blocks 中 Block 是不連續(xù)的（對應(yīng)預(yù)先分配的大塊物理內(nèi)存）。因此，可以像操作系統(tǒng)的虛擬內(nèi)存一樣，以更靈活的方式管理 KV Cache，當(dāng)然，也就需要 Block Table 來管理這種映射關(guān)系。這樣的好處是可以將整個(gè) Request 對應(yīng)的 KV Cache 劃分為相同大小的 Block，Block 可以遠(yuǎn)小于模型支持的序列長度，以此來明顯緩解內(nèi)存碎片化。

3.3.2 Prefix Caching

除了減少內(nèi)存碎片外，這種方式的另外一個(gè)好處是可以更好的進(jìn)行 Prefix Caching。如下圖所示，如果一個(gè) Request 要進(jìn)行并行采樣，或者兩個(gè) Request 有公共的前綴（比如 System Prompt），那么實(shí)際上只用計(jì)算、存儲一份 KV Cache 即可，即可減少計(jì)算，也可減少存儲。當(dāng)然，一個(gè) Block 中只會來自一個(gè) Request（Sample），因此都是按照整個(gè) Block 的粒度共享，如果 Block Size 過大，則可能影響可以共享的長度。比如，如果 Block Size 為 100，兩個(gè)序列在第 99 個(gè) Token 時(shí)出現(xiàn)不同，則無法實(shí)現(xiàn)共享。

3.3.3 Attention 計(jì)算

在實(shí)際的 Attention Kernel 計(jì)算時(shí)，會按照 Block 的粒度執(zhí)行，以 Query Token qi 為例（對應(yīng) “forth”），第一次會計(jì)算 qi 對應(yīng)的向量與 Block 0 中（對應(yīng) “Four score and seven”）Kj 的乘積，來計(jì)算 Attention Score Aij，第二次計(jì)算 Block 1，以此類推：

也就是說：PagedAttention 允許不同的 Block 在物理內(nèi)存中以不連續(xù)的方式存儲，可以充分增加內(nèi)存管理的靈活性。當(dāng)然，這也就導(dǎo)致實(shí)現(xiàn)的 CUDA Kernel 是和內(nèi)存管理耦合的，這也就增加了 Kernel 的計(jì)算開銷和實(shí)現(xiàn)的復(fù)雜度。

3.4 消融實(shí)驗(yàn)

3.4.1 Block Mapping 對性能的影響

PagedAttention 中的動態(tài) Block Mapping 會影響涉及存儲 KV Cache 的 GPU 操作的性能。與傳統(tǒng)連續(xù)存儲的方式相比，PagedAttention 中的 GPU Kernel 會引入訪問 Block Table、執(zhí)行額外的分支，以及處理可變序列長度的額外開銷。如下圖所示，作者與高度優(yōu)化的 FasterTransformer 實(shí)現(xiàn)相比，Attention Kernel 的延遲會高出 20-26%。因?yàn)橹粫绊?Attention Kernel，不會影響其他算子，比如 Linear Kernel，作者認(rèn)為其還可以接受（PS：這里用的 Context Length 非常短，隨著序列變長，差距可能會越來越大）：

3.4.2 Block Size 對性能的影響

Block 的大小可能會對 PagedAttention 的性能產(chǎn)生重大影響：

• 如果Block 太小，PagedAttention 可能無法充分利用 GPU 的并行性來讀取和處理 KV Cache。
• 如果Block 過大，則內(nèi)存碎片會增加，Prefix Cache 共享的可能性會降低。

如下圖所示，作者使用 ShareGPT 和 Alpaca 數(shù)據(jù)評估了不同 Block Size 下的 Latency（越低越好），可以看出，當(dāng) Block Size 為 16 和 32 時(shí)表現(xiàn)最好，因此作者在 vLLM 中將默認(rèn)的 Block Size 設(shè)置為 16：

3.5 問題

這里的 Block Size 應(yīng)該指的是 Token 數(shù)？那么不同的配置是否會有不同的最優(yōu)值，比如 LLaMA3 8B、70B、405B 的最優(yōu)配置是否相同？除此之外，使用了 GQA 的模型是否又會有不同的 Block Size？

四、GMLake

4.1 摘要

GMLake 是螞蟻和上海交通大學(xué)的工作，作者指出，GPU 原生的內(nèi)存分配器開銷很大，常見的 DNN 框架（比如 Pytorch 和 Tensorflow）會采用 Caching 機(jī)制，通過使用 GPU 原生的分配器分配大塊內(nèi)存，然后通過分割機(jī)制來實(shí)現(xiàn)快速分配和釋放。然而，Caching 機(jī)制會因?yàn)橹赜?jì)算、Offload、分布式訓(xùn)練以及低秩微調(diào)等方式而引入頻繁和不規(guī)則的內(nèi)存分配和釋放，導(dǎo)致存在嚴(yán)重的內(nèi)存碎片問題。

為了解決上述問題，作者提出了一種基于 Low-level 的 GPU 虛擬內(nèi)存管理的內(nèi)存分配框架，稱為 GMLake（GPU Memory Lake）。使用 GMLake，可以融合或連接非連續(xù)的內(nèi)存塊，并通過虛擬內(nèi)存地址進(jìn)行映射。在 A100 80GB GPU 上，通過這種方式可以顯著減少 GPU 內(nèi)存使用量（平均減少 9.2 GB，最多 25 GB）和內(nèi)存碎片（平均減少 15%，最多 33%）。

對應(yīng)的論文為：[2401.08156] GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching

4.2 背景

如下圖 Figure 2 所示為 3 種分配機(jī)制的區(qū)別：

• 原生的 GPU 內(nèi)存分配：每次都調(diào)用 cudaMalloc 分配，調(diào)用 cudaFree 釋放。
• Caching 分配：Pytorch 和 Tensorflow 采用的 BFC 算法，預(yù)先分配大塊內(nèi)存，然后通過查找、拆分、合并等方式實(shí)現(xiàn)最大化利用，減少 cudaMalloc 和 cudaFree 的巨大開銷。作者實(shí)驗(yàn)表明，Caching 分配的吞吐大概是原生 GPU 內(nèi)存分配的 10 倍。
• Virtual Memory：利用 GPU 的 Low-level 虛擬內(nèi)存管理方案分配。

4.3 方案

4.3.1 概覽

如下圖 Figure 7 所示為 GMLake 的方案概覽，它提供了與現(xiàn)有的 Caching Allocator 接口相同的 GMLake Allocator，但是在內(nèi)部集成了虛擬內(nèi)存拼接機(jī)制（Virtual Memory Stiching，VMS），其主要是依賴 CUDA 的 Low-level VM 管理 API實(shí)現(xiàn)。GMLake 主要是包含 3 個(gè)組件：

• Virtual memory API：用于指示 GPU 使用虛擬內(nèi)存地址分配和釋放內(nèi)存的 Low-level API。
• Virtual memory pool：作為基礎(chǔ)數(shù)據(jù)結(jié)構(gòu)，用于緩存虛擬內(nèi)存，提高效率。
• GMLake allocator：包括管理 VM 池所必須的所有函數(shù)、算法和策略。

4.3.2 Stitched Memory Pool & Primitive Memory Pool

原始的 Low-level VMM API 也非常耗時(shí)，因此減少其使用量對于實(shí)現(xiàn)高效率 GMLake 至關(guān)重要。作者從 Caching Allocator 中汲取靈感，設(shè)計(jì)了具有 Caching 功能的虛擬內(nèi)存池（VMP），從而顯著減少物理內(nèi)存分配的次數(shù)。如下圖 Figure 8 所示，作者設(shè)計(jì)了兩種內(nèi)存池：

• Primitive Memory Pool（pPool）：pPool 使用有序集合來存儲 pBlock，對于每個(gè) pBlock，pPool 首先構(gòu)造一個(gè)結(jié)構(gòu)來記錄指向 pBlock 的指針，其包含基礎(chǔ)屬性，比如 pBlock 的激活狀態(tài)。隨后，將新分配的 pBlock 插入到集合中，所有 pBlock 按照塊大小降序排列。pBlock 作為原始塊，代表 High-level Tensor 可訪問的最小單元（Low-level 能分配的最小 Physical Chunk 為 2MB，一個(gè) pBlock 可以包含多個(gè) Physical Block），它作為基本數(shù)據(jù)結(jié)構(gòu)，可以被多個(gè) sBlock 拼接或指向。
• Stitched Memory Pool（sPool）：sPool 也被組織成一個(gè)有序集合，類似于 pPool。它的元素包含拼接的 block 結(jié)構(gòu)，該結(jié)構(gòu)組合了多個(gè) pBlock。比如下圖中的 sBlock 3 包含一個(gè)組合在一起的 pBlock 3 和 pBlock 5。同時(shí)，pBlock 3 也可以指向 sBlock 2。

4.4 問題

LLM 推理與訓(xùn)練不同，訓(xùn)練中通常會將輸入序列拼接到模型支持的最大序列長度，比如 4096 Token，可以充分提高效率，并且是等價(jià)的。這種方式對于內(nèi)存分配相對比較友好，并且通常是一些大塊的內(nèi)存分配。而在 LLM 推理場景，輸入、輸出的序列長度可能差異很大，尤其是 Decoding 階段是一個(gè)一個(gè) Token 生成，導(dǎo)致分配的最小粒度變?yōu)閱蝹€(gè) Token，就會涉及很多小塊內(nèi)存分配，也就需要更精細(xì)化的內(nèi)存管理。

五、vTensor

5.1 摘要

vTensor 同樣由螞蟻和上海交大發(fā)表，和 GMLake 大部分作者相同，可以認(rèn)為是將 GMLake 由 DNN Training 擴(kuò)展到 LLM Inference 場景。其比 vAttention 晚發(fā)表幾個(gè)月，和 vAttention 的思路非常類似，不過兩個(gè)工作都是最近才開源的。

很自然，vTensor 也是一種基于 GPU 虛擬內(nèi)存管理（VMM）的 LLM 推理 Tensor 結(jié)構(gòu)。vTensor 通過將計(jì)算與內(nèi)存碎片整理解耦并提供動態(tài)可擴(kuò)展性來解決現(xiàn)有的限制（主要指 PagedAttention）。其采用 CPU-GPU 異構(gòu)方案，確保高效、無碎片的內(nèi)存管理（PS：近似？），同時(shí)可以適應(yīng)不同 LLM 架構(gòu)的各種計(jì)算 Kernel。

實(shí)驗(yàn)表明，vTensor 在不同的模型中實(shí)現(xiàn)了平均 1.86x 加速，在多輪聊天場景中最高可達(dá) 2.42x。此外，vTensor 在 Kernel 評估中，可以實(shí)現(xiàn)平均 2.12x 和 3.15x 加速，與 SGLang Triton prefix-prefilling Kernel 和 vLLM 的 PagedAttention Kernel 相比，在 A100 上可以釋放 71.25%（57GB）內(nèi)存，從而支持更多內(nèi)存密集型工作負(fù)載。

對應(yīng)的論文為：[2407.15309] vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving

對應(yīng)的代碼庫為：??https://github.com/intelligent-machine-learning/glake/tree/master/GLakeServe??

5.2 方案

5.2.1 vTensor

如下圖 Figure 1 所示為本文提出的 vTensor 與原始的 KV Cache 機(jī)制以及 Paged KV Cache 的區(qū)別，和 vAttention 類似，也是基于 Low-level 的 vMM API，有 3 個(gè)好處：

• 可以實(shí)現(xiàn)內(nèi)存管理與 CUDA Kernel 的解構(gòu)，更加通用，Kernel 實(shí)現(xiàn)更加簡單。
• 可以實(shí)現(xiàn)內(nèi)存的動態(tài)擴(kuò)展，減少浪費(fèi)。
• Attention Kernel 可以使用更強(qiáng)算力的Tensor Core（虛擬地址連續(xù)），而 PagedAttention 只能使用 CUDA Core。

?

如下圖 Figure 5 所示為具體的 vTensor 的實(shí)現(xiàn)，vTensor 指針由 vTensor Manager（VTM）生成。

• 當(dāng)向vTensor Scheduler（VTS）發(fā)送創(chuàng)建請求時(shí)，VTS 將創(chuàng)建分配策略，然后讓 vTensor Operation（VTO）分配虛擬內(nèi)存和相應(yīng)的 Physical Chunk（PC）。
• vTensor 指針 *A 指向 GPUVirtual Address（VA），該地址必須是連續(xù)的，才能與標(biāo)準(zhǔn) CUDA 分配的 Tensor 兼容。
• Virtual Memory Management（VMM）維護(hù) VTM 注冊的 Physical Chunk 的完整映射信息，VMM 允許 GPU Tensor Core 通過 Virtual Address 訪問 GPU 內(nèi)存中所需的數(shù)據(jù)。
• Physical Chunk在 GPU 內(nèi)存中分配，但是Physical Chunk Handle（PH）和 Virtual Address 可以由 CPU 訪問和管理。Physical Handle 和 Virtual Memory 僅具有 Physical Chunk 的索引信息，而不包含 Device-Host 傳輸。
• Physical Chunk同樣采用了 2MB 的 Physical Chunk，其對應(yīng)的 Handle 只有幾個(gè)字節(jié)，由 vTensor Pool（VTP）記錄并存儲在 CPU 內(nèi)存中。
• 作者也開發(fā)了一系列基于 vTensor 的方法來操作 KV Cache 的碎片整理，這是 vTensor 和 FlexInfer 設(shè)計(jì)的基礎(chǔ)。

5.2.2 FlexInfer

基于 vTensor，作者進(jìn)一步開發(fā)了 FlexInfer，它是一個(gè) CPU 和 GPU 異構(gòu)框架，將大多數(shù)內(nèi)存操作解耦并卸載到 CPU，并通過重疊 GPU 計(jì)算來隱藏它們。與之前在 GPU 上內(nèi)存操作（比如 PagedAttention）相比，CPU 更擅長與內(nèi)存相關(guān)的操作。當(dāng)請求發(fā)送到 FlexInfer 時(shí)，它會根據(jù)請求的配置（例如 Batch Size 和 Seq Length）解耦內(nèi)存和計(jì)算操作。

• 在計(jì)算方面，F(xiàn)lexInfer Scheduler 采用原始的高度優(yōu)化的 Kernel 在 GPU Tensor Core 上運(yùn)行 LLM，在不受算術(shù)強(qiáng)度限制的情況下保持計(jì)算靈活性和效率。
• 在內(nèi)存方面，F(xiàn)lexInfer 還提供了高度定制的調(diào)度方案，以在啟動、Prefill、Decoding 以及終止階段與計(jì)算重疊，隱藏內(nèi)存分配和釋放，可以顯著降低 LLM Service System 的內(nèi)存操作開銷。

5.3 消融實(shí)驗(yàn)

5.3.1 Decoding Kernel 評估

如下圖 Figure 7 所示，作者對比了不同場景下相應(yīng) Attention Kernel 的性能，其中 FlexInfer attn 表示使用 vTensor 的 FlashAttention，Paged flash attn 表示使用 Paged 的 FlashAttention，F(xiàn)lash attn 表示原始的 FlashAttention：

• Batch Size：隨著 Batch Size 增加，F(xiàn)lexInfer attn 和 Flash attn 始終保持最低 Latency，F(xiàn)lexInfer attn Latency 平均比 Paged attn 低 42%。
• Sequence Length：以 Batch Size 16 為例，隨著 Sequence Length 增加，F(xiàn)lexInfer attn 和 Flash attn 依然保持最低 Latency，在 1K 時(shí)加速最明顯（但是在 1K 序列長度時(shí)，Attention 在整個(gè)計(jì)算中的占比也比較?。?。
• KV Head：以 Batch Size 16 和 Sequence Length 16K 為例，其中 KV Head 為 1 對應(yīng) MQA，KV Head 32 對應(yīng) MHA（Yi-6B、Yi-9B），KV Head 4 和 8 對應(yīng) GQA?？梢钥闯?/span>

MQA（KV Head 1）時(shí)，加速最明顯，可以最充分發(fā)揮 Tensor Core 的算力。

MHA（KV Head 32）時(shí)，基本沒有加速，此時(shí)都無法充分發(fā)揮 Tensor Core 算力。

GQA（KV Head 4 和 8）時(shí)，有一定加速，可以部分發(fā)揮 Tensor Core 算力。

5.3.2 Prefix-prefilling Kernel 評估

如下圖 Figure 8 所示，作者進(jìn)一步對比了 Prefilling 階段 Kernel 的性能（對應(yīng)的 Sequence Length 固定為 16K），可以看出，在不同 Batch Size 和 Prefix/Prompt 比例下，PagedAttention 的性能都是最差的，其他幾種方式性能相當(dāng)。當(dāng)然，將 Paged KV Cache 適配到 FlashAttention 的代價(jià)也很高，而 FlexInfer attn 的實(shí)現(xiàn)代價(jià)小得多。

5.4 問題

六、vAttention

6.1 摘要

vAttention 是微軟的工作，其發(fā)表在 GMLake 和 vTensor 之間，思路和 vTensor 非常接近。同樣是使用 Low-level 的 VMM API 實(shí)現(xiàn)，也同樣是為了減少 KV Cache 內(nèi)存碎片，降低 Attention Kernel 的開發(fā)成本。與 vTensor 不同的是，作者還進(jìn)一步修改了 GPU 內(nèi)核 Driver，以提供更細(xì)粒度的 Physical Chunk 的分配，比如 64KB、128KB 和 256KB，而不局限于 2MB。

結(jié)果表明，vAttention 可以為各種 Attention Kernel 實(shí)現(xiàn)無縫的兼容，并提供動態(tài)內(nèi)存管理能力。vAttention 生成 Token 的速度比 vLLM 快 1.97x，同時(shí)處理 Prompt 的速度比 FlashAttention 和 FlashInfer 的 PagedAttention 快 3.92x 和 1.45x。

對應(yīng)的論文為：[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，作者認(rèn)為其依然會導(dǎo)致存在一部分的內(nèi)存碎片（PS：后文會介紹），因此決定修改 CUDA Low-level API，以允許分配更細(xì)粒度的 Physical Chunk，比如 64KB、128KB 和 256KB。并提供了一些新的 API：

如上只是一部分代碼修改，實(shí)際的修改有 200-300 行，具體可以參考 ??https://github.com/microsoft/vattention/tree/main/nvidia-vattn-uvm-driver??。需要指出的是，這個(gè)修改是針對 545.23.06 版本，其他版本需要進(jìn)一步適配；此外，因?yàn)樯婕傲说讓拥?nvidia-uvm Driver 的修改，因此需要替換系統(tǒng)已有 Driver 并且重新啟動 Server，相應(yīng)的代價(jià)也比較高。

如下圖 Table 3 所示，作者也進(jìn)一步對相關(guān) API 進(jìn)行了封裝，并測試了不同分配大小的時(shí)延：

6.2.2 vAttention

具體的思路和 vTensor 類似，不再贅述，這里階段介紹一個(gè) vAttention 中動態(tài)內(nèi)存管理的示例：

• a：兩個(gè)請求 R1 和 R2 的 virtural tensor，沒有使用 Physical Memory。
• b：R1 分配了一個(gè) Physical Page。
• c：R1 分配了兩個(gè) Physical Page，R2 分配了一個(gè) Physical Page。
• d：R1 處理完，但是不會立即釋放對應(yīng)的 Physical Page；R2 分配了兩個(gè) Physical Page。
• e：新的請求 R3 分配了兩個(gè) Physical Page，會利用之前 R2 分配的 Physical Page。

其實(shí) vLLM 也已經(jīng)在 PR（??https://github.com/vllm-project/vllm/pull/6102??）中提供了對 vAttention 的支持，但是目前還沒有合入，也沒有實(shí)現(xiàn)所有功能。比如，當(dāng)前還只支持 2MB 的 Physical Chunk，可能會存在比較大的顯存碎片（下面會介紹）。

6.3 消融實(shí)驗(yàn)

6.3.1 Physical Chunk 大小對 Prefill 的影響

在 Prefill 階段的 Token 數(shù)比較多，每層的 K 或 V Cache 可能遠(yuǎn)超 2M，此時(shí)使用過小的 Chunk 有可能會引入比較多的開銷。如下圖 Figure 11 所示，使用 64KB Chunk 最多會導(dǎo)致 Prefill Latency 增加 15%，不過通過計(jì)算和分配的 Overlap 可以隱藏掉這一部分開銷（對應(yīng) vAttention）。

6.3.2 Physical Chunk 大小對分配帶寬的影響

如下圖 Table 7 所示，使用更小的 Physical Chunk Size 會導(dǎo)致每秒分配的顯存大小降低，當(dāng) Physical Chunk Size 為 64KB 時(shí)，每秒分配的顯存量只有 2MB 時(shí)的 1/5 左右：

如果 64KB 時(shí)的分配速度無法滿足實(shí)際需要的分配速度，那么就可能成為瓶頸。作者也進(jìn)行了相應(yīng)的測試，如下圖所示，當(dāng) Batch Size 增加到 320 時(shí)（綠線 LLama3-8B - 2 A100，實(shí)線 Yi-6B - 1A100，藍(lán)線 Yi-34B - 2 A100），需要的最大分配帶寬也只有 600MB/s，遠(yuǎn)小于 64KB 對應(yīng)的 7.59GB/s，也證明 64KB 的 Physical Chunk 完全可以接受：

6.3.3 Memory 碎片

在進(jìn)行 Attention 計(jì)算時(shí)，各個(gè) Request 之間是沒有任何關(guān)系的，即使是 Continuous Batching，也是 N 個(gè)矩陣乘向量（MHA）或者 N 個(gè)矩陣乘矩陣（GQA、MQA）。因此，不管是 PagedAttention 還是 vTensor 或者 vAttention，其每個(gè) Chunk 中都只會存儲同一個(gè) Request 的 Token，此時(shí)，在每個(gè) Request 的最后一個(gè) Chunk 中就可能存來空閑未被使用的空間，Chunk 的 Size 越大，理論浪費(fèi)的空間就越大。

如下圖 Table 8 所示，作者以 Yi-6B、LLaMA-3-8B 和 Yi-34B 為例，統(tǒng)計(jì)了不同 Chunk 可以容納的 Token 數(shù)，以及浪費(fèi)的最大 Memory 空間：

• 模型結(jié)構(gòu)：

Yi-6B 的 Hidden Size 為 4096，總共 32 個(gè) Q-Head，4 個(gè) KV-Head，32 個(gè) Layer。

Yi-34B 的 Hidden Size 為 7168，總共 56 個(gè) Q-Head，8 個(gè) KV-Head， 60 個(gè) Layer。

LLaMA-3-8B 的 Hidden Size 為 4096，總共 32 個(gè) Q-Head，8 個(gè) KV-Head， 32 個(gè) Layer。

• 每一層一個(gè) Token 的 Key 或 Value Cache 的大小為（FP16 存儲，如果采用 TP，則一般會按照 Head 切分，所以每個(gè) GPU 上 1 個(gè) Token 相應(yīng)的存儲占用的空間減少）：
• 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
• 每一層一個(gè) Chunk 可以存儲的 Key 或 Value Cache 的 Token 數(shù)為（TP1）：
• Yi-6B 64KB：64KB/1KB=64
• Yi-34B 256KB：256KB/2KB=128
• 每一層最多浪費(fèi) Key 和 Value 兩個(gè) Chunk，則理論上一個(gè) Request 浪費(fèi)的最大空間為：
• Yi-6B 2MB：2MB*2(K/V)*32(Layer)=128MB
• LLaMA-3-8B 128KB：128KB*2*32(Layer)=8MB

從以上的統(tǒng)計(jì)數(shù)據(jù)可以看出，每個(gè)請求最大的 Memory 浪費(fèi)與 Chunk Size 成正比，以 Batch Size 100，TP=1 為例：

• 2MB 的 Chunk Size 最大浪費(fèi) 12.5GB-23.4GB（128MB-240MB * 100）。
• 64KB 的 Chunk Size 最大浪費(fèi) 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??

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談