6.7k Star量的vLLM出論文了，讓每個人都能輕松快速低成本地部署LLM服務

作者：機器之心 2023-09-25 12:01:42

現在離 vLLM 最初宣布時差不多已過去三個月，他們終于發布了這篇關于 vLLM 和 PagedAttention 的研究論文，其中詳細解釋了他們如何通過類似操作系統虛擬內存管理的機制來實現高效的 LLM 服務。

今年六月，來自加州大學伯克利分校等機構的一個研究團隊開源了 vLLM（目前已有 6700 多個 star），其使用了一種新設計的注意力算法 PagedAttention，可讓服務提供商輕松、快速且低成本地發布 LLM 服務。

在當時的博客文章中，該團隊宣稱 vLLM 能實現比 HuggingFace Transformers 高 24 倍的吞吐量！

八月底，該項目還與其它另外 7 個開源項目一道獲得了 a16z 開源人工智能資助計劃的首輪資助。

現在離 vLLM 最初宣布時差不多已過去三個月，他們終于發布了這篇關于 vLLM 和 PagedAttention 的研究論文，其中詳細解釋了他們如何通過類似操作系統虛擬內存管理的機制來實現高效的 LLM 服務。該論文已被將于十月底在德國舉辦的 ACM 操作系統原理研討會（SOSP 2023）接收。

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

GPT 和 PaLM 等大型語言模型（LLM）的出現催生出了開始對我們的工作和日常生活產生重大影響的新應用，比如編程助理和通用型聊天機器人。

許多云計算公司正競相以托管服務的方式提供這些應用。但是，運行這些應用的成本非常高，需要大量硬件加速器，如 GPU。根據最近的估計，相比于傳統的關鍵詞查詢方法，處理一個 LLM 請求的成本超過其 10 倍以上。考慮到成本如此之高，提高 LLM 服務系統的吞吐量（并由此降低單位請求的成本）就變得更為重要了。

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

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

舉個例子，圖 1（左）展示了一個 130 億參數的 LLM 在一臺 40GB RAM 的英偉達 A100 GPU 上的內存分布。

其 65% 的內存都分配給了模型權重，而模型權重在提供服務期間是不會變化的。

30% 的內存是用于存儲請求的動態狀態。對 Transformer 而言，這些狀態由與注意力機制關聯的鍵（key）和值（value）張量構成，通常被稱為 KV 緩存，其表示用于生成序列中新輸出 token 的之前 token 上下文。

其余占比很小的內存則是用于其它數據，包括激活 —— 評估 LLM 時創建的臨時張量。

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

來自 UC 伯克利等機構的這個研究團隊在論文中表示，他們觀察到當前的 LLM 服務系統都沒有高效地管理 KV 緩存內存。主要原因是它們會將請求的 KV 緩存保存在鄰接的內存空間中，因為大多數深度學習框架都需要將張量存儲在相鄰連續的內存中。

但是，不同于傳統深度學習工作負載中的張量，KV 緩存有其自己的獨特性質：它會在模型生成新 token 的過程中隨時間動態地增長和縮小，而且它的持續時間和長度是無法事先知曉的。

這些特性使得現有系統的方法極其低效，這主要體現在兩方面：

第一，現有系統存在內部和外部內存碎片的問題。

為了將請求的 KV 緩存存儲在相鄰連續的空間中，現有系統會預先分配一塊限制了請求最大長度的相鄰連續內存。這會導致嚴重的內部碎片化，因為請求的實際長度可能比其最大長度短很多。

此外，就算能事先知道實際長度，預先分配內存的方法依然很低效。因為那一整塊內存在請求的生命周期中都被占用了，其它更短的請求無法使用這一塊內存 —— 盡管可能其中很大一部分都未被使用。

此外，外部內存碎片可能也很顯著，因為每個請求預先分配的內存大小可能都不同。

事實上，我們可以在圖 2 的分析中看到，在現有系統中，僅有 20.4%-38.2% 的 KV 緩存內存會用于存儲實際的 token 狀態。

第二，現有系統無法利用內存共享的機會。

LLM 服務通常會使用先進的解碼算法，例如并行采樣和波束搜索，這些方法可為每個請求生成多個輸出。在這些場景中，由多個序列組成的請求可以部分共享它們的 KV 緩存。但是，現有系統不可能使用內存共享，因為不同序列的 KV 緩存存儲在分開的相鄰連續空間中。

為了解決上述限制，該團隊提出了一種注意力算法 PagedAttention；該算法的靈感來自操作系統（OS）解決內存碎片化和內存共享的方案：使用分頁機制的虛擬內存。

PagedAttention 會將請求的 KV 緩存分成一塊塊的，每一塊（block）都包含一定數量 token 的注意力鍵和值。在 PagedAttention 中，KV 緩存的塊不一定要存儲在相鄰連續空間中。

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

此外，它還能消除外部碎片，因為所有塊的大小都相同。

最后，它還能實現以塊為粒度的內存共享；這種內存共享支持與同一請求相關的不同序列，甚至也支持不同的請求。

立足于 PagedAttention，該團隊構建了一個高吞吐量的分布式 LLM 服務引擎 vLLM，其幾乎做到了 KV 緩存內存的零浪費。

vLLM 使用了塊級的內存管理和搶占式的請求調度 —— 這些機制都是配合 PagedAttention 一起設計的。vLLM 支持 GPT、OPT 和 LLaMA 等各種大小的常用 LLM，包括那些超出單個 GPU 內存容量的 LLM。

研究者基于多種模型和工作負載進行了實驗評估，結果表明：相比于當前最佳的系統，vLLM 能在完全不影響準確度的前提下將 LLM 服務吞吐量提升 2-4 倍。而且當序列更長、模型更大、解碼算法更復雜時，vLLM 帶來的提升還會更加顯著。

總體而言，該研究的貢獻如下：

確認了提供 LLM 服務時的內存分配挑戰并量化了它們對服務性能的影響。
提出了 PagedAttention，這是一種對存儲在非相鄰連續的分頁內存中的 KV 緩存進行操作的注意力算法，其靈感來自操作系統中的虛擬內存和分頁。
設計并實現了 vLLM，這是一個基于 PagedAttention 構建的分布式 LLM 服務引擎。
通過多種場景對 vLLM 進行了實驗評估，結果表明其表現顯著優于 FasterTransformer 和 Orca 等之前最佳方案。

方法概況

vLLM 的架構如圖 4 所示。

vLLM 采用一種集中式調度器（scheduler）來協調分布式 GPU 工作器（worker）的執行。KV 緩存管理器由 PagedAttention 驅動，能以分頁方式有效管理 KV 緩存。具體來說，KV 緩存管理器通過集中式調度器發送的指令來管理 GPU 工作器上的物理 KV 緩存內存。

PagedAttention

不同于傳統的注意力算法，PagedAttention 支持將連續的鍵和值存儲在非相鄰連續的內存空間中。

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

其中 A_{i,j} 是在第 j 個 KV 塊上的注意力分數的行向量。

在注意力計算期間，PagedAttention 核會分開識別并獲取不同的 KV 塊。

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

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

綜上所述，PagedAttention 算法能讓 KV 塊存儲在非相鄰連續的物理內存中，從而讓 vLLM 實現更為靈活的分頁內存管理。

KV 緩存管理器

使用 PagedAttention，該團隊將 KV 緩存組織為固定大小的 KV 塊，就像虛擬內存中的分頁。

對 KV 緩存的請求會被表示成一系列邏輯 KV 塊，在生成新 token 和它們的 KV 緩存時從左向右填充。最后一個 KV 塊中未填充的位置留給未來填充。

使用 PagedAttention 和 vLLM 進行解碼

圖 6 通過一個示例展示了 vLLM 在對單個輸入序列的解碼過程中執行 PagedAttention 和管理內存的方式。

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

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

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

圖 7 給出了 vLLM 管理兩個序列的內存的示例。

應用于其它解碼場景

在許多成功的 LLM 應用中，LLM 服務必須要能提供更為復雜的解碼場景 —— 有復雜的訪問模式和更多的內存共享機會。研究者展示了 vLLM 的一般適用性。他們在論文中討論過的場景包括并行采樣、波束搜索、共享前綴、混合解碼方法。

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

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

調度和搶占

LLM 面臨著一個獨特挑戰：輸入 LLM 的 prompt 的長度變化范圍很大，并且輸出結果的長度也無法事先知曉，而要依輸入 prompt 和模型的情況而定。隨著請求和相應輸出數量的增長，vLLM 可能會為了存儲新生成的 KV 緩存而耗盡 GPU 的物理塊。

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

通常來說，淘汰策略是使用啟發式方法來預測將在最遠的未來訪問的塊然后淘汰那個塊。由于這里已知是一起訪問一個序列的所有塊，所以該團隊實現了一種全清或不動的淘汰策略，即要么淘汰序列的所有塊，要么就全不淘汰。此外，一個請求中的多個序列會被組成一個序列組來一起調度。一個序列組中的序列總是會被一起搶占或重新調度，因為這些序列之間存在潛在的內存共享。

為了解答第二個有關如何恢復已淘汰塊的問題，研究者考慮了兩種技術：

交換。這是大多數虛擬內存實現使用的經典技術，即把已淘汰的分頁復制到磁盤上的一個交換空間。在這里，研究者的做法是把已淘汰塊復制到 CPU 內存。如圖 4 所示，除了 GPU 塊分配器之外，vLLM 還包括一個 CPU 塊分配器，用以管理交換到 CPU RAM 的物理塊。
重新計算。這個做法很簡單，就是當被搶占的序列被重新調度時，直接重新計算 KV 緩存。請注意，重新計算的延遲可能比原始延遲低得多。

交換和重新計算的性能取決于 CPU RAM 和 GPU 內存之間的帶寬以及 GPU 的計算能力。

分布式執行

vLLM 能有效用于分布式的硬件設置，因為其支持 Transformers 上廣泛使用的 Megatron-LM 式張量模型并行化策略。該策略遵循 SPMD（單程序多數據）執行調度方案，其中線性層會被分開以執行逐塊的矩陣乘法，并且 GPU 會通過 allreduce 操作不斷同步中間結果。

具體來說，注意力算子會被分散在注意力頭維度上，每個 SPMD 過程負責多頭注意力中一個注意力頭子集。