現代 AI 的「KV Cache」：大型語言模型如何透過「記憶」來加速生成？

在與大型語言模型（LLM）對話時，你可能會注意到一個現象：模型產生答案的速度通常是恆定的，但當你輸入一段極長的文本後，模型在吐出第一個字之前的「思考時間」（Time to First Token, TTFT）會明顯增加。

這種現象背後涉及到一個至關重要的工程優化技術——KV Cache（Key-Value Cache，鍵值快取）。它是目前所有主流 LLM（如 GPT-4、Claude、Llama）能夠實現高效即時生成的基石。

為什麼需要 KV Cache？

要理解 KV Cache，首先要理解 Transformer 架構中的 Attention（注意力機制）。

在產生文本時，LLM 是一個「自回歸」過程：它每次只預測下一個 token。為了預測第 $N+1$ 個 token，模型必須回顧之前所有的 $N$ 個 token。

在標準的 Attention 計算中，每個 token 都會被轉化為三個向量：Query (Q)、Key (K) 和 Value (V)。
- Query：代表「我現在想要尋找什麼資訊」。
- Key：代表「我包含什麼樣的資訊」。
- Value：代表「如果我被選中，我能提供什麼具體內容」。

如果沒有快取，每產生一個新 token，模型都要對之前所有的 token 重新計算一遍 K 和 V 向量。這意味著：
- 產生第 1 個詞時，計算 1 次 K/V；
- 產生第 2 個詞時，重新計算前 1 個詞 + 當前詞的 K/V；
- 產生第 100 個詞時，重新計算前 99 個詞 + 當前詞的 K/V。

這種計算量隨序列長度呈 平方級增長 ($\mathcal{O}(n^2)$)，會導致極大的計算浪費和嚴重的延遲。

KV Cache 的工作原理：用空間換取時間

KV Cache 的核心思想非常簡單：既然之前的 token 在輸入過程中沒有改變，那麼它們的 K 和 V 向量在每一輪迭代中都是相同的，沒必要重複計算。

具體流程如下：
1. Prefill（預填充階段）：當你發送 Prompt 時，模型一次性處理所有輸入 token，計算出它們的 K 和 V 向量，並將它們儲存在顯存（GPU Memory）中。
2. Decoding（解碼階段）：
- 模型只為當前最新的一個 token 計算 Q, K, V。
- 將這個新產生的 K 和 V 追加到快取中。
- 使用當前的 Q 與快取中所有的 K 進行點積計算權重 $\text{softmax}(\frac{QK^T}{\sqrt{d}})$。
- 根據權重對快取中所有的 V 進行加權求和，得出最終輸出。

透過這種方式，每一步的計算量從 $\mathcal{O}(n^2)$ 降低到了 $\mathcal{O}(n)$。

KV Cache 的代價：顯存的「吞噬者」

雖然 KV Cache 大幅提升了速度，但它引入了一個巨大的挑戰：顯存壓力。

KV Cache 需要儲存每一層、每個注意力頭的所有 token 的 K 和 V 向量。其佔用空間公式大致為：
$$\text{Memory} = 2 \times \text{layers} \times \text{heads} \times \text{dim} \times \text{seq_len} \times \text{precision}$$

舉個例子：對於一個典型的 Llama-3-8B 模型（FP16 精度），如果上下文長度達到 32K tokens，僅 KV Cache 就可能佔用數 GB 的顯存。這就是為什麼很多 AI 服務會對上下文長度設限，或者在長文本下出現 OOM（Out of Memory）的原因。

如何優化 KV Cache？

為了緩解顯存壓力，工業界開發了多種優化方案：

1. MQA (Multi-Query Attention) 與 GQA (Grouped-Query Attention)

傳統的 Multi-Head Attention 為每個 Query 頭配備一個 Key 頭和一個 Value 頭。
- MQA：所有 Query 頭共享同一個 Key 和 Value 頭。這直接將 KV Cache 的體積縮小了 $H$ 倍（$H$ 為頭數）。
- GQA：折衷方案。將 Query 頭分組，每組共享一對 KV 頭。這在保持模型效能的同時顯著降低了記憶體佔用（Llama-3 等主流模型均採用 GQA）。

2. PagedAttention (vLLM)

傳統的快取需要連續的記憶體空間，容易產生碎片化導致浪費（類似早期的作業系統記憶體管理）。vLLM 推出的 PagedAttention 將 KV Cache 分頁儲存在不連續的實體記憶體區塊中，極大提高了顯存利用率並支援更高的併發量（Throughput）。

3. 量化 (Quantization)

將 KV Cache 從 FP16 量化為 INT8 或 FP8，可以在不明顯損失精度的情況下直接將記憶體佔用減半。

總結

KV Cache 是 LLM 從「實驗室玩具」變成「實用產品」的關鍵工程突破之一。它透過犧牲顯存空間來消除冗餘計算，讓 AI 能以人類閱讀的速度流暢地輸出文字。當我們討論 AI 的上下文視窗大小和推論成本時，本質上就是在討論如何更高效地管理這塊名為 「KV Cache」 的記憶空間。