現代 AI 的「KV Cache」：為什麼 LLM 越聊越慢，且顯存會被吃光？

當你與大型語言模型（LLM）進行長對話時，可能會發現兩個現象：一是隨著對話長度增加，生成速度逐漸變慢；二是顯存佔用不斷攀升，即使模型本身的大小沒有變化。這背後的核心機制就是 KV Cache（Key-Value Cache）。

什麼是 KV Cache？

要理解 KV Cache，首先要理解 Transformer 的注意力機制（Attention）。在生成每一個新 Token 時，模型需要回顧之前所有已經生成的 Token。

在計算注意力時，每個 Token 都會被轉化為三個向量：Query (Q), Key (K), 和 Value (V)。
- Query：當前 Token 「想要尋找什麼」。
- Key：之前的 Token 「能提供什麼資訊」。
- Value：之前的 Token 「實際包含的內容」。

計算過程是：用當前的 $Q$ 與之前所有 Token 的 $K$ 做點積，計算權重，然後根據權重對 $V$ 進行加權求和。

關鍵點在於： 對於已經生成的 Token，$K$ 和 $V$ 向量在後續的生成步驟中是完全不變的。如果每次生成新 Token 都要重新計算一遍之前所有 Token 的 $K$ 和 $V$，那麼計算量將隨長度呈平方級增長（$O(n^2)$），這會導致極大的資源浪費。

為了優化，工程師將這些已經計算過的 $K$ 和 $V$ 儲存在顯存中，這就是 KV Cache。下次生成時，直接從記憶體讀取舊的 $K, V$，只計算當前新 Token 的 $K, V$ 即可。

KV Cache 的代價：顯存的「隱形殺手」

雖然 KV Cache 大幅提升了推論速度（將複雜度降為線性），但它引入了一個巨大的空間成本。

1. 顯存佔用公式

KV Cache 的大小取決於：
- 模型層數 (Layers)
- 隱藏層維度 (Hidden Dimension)
- 注意力頭數 (Heads)
- 序列長度 (Sequence Length)
- 精度 (Precision, 如 FP16/BF16)

一個簡單的估算公式（對於標準 Transformer）：
$\text{Memory} = 2 \times \text{Layers} \times \text{Heads} \times \text{Dim_per_head} \times \text{Seq_len} \times \text{Bytes_per_param}$

例如，一個 Llama-3-8B 模型（32層，32頭，128維/頭），使用 FP16（2位元組）：
每個 Token 需要的快取 $\approx 2 \times 32 \times 32 \times 128 \times 2 = 524,288$ 位元組 $\approx 0.5\text{ MB}$。
如果上下文達到 32k tokens，僅 KV Cache 就需要 $\approx 16\text{ GB}$ 的顯存！這還不包括模型本身的權重。

2. 「碎片化」問題

傳統的 KV Cache 分配像是在記憶體中申請連續的大塊空間。但由於對話長度不可預測，系統往往會預先分配一個最大長度的緩衝區。這導致了嚴重的記憶體碎片化——很多預留空間被浪費了，而真正需要時卻可能因為碎片過多而 OOM（Out of Memory）。

如何優化？（工業界的解決方案）

為了解決上述問題，現代 AI 系統採用了多種黑科技：

PagedAttention (vLLM)

這是目前最主流的方案。借鑑了作業系統的「虛擬記憶體」概念，將 KV Cache 分成固定大小的「頁」（Pages）。不再要求連續儲存，而是透過頁表映射。這樣可以幾乎消除記憶體碎片，允許更高的併發量（Throughput）。

MQA 與 GQA (Multi-Query / Grouped-Query Attention)

既然 $K, V$ 太佔空間，能不能少存點？
- MQA: 所有 Query 頭共享同一組 $K, V$ 頭。顯存佔用直接降低到原來的 $1/\text{Heads}$。但效果會有所下降。
- GQA: 將 Query 頭分組，每組共享一組 $K, V$ 頭（如 Llama-3 使用 GQA）。在效能和顯存之間取得了極佳的平衡。

量化 (Quantization)

將 KV Cache 從 FP16 量化到 INT8 或 FP8。這能直接將快取空間減半，且對模型精度的影響微乎其微。

總結

KV Cache 是 LLM 推論中的一種「空間換時間」策略。它讓即時對話成為可能，但也成為了限制上下文長度和併發數的瓶頸。從 PagedAttention 到 GQA 再到量化，AI 工程化的核心目標之一就是如何更高效地管理這塊昂貴的顯存快取。