如何在LLM中使用KV快取？

LLM 最近很熱門，但瞭解這些模型背後的工作原理總是很有意思的。可能有些人還不瞭解，LLM 自 2017 年著名的論文《注意力機制就是一切》（Attention is all you need）發表以來就一直在發展。但早期的基於 Transformer 的模型由於內部數學運算繁重，需要大量的記憶體，因此存在不少缺陷。

隨著 LLM 生成的文字越來越多，GPU 記憶體消耗也會越來越高。當達到一定程度時，GPU 會出現記憶體溢位（Out of Memory）問題，導致整個程式崩潰，LLM 也無法繼續生成文字。鍵值快取（Key-Value Cacheing）是一種可以緩解這個問題的技術。它本質上是記住之前步驟中的重要資訊。模型無需從頭開始重新計算所有內容，而是重用已計算的內容，從而大大提高文字生成速度和效率。這項技術已被應用於多個模型，例如 Mistral、Llama 2 和 Llama 3 模型。

那麼，讓我們來了解一下為什麼 KV 快取對 LLM 如此重要。

注意力機制

我們通常會生成三個不同的權重矩陣（W_q、W_k、W_v）來生成 Q、K 和 V 向量。這些權重矩陣源自資料。

我們可以將 Q 視為一個向量，將 K 和 V 視為二維矩陣。這是因為 K 和 V 分別儲存每個先前詞元的向量，堆疊起來就形成了矩陣。

Q 向量表示解碼器步驟中輸入的新詞元。

同時，K 矩陣表示新詞元可以查詢的所有先前詞元的資訊或“鍵”，以確定其相關性。

每當輸入一個新的查詢向量時，它都會將其自身與所有鍵向量進行比較。它本質上是找出哪些先前詞元最為重要。

這種相關性透過加權平均值來表示，即軟注意力得分。

V 矩陣表示每個先前詞元的內容/含義。注意力得分計算完成後，會用於對 V 向量進行加權求和。最終的上下文輸出將作為模型後續步驟的依據。

簡單來說，K 矩陣決定關注哪些內容，而 V 矩陣則決定從中提取哪些資訊。

預填充和解碼過程中的注意力機制

Source: Doubleworld

問題是什麼？

在典型的因果 Transformer 模型中，由於採用了自迴歸解碼，我們每次生成一個詞，前提是我們擁有所有先前的上下文資訊。隨著我們不斷生成新的詞元，K 矩陣和 V 矩陣會不斷更新。一旦計算出該詞元的嵌入向量，其對應的詞元值就不會再改變。但是，模型需要在每個步驟中對該詞元對應的 K 矩陣和 V 矩陣進行大量的計算。這會導致矩陣乘法運算次數呈平方級增長。這是一個非常繁重且耗時的任務。

為什麼我們需要鍵值快取？

讓我們透過一個實際例子來理解——考試準備場景

我們一週後就要期末考試了，需要在7天內學習20個模組。我們從第一個模組開始，可以把它看作是第一個詞元。學完之後，我們開始學習第二個模組，並自然而然地嘗試將它與之前學過的內容聯絡起來。這個新模組（第二個詞元）與之前模組的互動，類似於序列中新詞元與之前詞元的互動。

這種關係對於理解至關重要。當我們繼續學習更多模組時，我們仍然記得之前模組的內容，我們的記憶就像一個快取。因此，在學習一個新模組時，我們只需要思考它如何與之前的模組聯絡起來，而無需每次都從頭開始學習所有內容。

這就是鍵值快取的核心思想。鍵值對從一開始就儲存在記憶體中，每當一個新的詞元到來時，我們只需要計算它與已儲存上下文的互動。

鍵值快取的作用

正如您所見，鍵值快取節省了大量的計算資源，因為所有之前的鍵值對（K 和 V）都儲存在記憶體中。因此，當一個新的詞元到達時，模型只需要計算這個新詞元如何與已儲存的詞元互動。這避免了從頭開始重新計算所有內容。

這樣一來，模型的瓶頸就從計算密集型轉變為記憶體密集型，視訊記憶體成為主要的限制因素。在傳統的多頭自注意力機制中，每個新詞元都需要進行完整的張量運算，其中 Q、K 和 V 都與輸入相乘，導致計算複雜度為二次方。但使用鍵值快取，我們只需對儲存的 K 和 V 矩陣進行一次行（向量-矩陣）乘法運算即可獲得注意力分數。這大大減少了計算量，使得該過程更多地受限於記憶體頻寬，而不是原始的 GPU 計算。

鍵值快取

鍵值快取的核心思想很簡單：我們能否避免對模型在推理過程中已經處理過的詞元重複計算？

Source: YouTube (Tensordroid)

在上圖中，我們可以看到紅色區域是冗餘的，因為我們只關心 Q-K(T) 中下一行的計算。

這裡，K 對應於鍵矩陣，V 對應於值矩陣。

在因果轉換器中計算下一個詞元時，我們只需要最近的詞元來預測下一個詞元。這是下一個詞預測的基礎。由於我們只查詢第 n 個詞元來生成第 n+1 個詞元，因此我們不需要舊的 Q 值，所以不會儲存它們。但在標準的多頭注意力機制中，由於我們沒有快取鍵和值，所有詞元的 Q、K 和 V 都會被重複計算。

Source: YouTube (Tensordroid)

在上圖中，K 和 V 代表了之前所有詞元累積的上下文資訊。在 KV 快取出現之前，模型在每次解碼步驟中都會重新計算從頭到尾所有詞元的 K 和 V 矩陣。這也意味著需要反覆重新計算所有軟注意力分數和最終注意力輸出，計算量非常大。然而，有了 KV 快取，我們儲存了序列長度減 1 之前的所有 K 向量，並對從 1 到 n-1 的 V 矩陣執行相同的操作。

由於我們只需要查詢最新的詞元，因此不再需要重新計算之前的詞元。第 n 個詞元已經包含了生成第 (n+1) 個詞元所需的所有資訊。

KV快取：核心思想

KV 快取的核心思想由此開始……

每當一個新的詞元到達時，我們計算其新的 K 向量，並將其作為附加列新增到現有的 K 矩陣中，形成 K’。類似地，我們計算最新標記的 V 向量，並將其作為新行新增到 V’ 中。這是一個輕量級的張量操作。

示例

如果這讓你感到困惑，讓我們用一個簡單的示例來解釋一下。

如果沒有 KV 快取，模型每次都會重新計算整個 K 和 V 矩陣。這會導致一次完整的二次方規模的注意力計算，以生成軟注意力分數，然後將其乘以整個 V 矩陣。這在記憶體和計算方面都非常耗費資源。

有了 KV 快取，我們只使用最新標記的查詢向量 (q_new)。將 q_new 與更新後的 K 矩陣相乘，該矩陣包含 K_prev（已快取）和 K_new（剛剛計算）。同樣的邏輯也適用於 V_prev，它會獲得一個新的 V_new 行。這透過將大型矩陣乘法轉換為小得多的向量乘法，大大減少了計算量。生成的注意力分數被新增到注意力矩陣中，並在 Transformer 流程​​中正常使用。

然後，這第四個注意力值（在本例中）用於計算下一個標記。使用鍵值快取（KV 快取）時，我們不會儲存或重新計算之前的 Q 值，但會儲存所有之前的 K 向量。新的 Q 值乘以 Kᵀ 得到 Q·Kᵀ 向量（軟注意力分數）。這些分數乘以 V 矩陣，生成儲存上下文的加權平均值，該平均值即為用於預測的第四個標記的輸出。

注意力僅依賴於前面的標記。此操作依賴於 GPU 的視訊記憶體（VRAM），因此我們需要儲存這些 K 和 V 矩陣。因此，實際上，這就是我們記憶體容量的上限。

在 Llama 2 等研究論文中也多次提到了 KV 頭，因為我們知道 Transformer 模型有多個層，每個層有多個頭。每個頭都會應用 KV 快取，因此得名 KV 頭。

鍵值快取的挑戰

鍵值快取面臨的挑戰包括：

• GPU 利用率低：儘管 GPU 效能強大，但鍵值快取通常僅使用 20-40% 的 GPU 視訊記憶體，這主要是由於記憶體分配方式無法充分利用。
• 連續記憶體需求：鍵值快取塊必須儲存在連續的記憶體中，這使得記憶體分配更加嚴格，並導致許多小的空閒空間無法使用。

即使 GPU 非常擅長處理快速迭代操作，但使用鍵值快取時，GPU 記憶體也無法得到充分利用。問題在於記憶體的劃分和預留方式，即如何用於生成令牌。

內部碎片

內部碎片是指分配的記憶體超過實際使用的記憶體量。例如，如果一個模型支援的最大序列長度為 4096 個標記，則必須預先分配所有 4096 個位置的空間。但實際上，大多數迭代可能只處理 200-500 個標記。因此，即使預留了 4096 個標記的記憶體，也只有一小部分被實際使用。剩餘的空間仍然空置。這種浪費的空間就是內部碎片。

外部碎片

外部碎片是指記憶體可用，但並非以連續的塊形式存在。不同的請求開始和結束的時間不同。這會在記憶體中留下太小的“空洞”，無法有效地重用。即使 GPU 有剩餘的視訊記憶體，也無法容納新的鍵值快取塊，因為空閒記憶體不連續。這會導致記憶體利用率低，即使理論上記憶體是存在的。

記憶體預留問題

記憶體預留問題本質上是預先規劃出錯導致的。我們預留記憶體時，假設模型會始終生成最長的序列。但實際上，生成過程可能很早就停止（例如，EOS 標記或達到提前停止標準）。這意味著預留的鍵值快取記憶體中很大一部分從未被使用，卻仍然無法供其他請求使用。

Source: Efficient Memory Management for Large Language Model Serving with PagedAttention

分頁注意力機制

分頁注意力機制（vLLM 中使用的機制）借鑑了作業系統管理記憶體的思路，解決了記憶體碎片問題。就像作業系統進行分頁一樣，vLLM 將鍵值快取分割成固定大小的小塊，而不是預留一個連續的大塊。

• 邏輯塊：每個請求在邏輯上看到的都是一個乾淨、連續的鍵值記憶體序列。它認為其標記是按順序儲存的，就像普通的注意力機制一樣。
• 物理塊：實際上，這些“邏輯塊”分散在 GPU 視訊記憶體 (VRAM) 中許多不連續的實體記憶體頁上。vLLM 使用塊表（類似於作業系統中的頁表）來對映邏輯位置和物理位置。

Source: Efficient Memory Management for Large Language Model Serving with PagedAttention

因此，PagedAttention 不需要一個 4096 格的大型連續緩衝區，而是將標記儲存在小塊（類似於頁）中，並在記憶體中靈活地排列它們。

這意味著：

• 不再有內部碎片
• 無需連續的 VRAM
• 可以立即重用空閒塊
• GPU 利用率從約 20-40% 躍升至約 96%

但我們不會在這裡深入探討這個概念。也許會在以後的部落格文章中介紹。

鍵值快取的記憶體消耗

要了解特定模型的鍵值快取消耗多少記憶體，我們需要了解以下四個變數：

• num_layers（Transformer 模型中的層數）
• num_heads（多頭自注意力層中每個頭的鍵值頭數量）
• head_dim（每個鍵值頭的維度）
• precision_in_bytes（假設精度為 16 位，即 2 位元組）（FP16 => 16 位 = 2 位元組）

假設我們的用例使用以下值：num_layers = 32，num_heads = 32，head_dim = 128，batch_size = 1（實際部署通常會使用更大的批處理大小）。

現在，讓我們計算每個 token 的鍵值快取。

每個 token 的 KV 快取大小 = 2 * (層數) * (頭數 * 頭維度) * 精度（位元組） * 批次大小

為什麼是 2？

因為每個 token 需要儲存兩個矩陣——K 矩陣和 V 矩陣。

KV cache per token = 2 * 32 * (32 * 128) * 2 * 1= 524288 B= 0.5 MB

我們需要 0.5 MB 的空間來儲存每個 token 在所有層和頭部的鍵值對資訊。真是令人震驚！

以 Llama 2 模型為例，我們知道，如果只傳送一個請求，序列長度為 4096。如果我們每次請求都提前傳送，並且使用鍵值快取，那麼我們需要儲存所有這些資訊。

Total KV cache per request = seq_len * KV cache per token= 4096 * 524288= 2147483648 / (1024 * 1024 * 1024)= 2 GB

所以基本上，每個請求需要 2 GB 記憶體。KV 快取用於處理單個請求。

小結

在本文章中，我們討論了 KV 快取的底層工作原理及其用途。透過討論，我們瞭解了為什麼它對於現代 LLM 來說是一項如此重要的最佳化。我們探討了 Transformer 最初如何應對高計算量和記憶體需求，以及 KV 快取如何透過儲存先前計算的鍵值向量來顯著減少冗餘工作。

然後，我們研究了 KV 快取的實際問題，例如 GPU 利用率低、對連續記憶體的要求，以及諸如記憶體預留、內部碎片和外部碎片等問題如何導致大量的視訊記憶體浪費。

我們還簡要討論了分頁注意力機制 (vLLM) 如何利用作業系統風格的分頁邏輯來解決這些問題。它允許 KV 快取以小的、可重用的塊而不是大的、固定的塊來增長，從而顯著提高 GPU 記憶體利用率。

最後，我們分析了 KV 快取記憶體消耗的計算。我們觀察到，在像 Llama-2 這樣的模型上，單個請求很容易就會佔用近 2 GB 的視訊記憶體用於鍵值儲存。真是驚人。

如果您想讓我更詳細地講解分頁注意力機制、塊表以及 vLLM 快速遍歷的原因……請在下方評論區留言，我會在下一篇部落格中嘗試解答。下次見！