現代 AI 的「記憶碎片」：KV Cache 深度解析

在討論 LLM 推論最佳化時，PagedAttention 和 Speculative Decoding 經常被提及，但它們共同作用的核心物件只有一個：KV Cache（Key-Value Cache）。如果你想理解為什麼大型語言模型推論如此吃顯示記憶體，以及為什麼推論速度會隨著上下文增加而變慢，KV Cache 是唯一的入口。

什麼是 KV Cache？

LLM 的核心是 Transformer 架構。在生成文字時，模型採用的是自回歸（Autoregressive）模式：每生成一個新 token，都需要將之前所有生成的 token 作為輸入重新餵給模型。

在 Transformer 的注意力機制（Attention）中，每個 token 都會產生三個向量：Query (Q)、Key (K) 和 Value (V)。
- Query：當前 token「想要尋找什麼」。
- Key：歷史 token「能提供什麼」。
- Value：歷史 token「實際包含的內容」。

計算過程是：$\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$。

關鍵點在於：對於已經生成過的歷史 token，它們的 $K$ 和 $V$ 向量在後續的生成步驟中是完全不變的。如果每次生成新 token 都要重新計算一遍所有歷史 token 的 $K$ 和 $V$，那麼計算量將隨序列長度呈平方級增長 $\mathcal{O}(n^2)$。

為了避免這種重複計算，我們將每一步產生的 $K$ 和 $V$ 儲存在顯示記憶體中，這就是 KV Cache。這樣，每一步只需要計算當前新 token 的 $Q, K, V$，然後直接從快取中讀取之前的 $K, V$ 進行矩陣乘法即可。計算複雜度因此降為 $\mathcal{O}(n)$。

KV Cache 的代價：顯示記憶體黑洞

雖然 KV Cache 極大地提升了速度，但它帶來了巨大的顯示記憶體壓力。

1. 計算公式

一個模型的 KV Cache 大小取決於：
- $\text{batch_size}$ (並發數)
- $\text{seq_len}$ (序列長度)
- $\text{num_layers}$ (層數)
- $\text{num_heads}$ (注意力頭數)
- $\text{head_dim}$ (每個頭的維度)
- $\text{precision}$ (精度，如 FP16 為 2 bytes)

公式為：$\text{Memory} = 2 \times \text{batch_size} \times \text{seq_len} \times \text{num_layers} \times \text{num_heads} \times \text{head_dim} \times \text{precision}$
(乘以 2 是因為要同時儲存 Key 和 Value)。

2. 實戰量化

以 Llama-3-8B 為例（假設 FP16）：
- 層數: 32, 頭數: 32, 每個頭維度: 128.
- 單個 token 在單層產生的 KV 大小 = $2 \times 32 \times 128 \times 2\text{ bytes} = 16\text{ KB}$。
- 全模型單 token KV 大小 = $32\text{ layers} \times 16\text{ KB} = 512\text{ KB}$。

看起來不多？但如果 batch_size=32 且 seq_len=4096：
$32 \times 4096 \times 512\text{ KB} \approx 67\text{ GB}$。

這意味著即使模型權重本身只佔約 15GB，為了支援高並發的長文本推論，你可能需要一張 A100 (80GB) 或更多顯示卡，僅僅是為了存放這些「記憶碎片」。

如何最佳化 KV Cache？

面對顯示記憶體壓力，工業界演進出了三種主流方案：

MQA 與 GQA（結構最佳化）

傳統的 Multi-Head Attention (MHA) 每個 Query 頭對應一個 Key/Value 頭。
- MQA (Multi-Query Attention)：所有 Query 頭共享一對 KV 頭。顯示記憶體佔用直接降低到原來的 $1/\text{num_heads}$，但精度有損。
- GQA (Grouped-Query Attention)：折衷方案。將 Query 分組，每組共享一對 KV 頭（如 Llama-3 使用）。在保持效能的同時顯著降低顯示記憶體佔用。

PagedAttention（記憶體管理最佳化）

傳統的 KV Cache 要求連續記憶體空間，導致嚴重的記憶體碎片化（Internal Fragmentation）。vLLM 推出的 PagedAttention 將 KV Cache 分頁儲存（類似作業系統的虛擬記憶體），允許非連續儲存且動態按需分配，將顯示記憶體利用率提升至接近 100%。

量化（精度最佳化）

將 KV Cache 從 FP16 量化到 INT8 或 FP8，甚至 INT4。這可以直接將顯示記憶體佔用減半或更多，而對模型生成品質的影響在可接受範圍內。

總結

KV Cache 是 LLM 推論的「空間換時間」典範。它解決了自回歸生成的冗餘計算問題，但也成為了限制吞吐量和上下文長度的最大瓶頸。從 GQA 到 PagedAttention 再到量化，AI 系統工程的演進本質上就是在與這塊巨大的「記憶碎片」做鬥爭。