從頭講解 vLLM 推理加速原理

大型語言模型（LLM）在自然語言處理領域展現出了強大的能力，但同時也帶來了巨大的計算挑戰，尤其是在推理階段。這些模型通常包含數十億甚至數千億個參數，在生成文本時需要進行大量的計算操作。爲了提高推理效率，研究人員開發了各種優化技術，其中 vLLM 推理加速技術是一種效果顯著的方法。vLLM 通過綜合運用多種優化策略，顯著提高了大型語言模型的推理速度和資源利用率。本文將詳細探討 vLLM 推理加速的核心原理，包括自迴歸模型的基本概念、注意力機制中的 Q、K、V，以及關鍵的 KV 緩存技術，並深入分析 vLLM 所採用的各種優化策略。

一. 自迴歸模型基礎

自迴歸模型是大型語言模型的基礎架構之一。這類模型的核心特徵是逐個生成序列中的元素，每個新元素的生成都依賴於之前已生成的所有元素。在語言模型中，這些元素通常是詞元（tokens）。

自迴歸過程可以描述如下：

1. 模型接收初始輸入序列。

2. 基於當前序列，預測下一個最可能的詞元。

3. 將預測的詞元添加到序列末尾。

4. 重複步驟 2 和 3，直到達到預定的終止條件（如生成特定的結束標記或達到最大長度）。

這個過程雖然直觀，但在處理長序列時會面臨嚴重的性能瓶頸。因爲對於序列中的每個新位置，模型都需要處理從開始到當前位置的所有信息。這就引出了優化的需求，而理解注意力機制是優化的關鍵。

二、 注意力機制

注意力機制，特別是自注意力（Self-Attention）機制，是現代大型語言模型的核心組成部分。它允許模型在處理序列數據時，動態地關注不同部分的信息。注意力機制的核心概念是 Query（查詢）、Key（鍵）和 Value（值）。

• Query (Q)：當前單詞的一種表示，用於對所有其他單詞進行評分（使用它們的 Key）。我們只關心當前正在處理的 token 的 Query 。在生成任務中，通常是最後一個 token 的表示。

• Key (K)：是該句段中所有單詞的標籤。它們是我們在搜索相關單詞時所匹配的內容。用於與 Query 進行匹配，決定應該關注哪些信息。

• Value (V)：是實際的單詞表示，一旦我們對每個單詞的相關性進行了評分，這些就是我們聚合起來用來表示當前單詞的值。

1. 工作原理類比：

假設輸入爲 “A robot must obey the orders given it by human beings except where such orders would conflict with the First Law.”

Query (查詢)：你手裏拿着一張便利貼, 上面寫着 "it"。這就像是當前詞的查詢向量。在我們的例子中, 這是第 9 個位置的查詢向量。

Key (鍵)：文件櫃中每個文件夾的標籤就像是鍵向量。它們代表了序列中每個詞的 " 標識 "。在圖中, 我們看到了 4 個文件夾, 分別標記爲 "a", "robot", "must", "obey"。

Value (值)：每個文件夾裏的實際內容對應着值向量, 包含了詞的具體信息。

2. 注意力計算過程：

你拿着寫有 "it" 的便利貼 (Query #9), 與每個文件夾的標籤 (Key #1 到#4) 進行比較。

比較的結果決定了你對每個文件夾內容的關注程度。

在第二張圖中, 我們可以看到比較的結果:

"a" 文件夾得到 30% 的關注

"robot" 文件夾得到 50% 的關注

"must" 和 "obey" 文件夾各得到 0.1% 的關注

這些百分比實際上就是注意力權重, 決定了你從每個文件夾中提取信息的比例。

信息綜合：

最後, 你會根據這些權重從各個文件夾中提取信息, 並將它們綜合起來, 形成對 "it" 這個詞的理解。

最終的表示是多個信息源的加權組合。這個加權組合可以用一個簡單的形式表達：

it 的表示 = 0.3a + 0.5robot + 0.001must + 0.001obey

在這個表達式中，每個詞前的係數代表了注意力權重，而詞本身代表了其 Value 向量。這種加權求和的方式使得模型能夠根據當前上下文的需求，靈活地整合來自不同位置的信息，從而形成對當前詞 "it" 的理解。這個簡化的表達雖然忽略了很多細節，但基本表達了注意力機制中信息綜合的核心思想。

3. 注意力機制的本質：

Query 向量代表當前 token 的 " 問題 " 或 " 需求 "。

Key 向量代表每個 token 可能提供的 " 信息 "。

Value 向量是實際傳遞的 " 信息內容 "。

在生成新 token 時，我們需要新的 " 問題 "（Query）來查詢所有歷史 " 信息 "（Key）並獲取相關的 " 內容 "（Value）。

注意力機制允許模型動態地 " 查閱 " 之前的信息。不同的信息源 (早先的詞) 會根據其相關性獲得不同程度的 " 注意力 "。最終的表示是多個信息源的加權組合。

三. KV 緩存：自迴歸模型的關鍵優化

1. 爲什麼需要 KV Cache:

1. 計算效率問題:

在自迴歸生成過程中, 模型需要爲每個新 token 重新計算整個序列的 Key 和 Value。這意味着, 隨着生成的 token 數量增加, 計算量呈二次方增長。例如, 在生成第 20 個 token 時, 我們需要計算 20 次 Key 和 Value; 而在生成第 21 個 token 時, 我們又需要重新計算 21 次。這種重複計算極大地降低了推理效率。

2. 資源消耗:

反覆計算不僅耗時, 還會佔用大量的計算資源, 尤其是在處理長文本時。

3. 延遲問題:

在實時對話或文本生成應用中, 用戶體驗很大程度上取決於系統的響應速度。頻繁的重複計算會導致明顯的延遲。

2.KV Cache 的原理:

KV Cache 的核心思想是緩存並重用之前計算過的 Key 和 Value, 從而避免重複計算。

工作原理:

初始計算:

對於輸入序列 "A robot must obey", 我們首先計算每個 token 的 Key 和 Value。

緩存存儲:

將計算得到的 Key 和 Value 存儲在緩存中。例如:

• "A" 的 Key 和 Value 存儲在位置 1

• "robot" 的 Key 和 Value 存儲在位置 2

• "must" 的 Key 和 Value 存儲在位置 3

• "obey" 的 Key 和 Value 存儲在位置 4

生成新 token:

當模型需要生成下一個 token "the" 時:

a) 使用緩存中的 Key 和 Value, 無需重新計算。

b) 只需爲新 token "the" 計算 Query。

c) 用這個新 Query 與緩存中的所有 Key 計算注意力分數。

d) 基於注意力分數和緩存的 Value 生成新的表示。

更新緩存:

生成 "the" 後, 計算它的 Key 和 Value, 並將其添加到緩存的位置 5。

重複過程:

對於後續的每個新 token (如 "orders", "given", "it" 等), 重複步驟 3-4。

實例說明:

以生成 "it" (序列中的第 9 個 token) 爲例:

緩存狀態:

此時緩存中已有 8 個位置, 存儲了 "A robot must obey the orders given" 的 Key 和 Value。

計算過程:

只需爲 "it" 計算新的 Query。

這個 Query 與緩存中的 8 個 Key 計算注意力分數。

使用這些分數和緩存的 Value 計算 "it" 的表示。

緩存更新:

計算 "it" 的 Key 和 Value, 添加到緩存的第 9 個位置。

3.KV Cache 的優勢:

1) 顯著減少計算量:

對於長度爲 N 的序列, 傳統方法需要 O(N^2) 的計算量, 而使用 KV Cache 後僅需 O(N)。

2. 加速推理:

由於避免了重複計算, 推理速度大幅提升。

3. 內存權衡:

雖然需要額外的內存來存儲緩存, 但相比計算效率的提升, 這是值得的權衡。

4) 適合長文本生成:

對於長文本生成任務,KV Cache 的優勢更加明顯。

四、vLLM 的核心技術：PagedAttention

在介紹 vLLM 的核心技術之前，我們首先需要理解傳統 LLM 服務系統面臨的主要挑戰。這些挑戰正是 vLLM 試圖解決的問題，也是理解 PagedAttention 重要性的關鍵。

1. 傳統 LLM 服務的效率瓶頸

傳統的 LLM 服務系統在處理大規模語言模型時面臨以下幾個主要問題：

a) 內存碎片化

• 內部碎片：系統往往爲每個請求預分配最大可能長度的連續內存，導致大量內存被浪費。

• 外部碎片：不同長度的請求導致內存中出現難以利用的小塊空閒空間。

b) 內存利用率低

由於預分配和碎片化問題，大量 GPU 內存被浪費，限制了系統能夠同時處理的請求數量。

c) 缺乏靈活性

將整個 KV 緩存存儲在連續內存空間中，難以適應動態變化的工作負載。

d) 內存共享困難

在並行採樣或束搜索等場景中，難以實現高效的內存共享。

e) 計算冗餘

無法有效共享和重用計算結果，導致大量重複計算。

f) 批處理效率低下

靜態批處理策略難以在延遲和吞吐量之間取得良好平衡。

g) 調度不靈活

粗粒度的調度策略難以根據實時工作負載動態調整資源分配。

h) 擴展性限制

在模型並行和數據並行方面的支持不夠優化，限制了系統的擴展性。

i) 長序列處理效率低

隨着序列長度增加，性能急劇下降。 

j) 動態長度處理能力差

難以高效處理動態變化的輸入和輸出序列長度。

2. PagedAttention: vLLM 的核心創新

加州大學伯克利分校的研究團隊開發了 vLLM，一個開源的 LLM 推理和服務庫。vLLM 的核心創新在於其獨特的注意力算法——PagedAttention，它徹底改變了 LLM 服務的效率和性能。

PagedAttention 的靈感來源於操作系統中的虛擬內存和分頁概念。傳統的注意力算法要求將鍵 (key) 和值 (value) 張量連續存儲在內存中，而 PagedAttention 允許將它們存儲在非連續的內存空間中。簡而言之，PagedAttention 背後的想法是創建映射到 GPU 內存中的物理塊的連續虛擬塊。每個塊都旨在存儲預定義數量的標記的鍵值對張量。所有塊都是虛擬連續的，並映射到碎片化 GPU 內存中在推理期間按需分配的物理非連續塊。內存中還會創建一個簡單的索引表，以將虛擬塊與物理塊關聯起來。PagedAttention 的內核會根據需要獲取這些塊。由於塊的大小有限，系統會獲取較少數量的鍵值張量，因此這種方法非常高效。

PagedAttention 的工作原理可以通過以下幾個關鍵方面來理解：

1. 非連續內存訪問和查詢機制 

PagedAttention 允許從非連續的內存位置高效地檢索和使用相關信息。

1). Query vector（查詢向量）：

   圖中左側綠色方框中的 "for" 是當前的查詢向量。在注意力機制中，查詢向量用於與 key vectors 進行比較，以確定應該關注哪些信息。

2). Key and value vectors（鍵值向量）：

   右側的表格代表了存儲在內存中的 key 和 value 向量。這些向量被分成了多個塊（Block 0, Block 1, Block 2）。

3). 內存塊（Blocks）：

   - Block 0 包含 "Alan Turing is a"

   - Block 1 包含 "computer scientist and mathematician"

   - Block 2 包含 "renowned for"

4). 非連續內存訪問：

   圖中的箭頭顯示了查詢向量 "for" 如何同時訪問這三個非連續的內存塊。這正是 PagedAttention 的核心創新——它允許從非連續的內存位置高效地檢索和使用相關信息。 

5). 靈活的內存管理：

   這種設計使得系統可以更靈活地管理內存，不需要所有相關的 key 和 value 向量都連續存儲在一起。

6). 高效計算：

   儘管內存塊是非連續的，PagedAttention 算法仍然能夠高效地計算注意力分數，將查詢向量與所有相關的 key 向量進行比較。

這種方法的優勢在於它可以更有效地利用 GPU 內存，減少內存碎片，並允許更靈活的內存分配和管理策略。這對於處理變長序列和批量處理多個請求特別有利，從而顯著提高了大型語言模型的服務效率和吞吐量。

2. 動態內存管理和生成過程

1). 初始狀態:

   - 提示詞是 "Alan Turing is a computer scientist"

   - 邏輯 KV 緩存塊、物理 KV 緩存塊和塊表都是空的

2). 分配空間並存儲提示詞的 KV 緩存:

   - 邏輯塊 0 和 1 被填充了提示詞的 tokens

   - 物理塊 7 和 1 被分配並填充相應的數據

   - 塊表記錄了邏輯塊到物理塊的映射關係

3). 生成第一個 token "and":

   - "and" 被添加到邏輯塊 1

   - 物理塊 1 中相應位置也被填充

   - 塊表中填充槽數量增加

4). 生成第二個 token "mathematician":

   - "mathematician" 被添加到邏輯塊 1

   - 物理塊 1 被完全填滿

5). 生成第三個 token "renowned":

   - 由於邏輯塊 1 已滿, 新的邏輯塊 2 被創建

   - 物理塊 3 被按需分配

   - "renowned" 被添加到新的邏輯和物理塊中

6). 生成第四個 token "for":

   - "for" 被添加到邏輯塊 2

   - 物理塊 3 中相應位置也被填充

這個過程展示了 PagedAttention 如何靈活地管理內存, 通過將邏輯塊映射到非連續的物理塊, 實現了高效的內存利用和動態擴展。

3. 內存共享

1). 輸入 prompt：

   "The future of artificial intelligence is" 作爲輸入提示。

2). 共享 prompt 處理：

   - LLM（大語言模型）首先處理這個 prompt。

   - prompt 的處理結果（key 和 value 向量）被存儲在 KV 緩存中。

   - 這部分內存是共享的，意味着所有後續的並行輸出都可以使用這些相同的計算結果。

3). 多個輸出：

   - LLM 基於相同的 prompt 生成多個不同的輸出。

   - 每個輸出代表一個不同的續寫或完成方案。

4. 內存共享實現機制演示

1). 初始狀態，序列 A 和 B 共享相同的初始提示："The future of artificial intelligence is"。物理 KV 緩存塊存儲了這個共享提示，兩個序列的邏輯 KV 緩存塊都指向相同的物理塊，引用計數爲 2。

2). 序列 A 生成了第一個新 token "likely"。由於這是一個寫操作，觸發了 Copy-on-Write 機制。系統創建了一個新的物理塊來存儲 "intelligence is likely"。序列 A 的邏輯塊現在指向這個新的物理塊。

3).Copy-on-Write 過程的細節。新的物理塊被創建，包含 "intelligence is" 和新生成的 "likely"。序列 A 的邏輯塊更新以指向這個新塊。原始物理塊的引用計數從 2 減少到 1。

4). 圖 3Copy-on-Write 過程完成。序列 A 的邏輯塊完全指向新的物理塊，包含 "intelligence is likely"。原始物理塊仍然被序列 B 引用。

5). 序列 B 生成了它的第一個新 token "a"。由於這個 token 與序列 A 不同，不需要複製共享的物理塊。序列 B 在自己的邏輯塊中添加了新 token "a"。

6). 兩個序列繼續生成新的 token。序列 A 生成 "to"，序列 B 生成 "subject"。每個序列在自己的邏輯塊中繼續添加新的 token。

7). 序列繼續生成。序列 A 生成 "be"，序列 B 生成 "of"。新的物理塊被創建來存儲這些不同的 token。

8). 最後一步。序列 A 生成 "profound"，序列 B 生成 "great"。每個序列繼續在各自的邏輯塊中添加新 token。

這個過程展示了 PagedAttention 如何高效地管理內存，通過共享初始內容和僅在必要時創建新的物理塊來最大化內存利用率，同時允許不同序列獨立生成內容。

四、vLLM 其它優化技術

• 連續批處理：傳入的請求被連續批處理在一起，以最大限度地提高硬件利用率並減少計算浪費，最大限度地減少空閒時間。

• 量化：vLLM 利用 FP16 等量化技術，通過以較低的精度表示 KV 緩存來優化內存使用，從而減少內存佔用並加快計算速度。

• 優化的 CUDA 內核：vLLM 手動調整在 GPU 上執行的代碼以實現最佳性能。例如，對於融合重塑和塊寫入，開發了優化的內核，將新的 KV 緩存拆分爲塊，重塑它們以實現高效的內存訪問，並根據塊表保存它們，所有這些都融合到單個內核中以減少開銷。

五、總結

vLLM 解決了 LLM 部署中的一個關鍵瓶頸：推理和服務效率低下。使用創新的 PagedAttention 技術，vLLM 優化了核心注意操作期間的內存使用，從而顯著提高了性能。這意味着更快的推理速度和在資源受限的硬件上運行 LLM 的能力。

參考及附圖來源：

1.https://jalammar.github.io/illustrated-gpt2/

2.https://blog.vllm.ai/2023/06/20/vllm.html

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/FFcZ1c_a3Ua0vLIj3DGaCQ