今年六月，來自加州大學(xué)伯克利分校等機(jī)構(gòu)的一個(gè)研究團(tuán)隊(duì)開源了 vLLM（目前已有 6700 多個(gè) star），其使用了一種新設(shè)計(jì)的注意力算法 PagedAttention，可讓服務(wù)提供商輕松、快速且低成本地發(fā)布 LLM 服務(wù)。

在當(dāng)時(shí)的博客文章中，該團(tuán)隊(duì)宣稱 vLLM 能實(shí)現(xiàn)比 HuggingFace Transformers 高 24 倍的吞吐量！

八月底，該項(xiàng)目還與其它另外 7 個(gè)開源項(xiàng)目一道獲得了 a16z 開源人工智能資助計(jì)劃的首輪資助。

現(xiàn)在離 vLLM 最初宣布時(shí)差不多已過去三個(gè)月，他們終于發(fā)布了這篇關(guān)于 vLLM 和 PagedAttention 的研究論文，其中詳細(xì)解釋了他們?nèi)绾瓮ㄟ^類似操作系統(tǒng)虛擬內(nèi)存管理的機(jī)制來實(shí)現(xiàn)高效的 LLM 服務(wù)。該論文已被將于十月底在德國(guó)舉辦的 ACM 操作系統(tǒng)原理研討會(huì)（SOSP 2023）接收。

• 論文：https://arxiv.org/abs/2309.06180
• 代碼：https://github.com/vllm-project/vllm
• 文檔：https://vllm.readthedocs.io/

GPT 和 PaLM 等大型語言模型（LLM）的出現(xiàn)催生出了開始對(duì)我們的工作和日常生活產(chǎn)生重大影響的新應(yīng)用，比如編程助理和通用型聊天機(jī)器人。

許多云計(jì)算公司正競(jìng)相以托管服務(wù)的方式提供這些應(yīng)用。但是，運(yùn)行這些應(yīng)用的成本非常高，需要大量硬件加速器，如 GPU。根據(jù)最近的估計(jì)，相比于傳統(tǒng)的關(guān)鍵詞查詢方法，處理一個(gè) LLM 請(qǐng)求的成本超過其 10 倍以上?？紤]到成本如此之高，提高 LLM 服務(wù)系統(tǒng)的吞吐量（并由此降低單位請(qǐng)求的成本）就變得更為重要了。

LLM 的核心是自回歸 Transformer 模型。該模型可基于輸入（prompt）和其之前輸出的 token 序列生成詞（token），一次生成一個(gè)。對(duì)于每次請(qǐng)求，這個(gè)成本高昂的過程都會(huì)重復(fù)，直到模型輸出終止 token。這種按序列的生成過程會(huì)讓工作負(fù)載受到內(nèi)存限制，從而無法充分利用 GPU 的計(jì)算能力，并會(huì)限制服務(wù)的吞吐量。

通過批量方式同時(shí)處理多個(gè)請(qǐng)求可以提高吞吐量。但是，要在單一批次中處理許多請(qǐng)求，就需要高效地管理每個(gè)請(qǐng)求所占用的內(nèi)存空間。

舉個(gè)例子，圖 1（左）展示了一個(gè) 130 億參數(shù)的 LLM 在一臺(tái) 40GB RAM 的英偉達(dá) A100 GPU 上的內(nèi)存分布。

其 65% 的內(nèi)存都分配給了模型權(quán)重，而模型權(quán)重在提供服務(wù)期間是不會(huì)變化的。

30% 的內(nèi)存是用于存儲(chǔ)請(qǐng)求的動(dòng)態(tài)狀態(tài)。對(duì) Transformer 而言，這些狀態(tài)由與注意力機(jī)制關(guān)聯(lián)的鍵（key）和值（value）張量構(gòu)成，通常被稱為 KV 緩存，其表示用于生成序列中新輸出 token 的之前 token 上下文。

其余占比很小的內(nèi)存則是用于其它數(shù)據(jù)，包括激活 —— 評(píng)估 LLM 時(shí)創(chuàng)建的臨時(shí)張量。

由于模型權(quán)重恒定不變，激活也只會(huì)占用少量 GPU 內(nèi)存，因此對(duì) KV 緩存的管理方式就成了決定最大批量大小的關(guān)鍵。如果管理方式很低效，KV 緩存內(nèi)存就會(huì)極大限制批量大小，并由此限制 LLM 的吞吐量，如圖 1（右）所示。

來自 UC 伯克利等機(jī)構(gòu)的這個(gè)研究團(tuán)隊(duì)在論文中表示，他們觀察到當(dāng)前的 LLM 服務(wù)系統(tǒng)都沒有高效地管理 KV 緩存內(nèi)存。主要原因是它們會(huì)將請(qǐng)求的 KV 緩存保存在鄰接的內(nèi)存空間中，因?yàn)榇蠖鄶?shù)深度學(xué)習(xí)框架都需要將張量存儲(chǔ)在相鄰連續(xù)的內(nèi)存中。

但是，不同于傳統(tǒng)深度學(xué)習(xí)工作負(fù)載中的張量，KV 緩存有其自己的獨(dú)特性質(zhì)：它會(huì)在模型生成新 token 的過程中隨時(shí)間動(dòng)態(tài)地增長(zhǎng)和縮小，而且它的持續(xù)時(shí)間和長(zhǎng)度是無法事先知曉的。

這些特性使得現(xiàn)有系統(tǒng)的方法極其低效，這主要體現(xiàn)在兩方面：

第一，現(xiàn)有系統(tǒng)存在內(nèi)部和外部?jī)?nèi)存碎片的問題。

為了將請(qǐng)求的 KV 緩存存儲(chǔ)在相鄰連續(xù)的空間中，現(xiàn)有系統(tǒng)會(huì)預(yù)先分配一塊限制了請(qǐng)求最大長(zhǎng)度的相鄰連續(xù)內(nèi)存。這會(huì)導(dǎo)致嚴(yán)重的內(nèi)部碎片化，因?yàn)檎?qǐng)求的實(shí)際長(zhǎng)度可能比其最大長(zhǎng)度短很多。

此外，就算能事先知道實(shí)際長(zhǎng)度，預(yù)先分配內(nèi)存的方法依然很低效。因?yàn)槟且徽麎K內(nèi)存在請(qǐng)求的生命周期中都被占用了，其它更短的請(qǐng)求無法使用這一塊內(nèi)存 —— 盡管可能其中很大一部分都未被使用。

此外，外部?jī)?nèi)存碎片可能也很顯著，因?yàn)槊總€(gè)請(qǐng)求預(yù)先分配的內(nèi)存大小可能都不同。

事實(shí)上，我們可以在圖 2 的分析中看到，在現(xiàn)有系統(tǒng)中，僅有 20.4%-38.2% 的 KV 緩存內(nèi)存會(huì)用于存儲(chǔ)實(shí)際的 token 狀態(tài)。

第二，現(xiàn)有系統(tǒng)無法利用內(nèi)存共享的機(jī)會(huì)。

LLM 服務(wù)通常會(huì)使用先進(jìn)的解碼算法，例如并行采樣和波束搜索，這些方法可為每個(gè)請(qǐng)求生成多個(gè)輸出。在這些場(chǎng)景中，由多個(gè)序列組成的請(qǐng)求可以部分共享它們的 KV 緩存。但是，現(xiàn)有系統(tǒng)不可能使用內(nèi)存共享，因?yàn)椴煌蛄械?KV 緩存存儲(chǔ)在分開的相鄰連續(xù)空間中。

為了解決上述限制，該團(tuán)隊(duì)提出了一種注意力算法 PagedAttention；該算法的靈感來自操作系統(tǒng)（OS）解決內(nèi)存碎片化和內(nèi)存共享的方案：使用分頁(yè)機(jī)制的虛擬內(nèi)存。

PagedAttention 會(huì)將請(qǐng)求的 KV 緩存分成一塊塊的，每一塊（block）都包含一定數(shù)量 token 的注意力鍵和值。在 PagedAttention 中，KV 緩存的塊不一定要存儲(chǔ)在相鄰連續(xù)空間中。

這樣一來，就能以一種更為靈活的方式來管理 KV 緩存，就像是操作系統(tǒng)的虛擬內(nèi)存：你可以將那些塊看作是分頁(yè)，將 token 看作是字節(jié)，將請(qǐng)求視為進(jìn)程。通過使用相對(duì)較小的塊并按需分配它們，這種設(shè)計(jì)可以減少內(nèi)部碎片。

此外，它還能消除外部碎片，因?yàn)樗袎K的大小都相同。

最后，它還能實(shí)現(xiàn)以塊為粒度的內(nèi)存共享；這種內(nèi)存共享支持與同一請(qǐng)求相關(guān)的不同序列，甚至也支持不同的請(qǐng)求。

立足于 PagedAttention，該團(tuán)隊(duì)構(gòu)建了一個(gè)高吞吐量的分布式 LLM 服務(wù)引擎 vLLM，其幾乎做到了 KV 緩存內(nèi)存的零浪費(fèi)。

vLLM 使用了塊級(jí)的內(nèi)存管理和搶占式的請(qǐng)求調(diào)度 —— 這些機(jī)制都是配合 PagedAttention 一起設(shè)計(jì)的。vLLM 支持 GPT、OPT 和 LLaMA 等各種大小的常用 LLM，包括那些超出單個(gè) GPU 內(nèi)存容量的 LLM。

研究者基于多種模型和工作負(fù)載進(jìn)行了實(shí)驗(yàn)評(píng)估，結(jié)果表明：相比于當(dāng)前最佳的系統(tǒng)，vLLM 能在完全不影響準(zhǔn)確度的前提下將 LLM 服務(wù)吞吐量提升 2-4 倍。而且當(dāng)序列更長(zhǎng)、模型更大、解碼算法更復(fù)雜時(shí)，vLLM 帶來的提升還會(huì)更加顯著。

總體而言，該研究的貢獻(xiàn)如下：

• 確認(rèn)了提供 LLM 服務(wù)時(shí)的內(nèi)存分配挑戰(zhàn)并量化了它們對(duì)服務(wù)性能的影響。 
• 提出了 PagedAttention，這是一種對(duì)存儲(chǔ)在非相鄰連續(xù)的分頁(yè)內(nèi)存中的 KV 緩存進(jìn)行操作的注意力算法，其靈感來自操作系統(tǒng)中的虛擬內(nèi)存和分頁(yè)。 
• 設(shè)計(jì)并實(shí)現(xiàn)了 vLLM，這是一個(gè)基于 PagedAttention 構(gòu)建的分布式 LLM 服務(wù)引擎。 
• 通過多種場(chǎng)景對(duì) vLLM 進(jìn)行了實(shí)驗(yàn)評(píng)估，結(jié)果表明其表現(xiàn)顯著優(yōu)于 FasterTransformer 和 Orca 等之前最佳方案。

方法概況

vLLM 的架構(gòu)如圖 4 所示。

vLLM 采用一種集中式調(diào)度器（scheduler）來協(xié)調(diào)分布式 GPU 工作器（worker）的執(zhí)行。KV 緩存管理器由 PagedAttention 驅(qū)動(dòng)，能以分頁(yè)方式有效管理 KV 緩存。具體來說，KV 緩存管理器通過集中式調(diào)度器發(fā)送的指令來管理 GPU 工作器上的物理 KV 緩存內(nèi)存。

PagedAttention

不同于傳統(tǒng)的注意力算法，PagedAttention 支持將連續(xù)的鍵和值存儲(chǔ)在非相鄰連續(xù)的內(nèi)存空間中。

具體來說，PagedAttention 會(huì)將每個(gè)序列的 KV 緩存分成 KV 塊。每一塊都包含固定數(shù)量 token 的鍵和值的向量；這個(gè)固定數(shù)量記為 KV 塊大小（B）。令第 j 個(gè) KV 塊的鍵塊為 K_j，值塊為 V_j。則注意力計(jì)算可以轉(zhuǎn)換為以下形式的對(duì)塊的計(jì)算：

其中 A_{i,j} 是在第 j 個(gè) KV 塊上的注意力分?jǐn)?shù)的行向量。

在注意力計(jì)算期間，PagedAttention 核會(huì)分開識(shí)別并獲取不同的 KV 塊。

圖 5 給出了 PagedAttention 的一個(gè)示例：其鍵和值向量分布在三個(gè)塊上，并且這三個(gè)塊在物理內(nèi)存上并不相鄰連續(xù)。

每一次，這個(gè) PagedAttention 核都會(huì)將查詢 token（forth）的查詢向量 q_i 與一個(gè)塊（比如 0 塊中的 Four score and seven 的鍵向量）中鍵向量 K_j 相乘，以計(jì)算注意力分?jǐn)?shù) A_{i,j}；然后再將 A_{i,j} 與塊中的值向量 V_j 相乘，得到最終的注意力輸出 o_i。

綜上所述，PagedAttention 算法能讓 KV 塊存儲(chǔ)在非相鄰連續(xù)的物理內(nèi)存中，從而讓 vLLM 實(shí)現(xiàn)更為靈活的分頁(yè)內(nèi)存管理。

KV 緩存管理器

使用 PagedAttention，該團(tuán)隊(duì)將 KV 緩存組織為固定大小的 KV 塊，就像虛擬內(nèi)存中的分頁(yè)。

對(duì) KV 緩存的請(qǐng)求會(huì)被表示成一系列邏輯 KV 塊，在生成新 token 和它們的 KV 緩存時(shí)從左向右填充。最后一個(gè) KV 塊中未填充的位置留給未來填充。

使用 PagedAttention 和 vLLM 進(jìn)行解碼

圖 6 通過一個(gè)示例展示了 vLLM 在對(duì)單個(gè)輸入序列的解碼過程中執(zhí)行 PagedAttention 和管理內(nèi)存的方式。

從全局來看，在每次解碼迭代中，vLLM 首先會(huì)選取一組候選序列來批處理，并為新請(qǐng)求的邏輯塊分配物理塊。

然后，vLLM 會(huì)將當(dāng)前迭代的所有輸入 token 連接起來，組成一個(gè)序列并將其輸入到 LLM。在 LLM 的計(jì)算過程中，vLLM 使用 PagedAttention 核來訪問以邏輯 KV 塊形式存儲(chǔ)的之前的 KV 緩存，然后將新生成的 KV 緩存保存到物理 KV 塊中。

在一個(gè) KV 塊中存儲(chǔ)多個(gè) token（塊大小 > 1）可讓 PagedAttention 核并行處理多個(gè)位置的 KV 緩存，由此可以提升硬件使用率并降低延遲。

圖 7 給出了 vLLM 管理兩個(gè)序列的內(nèi)存的示例。

應(yīng)用于其它解碼場(chǎng)景

在許多成功的 LLM 應(yīng)用中，LLM 服務(wù)必須要能提供更為復(fù)雜的解碼場(chǎng)景 —— 有復(fù)雜的訪問模式和更多的內(nèi)存共享機(jī)會(huì)。研究者展示了 vLLM 的一般適用性。他們?cè)谡撐闹杏懻撨^的場(chǎng)景包括并行采樣、波束搜索、共享前綴、混合解碼方法。

比如圖 8 展示了一個(gè)帶有兩個(gè)輸出的并行編碼的例子。

由于兩個(gè)輸出共享同一個(gè) prompt，所以在 prompt 階段只需為 prompt 的狀態(tài)在內(nèi)存空間中保留一個(gè)副本；這兩個(gè)序列的 prompt 的邏輯塊都會(huì)映射到同一個(gè)物理塊。更多分析和示例請(qǐng)參閱原論文。

調(diào)度和搶占

LLM 面臨著一個(gè)獨(dú)特挑戰(zhàn)：輸入 LLM 的 prompt 的長(zhǎng)度變化范圍很大，并且輸出結(jié)果的長(zhǎng)度也無法事先知曉，而要依輸入 prompt 和模型的情況而定。隨著請(qǐng)求和相應(yīng)輸出數(shù)量的增長(zhǎng)，vLLM 可能會(huì)為了存儲(chǔ)新生成的 KV 緩存而耗盡 GPU 的物理塊。

對(duì)于這種情況，vLLM 需要回答兩個(gè)經(jīng)典問題：(1) 應(yīng)該淘汰（evict）哪些塊？(2) 如果再次需要，如何恢復(fù)已淘汰的塊？

通常來說，淘汰策略是使用啟發(fā)式方法來預(yù)測(cè)將在最遠(yuǎn)的未來訪問的塊然后淘汰那個(gè)塊。由于這里已知是一起訪問一個(gè)序列的所有塊，所以該團(tuán)隊(duì)實(shí)現(xiàn)了一種全清或不動(dòng)的淘汰策略，即要么淘汰序列的所有塊，要么就全不淘汰。此外，一個(gè)請(qǐng)求中的多個(gè)序列會(huì)被組成一個(gè)序列組來一起調(diào)度。一個(gè)序列組中的序列總是會(huì)被一起搶占或重新調(diào)度，因?yàn)檫@些序列之間存在潛在的內(nèi)存共享。

為了解答第二個(gè)有關(guān)如何恢復(fù)已淘汰塊的問題，研究者考慮了兩種技術(shù)：

• 交換。這是大多數(shù)虛擬內(nèi)存實(shí)現(xiàn)使用的經(jīng)典技術(shù)，即把已淘汰的分頁(yè)復(fù)制到磁盤上的一個(gè)交換空間。在這里，研究者的做法是把已淘汰塊復(fù)制到 CPU 內(nèi)存。如圖 4 所示，除了 GPU 塊分配器之外，vLLM 還包括一個(gè) CPU 塊分配器，用以管理交換到 CPU RAM 的物理塊。
• 重新計(jì)算。這個(gè)做法很簡(jiǎn)單，就是當(dāng)被搶占的序列被重新調(diào)度時(shí)，直接重新計(jì)算 KV 緩存。請(qǐng)注意，重新計(jì)算的延遲可能比原始延遲低得多。

交換和重新計(jì)算的性能取決于 CPU RAM 和 GPU 內(nèi)存之間的帶寬以及 GPU 的計(jì)算能力。

分布式執(zhí)行

vLLM 能有效用于分布式的硬件設(shè)置，因?yàn)槠渲С?Transformers 上廣泛使用的 Megatron-LM 式張量模型并行化策略。該策略遵循 SPMD（單程序多數(shù)據(jù)）執(zhí)行調(diào)度方案，其中線性層會(huì)被分開以執(zhí)行逐塊的矩陣乘法，并且 GPU 會(huì)通過 allreduce 操作不斷同步中間結(jié)果。

具體來說，注意力算子會(huì)被分散在注意力頭維度上，每個(gè) SPMD 過程負(fù)責(zé)多頭注意力中一個(gè)注意力頭子集。

實(shí)現(xiàn)

vLLM 是一個(gè)端到端服務(wù)提供系統(tǒng)，具有 FastAPI 前端和基于 GPU 的推理引擎。

該前端擴(kuò)展了 OpenAI API 接口，允許用戶為每個(gè)請(qǐng)求自定義采樣參數(shù)，例如最大序列長(zhǎng)度和波束寬度。

vLLM 引擎由 8500 行 Python 代碼和 2000 行 C++/CUDA 代碼寫成。

該團(tuán)隊(duì)也用 Python 開發(fā)了一些與控制相關(guān)的組件，包括調(diào)度器和塊管理器，還為 PagedAttention 等關(guān)鍵操作開發(fā)了定制版 CUDA 核。

至于模型執(zhí)行器，他們使用 PyTorch 和 Transformer 實(shí)現(xiàn)了常用的 LLM，比如 GPT、OPT 和 LLaMA。

他們?yōu)榉植际?GPU 工作器之間的張量通信使用了 NCCL。

評(píng)估

為了驗(yàn)證 vLLM 的效果，該團(tuán)隊(duì)進(jìn)行了一系列實(shí)驗(yàn)評(píng)估。表 1 給出了模型大小和服務(wù)器配置情況。

基礎(chǔ)采樣

在 ShareGPT 數(shù)據(jù)集上，vLLM 可以維持比 Orca (Oracle) 高 1.7-2.7 倍的請(qǐng)求率，比 Orca (Max) 則高 2.7-8 倍，同時(shí)還能維持相近的延遲，如圖 12 上面一行和圖 13a 所示。

在 Alpaca 數(shù)據(jù)集上的結(jié)果也類似。

并行采樣和波束搜索

如圖 14 上面一行所示，如果要采樣的序列數(shù)量很多，則 vLLM 能在 Orca 基準(zhǔn)的基礎(chǔ)上帶來更大的提升。

圖中下面一行則展示了不同波束寬度的波束搜索結(jié)果。由于波束搜索支持更多共享，所以 vLLM 帶來的性能優(yōu)勢(shì)還要更大。

圖 15 展示了內(nèi)存節(jié)省量，計(jì)算方法是用因?yàn)楣蚕矶?jié)省的塊數(shù)除以不使用共享時(shí)的總塊數(shù)。

共享前綴

如圖 16a 所示，當(dāng)共享單樣本前綴時(shí)，vLLM 的吞吐量比 Orca (Oracle) 高 1.67 倍。此外，當(dāng)共享更多樣本時(shí)（圖 16b），vLLM 能實(shí)現(xiàn)比 Orca (Oracle) 高 3.58 倍的吞吐量。

聊天機(jī)器人

從圖 17 可以看出，相比于三個(gè) Orca 基準(zhǔn)模型，vLLM 可以維持高 2 倍的請(qǐng)求率。

該團(tuán)隊(duì)也進(jìn)行了消融實(shí)驗(yàn)，表明 PagedAttention 確實(shí)有助于 vLLM 提升效率。他們也研究了塊大小與重新計(jì)算和交換技術(shù)的影響。詳情參閱原論文。