現代 AI 系統中的 KV Cache：從記憶體瓶頸到推論加速的煉金術

在 LLM（大型語言模型）的推論過程中，最昂貴的成本之一並非計算量，而是記憶體頻寬。當你與 AI 對話時，模型需要回顧之前的所有上下文。如果每次生成新 Token 都要重新計算一遍之前的所有 Key 和 Value 向量，推論速度將隨著序列長度的增加而呈平方級下降。

為了解決這個問題，工業界引入了 KV Cache（Key-Value Cache）。簡單來說，KV Cache 就是一種「以空間換取時間」的策略：將已經計算過的 Token 的 K 和 V 向量儲存在顯示記憶體中，下次預測時直接讀取，而無需重新計算。

KV Cache 的物理本質

在 Transformer 的 Attention 機制中，每個 Token 會生成一個 Query (Q)、Key (K) 和 Value (V)。
- Query：當前要預測的 Token「在找什麼」。
- Key：歷史 Token「擁有什麼特徵」。
- Value：歷史 Token「包含的具體資訊」。

Attention 的核心是 $Softmax(QK^T / \sqrt{d})V$。注意到對於歷史 Token 而言，它們的 K 和 V 在模型權重不變的情況下是靜態的。這意味著一旦一個 Token 被處理過，它的 K 和 V 就不再改變。KV Cache 將這些向量快取起來，使得每一步推論只需要計算當前最新 Token 的 QKV，然後與快取中的歷史 KV 進行矩陣乘法。

記憶體牆：KV Cache 的代價

雖然 KV Cache 加速了推論，但它帶來了巨大的顯示記憶體壓力。其佔用空間由以下公式決定：
$\text{Memory} = 2 \times \text{layers} \times \text{heads} \times \text{dim} \times \text{seq_len} \times \text{precision}$

以 Llama-3-8B 為例（假設 FP16 精度）：
- 層數 (layers): 32
- 頭數 (heads): 32
- 每個頭的維度 (dim): 128
- 每個元素大小: 2 bytes

對於一個長度為 4096 的上下文，單個請求的 KV Cache 大約需要：
$2 \times 32 \times 32 \times 128 \times 4096 \times 2 \approx 2.1\text{ GB}$

在併發請求（Batch Size）增加時，顯示記憶體佔用會線性增長。如果 Batch Size 為 32，僅 KV Cache 就需要超過 60GB 顯示記憶體，這直接導致了許多模型在長文本場景下出現 OOM（Out of Memory，記憶體不足）。

工程優化路徑：從 PagedAttention 到 GQA

為了打破記憶體瓶頸，AI 系統工程演進出了三種主流方案：

1. PagedAttention (vLLM)

傳統的 KV Cache 要求連續的記憶體空間，這導致了嚴重的「碎片化」——很多預分配的空間被浪費掉。vLLM 借鑑了作業系統虛擬記憶體的分頁機制（Paging），將 KV Cache 分成固定大小的塊（Blocks），透過一個 Block Table 進行映射。這樣可以實現近乎零浪費的顯示記憶體利用率，極大地提升了吞吐量。

2. Grouped-Query Attention (GQA)

全量注意力（MHA）中每個 Q 頭對應一個 K/V 頭。GQA 則讓多個 Q 頭共享一組 K/V 頭。例如 Llama-3 使用 GQA 將 KV 頭數減少到 Q 頭數的 $1/8$。這直接將 KV Cache 的顯示記憶體佔用降低到原來的 $12.5\%$，且對模型精度的影響微乎其微。

3. 量化與壓縮 (Quantization)

將 FP16 的 KV Cache 量化為 INT8 或 FP8。這不僅能減半顯示記憶體佔用，還能利用硬體加速指令提升讀取速度。目前的趨勢是採用 Per-channel 或 Per-token 量化來維持精度。

實戰啟示：如何選擇推論框架？

對於開發者而言，理解 KV Cache 有助於優化部署策略：
- 追求高吞吐量 $\rightarrow$ 選擇支援 PagedAttention 的框架（如 vLLM, TensorRT-LLM）。
- 受限於顯示記憶體容量 $\rightarrow$ 選擇支援 GQA 或量化 KV Cache 的模型版本（如 Mistral, Llama-3）。
- 超長文本處理 $\rightarrow$ 注意監控 kv_cache_usage 指標，合理設定 max_model_len 以防止突發 OOM。

KV Cache 是 LLM 從「實驗室玩具」走向「工業級產品」的關鍵工程基石。它揭示了一個深刻的真理：在 AI 系統中，演算法的優雅往往需要透過極致的記憶體管理來落地。