現代 AI 的「KV 快取」(KV Cache)：為什麼大型語言模型生成文字時速度會變慢？

當你與大型語言模型（LLM）對話時，你會發現它生成文字的方式是「一個詞一個詞」蹦出來的。在技術底層，這種生成模式被稱為 自回歸生成 (Autoregressive Generation)。而為了讓這個過程不至於慢到無法忍受，工程師們引入了一個至關重要的最佳化機制：KV 快取 (Key-Value Cache)。

簡單來說，KV 快取就像是 AI 的「短期記憶快照」，它避免了 AI 在每寫一個新詞時，都要把之前寫過的所有內容重新讀一遍。

1. 問題的核心：重複計算的浪費

要理解 KV 快取，首先要理解 Transformer 架構中的 注意力機制 (Attention)。

在生成第 $N$ 個詞時，AI 需要回顧前 $N-1$ 個詞，計算當前詞與之前所有詞之間的相關性（即 Attention Score）。這個計算過程涉及兩個關鍵矩陣：Key (K) 和 Value (V)。

如果沒有快取，生成過程會變成這樣：
- 生成第 1 個詞：計算 $\text{Token}_1$ 的 K 和 V $\rightarrow$ 輸出 $\text{Word}_1$。
- 生成第 2 個詞：重新計算 $\text{Token}_1$ 的 K, V $\rightarrow$ 計算 $\text{Token}_2$ 的 K, V $\rightarrow$ 輸出 $\text{Word}_2$。
- 生成第 3 個詞：重新計算 $\text{Token}_1, \text{Token}_2$ 的 K, V $\rightarrow$ 計算 $\text{Token}_3$ 的 K, V $\rightarrow$ 輸出 $\text{Word}_3$。

你會發現，$\text{Token}_1$ 的 K 和 V 在每一次迭代中都被重複計算了無數次。隨著文本長度增加，這種冗餘計算量呈平方級增長，導致生成速度劇烈下降。

2. KV Cache 是如何運作的？

KV 快取的邏輯非常簡單：既然之前的 Token 在後續生成中其 K 和 V 值是不變的，那麼我們只需要在第一次計算時把它們存起來即可。

具體流程如下：
1. 預填充階段 (Prefill Phase)：AI 處理你輸入的 Prompt。它一次性計算所有輸入 Token 的 K 和 V 值，並將它們儲存在顯示記憶體（VRAM）中。
2. 解碼階段 (Decoding Phase)：
- AI 生成新 Token 時，只計算這 一個 新 Token 的 K 和 V 值。
- 將這個新值追加到之前的 KV Cache 中。
- 利用完整的 KV Cache（歷史 + 當前）進行注意力計算，得出下一個詞。

透過這種方式，每次迭代的計算量從 $O(N^2)$ 降低到了 $O(N)$。這就是為什麼即使對話很長，AI 生成單個詞的速度依然能保持相對穩定的原因。

3. 代價：顯示記憶體的「吞噬者」

天下沒有免費的午餐。KV 快取雖然節省了計算時間（Compute），但極大地消耗了記憶體（Memory）。

KV Cache 的大小取決於幾個因素：
- 模型規模（層數、隱藏層維度）。
- 序列長度（對話越長，快取越大）。
- Batch Size（同時服務多少個使用者）。

對於一個中型模型（如 Llama-3-8B），在 FP16 精度下，每 1000 個 Token 可能需要數百 MB 的顯示記憶體。當併發使用者增加或上下文視窗擴大到 128K 時，KV Cache 會迅速填滿 GPU 顯示記憶體，導致 OOM (Out of Memory) 或強制降低 Batch Size。

4. 如何最佳化 KV Cache？

為了緩解顯示記憶體壓力，工業界開發了多種黑科技：

• Multi-Query Attention (MQA) / Grouped-Query Attention (GQA)：不再為每個 Attention Head 分配獨立的 K 和 V，而是讓多個 Head 共用一組 KV 對。這能將 KV Cache 的體積直接壓縮數倍（例如 GQA 將其減少到原來的 $1/8$），而效能損失極小。
• PagedAttention (vLLM)：借鑑作業系統的虛擬記憶體管理方案。不再為每個請求分配連續的顯示記憶體區塊，而是將 KV Cache 分頁儲存在不連續的記憶體空間中，極大提高了顯示記憶體利用率並支援更高的併發量。
• 量化 (Quantization)：將 KV Cache 從 FP16 量化為 INT8 或 FP8，直接將記憶體佔用減半且幾乎不影響精度。

總結

KV 快取是現代 LLM 能實現即時互動的基石。它透過「以空間換取時間」的策略，將昂貴的重複計算轉化為簡單的記憶體讀取。當我們討論大型語言模型的「上下文視窗」上限時，本質上討論的不僅是模型的注意力能力，更是硬體顯示記憶體能否裝下那龐大的 KV Cache 快照。