現代 AI 的「推理加速之謎」：從 KV Cache 到 PagedAttention 的演進

在 LLM 推理優化領域，人們經常討論 PagedAttention 或 Speculative Decoding，但這些技術解決的底層問題其實只有一個：如何高效地管理和利用 KV Cache（Key-Value Cache）。如果你想理解為什麼大模型推理如此吃顯示記憶體，以及為什麼上下文越長速度越慢，KV Cache 是唯一的答案。

1. 為什麼需要 KV Cache？

LLM 生成文字是自回歸（Autoregressive）的：每生成一個新 token，都需要將之前所有生成的 token 作為輸入重新餵給模型。

在 Transformer 的注意力機制中，每個 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)$。

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

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

計算公式

一個模型的 KV Cache 大小取決於：$\text{batch_size} \times \text{seq_len} \times \text{num_layers} \times \text{num_heads} \times \text{head_dim} \times \text{precision} \times 2$ (Key 和 Value 各一份)。

實戰量化

以 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) 或更多顯示卡，僅僅是為了存放這些「記憶碎片」。

3. 從「暴力儲存」到「智慧管理」的演進

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

A. GQA (Grouped-Query Attention) —— 從結構上減負

傳統的 MHA 每個 Query 頭對應一個 KV 頭。GQA 將 Query 分組，每組共享一對 KV 頭（如 Llama-3 使用）。這在保持效能的同時顯著降低了 KV Cache 的物理尺寸。

B. PagedAttention —— 從記憶體管理上優化

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

C. 量化 (Quantization) —— 從精度上壓縮

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

總結

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