現代 AI 的「推理速度」：KV Cache 的工程真相

在 LLM 的宣傳手冊中，我們經常看到「每秒生成 100 個 token」這樣的速度指標。但對於開發者來說，決定這個速度的核心不在於 GPU 的算力（TFLOPS），而在於一個極其關鍵的記憶體最佳化機制：KV Cache（Key-Value Cache）。

如果你在呼叫 API 時感覺到首字回應快但後續生成慢，或者在部署模型時發現 VRAM 被迅速吃光，那麼你實際上正在與 KV Cache 打交道。

為什麼需要 KV Cache？

要理解 KV Cache，首先要理解 Transformer 的自回歸（Autoregressive）特性。

LLM 生成文字是一個一個 token 產生的。當你輸入「今天天氣」，模型預測出「很」；然後模型將「今天天氣很」作為輸入，預測出「好」。

在計算第 $N$ 個 token 時，模型需要計算當前 token 與之前所有 $N-1$ 個 token 的注意力（Attention）。這意味著：
1. 重複計算：如果沒有快取，每次生成新 token，都要重新計算前面所有 token 的 $Q$ (Query), $K$ (Key), $V$ (Value) 向量。
2. 複雜度爆炸：計算量隨序列長度呈平方級增長 $\mathcal{O}(N^2)$。

KV Cache 的核心邏輯很簡單：既然之前的 token 不會改變，那麼它們的 $K$ 和 $V$ 向量在每一輪迭代中都是相同的。我們只需要把它們存起來，下次直接用。

KV Cache 是如何運作的？

在推理過程中，模型分為兩個階段：

1. Prefill 階段（預填充）

當你發送 Prompt 時，模型一次性處理所有輸入 token。此時它會計算所有輸入 token 的 $K$ 和 $V$，並將它們寫入 VRAM 中的 KV Cache 區域。這個階段是計算密集型的（Compute-bound），因為它可以利用 GPU 的平行處理能力。

2. Decoding 階段（解碼）

生成每個新 token 時，模型只需要為這一個新 token 計算其 $Q, K, V$。然後，它將這個新的 $K, V$ 追加到快取中，並利用快取中的歷史 $K, V$ 來計算注意力權重。這個階段是記憶體頻寬密集型的（Memory-bound），因為 GPU 大部分時間在等待從 VRAM 中讀取巨大的 KV Cache 矩陣。

工程上的殘酷代價：VRAM 壓力

KV Cache 雖然解決了計算冗餘，但它引入了巨大的記憶體開銷。

KV Cache 的大小取決於：$\text{層數} \times \text{頭數} \times \text{維度} \times \text{序列長度} \times \text{精度}$。
以 Llama-3-8B 為例（FP16 精度）：
- 每增加一個 token，每個請求大約需要消耗數百 KB 到數 MB 的 VRAM。
- 當併發使用者增加或上下文視窗擴大到 128K 時，KV Cache 會迅速撐爆 A100/H100 的 VRAM，導致 OOM (Out of Memory)。

如何最佳化 KV Cache？（業界方案）

為了在不犧牲效能的前提下支援更長上下文和更高併發，業界採用了三種主流方案：

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

這是從架構層面減少快取量。傳統的 MHA 每個 Query 頭都有對應的 Key/Value 頭；而 GQA 讓多個 Query 頭共享一組 KV 頭。這直接將 KV Cache 的體積壓縮了數倍（例如 Llama-3 就使用了 GQA）。

2. PagedAttention (vLLM)

這是目前最主流的系統級最佳化方案。傳統的 KV Cache 要求連續的記憶體空間，導致嚴重的碎片化（類似早期的作業系統記憶體管理）。PagedAttention 將 KV Cache 分頁儲存在不連續的實體區塊中，實現了類似虛擬記憶體的管理方式，極大地提升了吞吐量並降低了浪費。

3. 量化 (Quantization)

將 KV Cache 從 FP16 量化到 INT8 或 FP8。這可以將 VRAM 佔用直接減半，且對模型精度的影響微乎其微。

總結

KV Cache 是 LLM 推理從「實驗室玩具」變成「工業產品」的關鍵工程基石。它將 $\mathcal{O}(N^2)$ 的重複計算轉化為 $\mathcal{O}(N)$ 的空間換時間策略。當我們討論 AI 推理成本和延遲時，本質上是在討論如何更高效地管理這塊昂貴的 VRAM 快取。