現代 AI 系統的「記憶力」之謎：從上下文視窗到 KV Cache 的工程真相

在與大型語言模型（LLM）對話時，我們經常聽到「上下文視窗（Context Window）」這個詞。例如某個模型支援 128K 或 1M 個 token。許多人的直覺認為，這就像是給 AI 分配了一塊記憶體，只要在這個範圍內，AI 就能「記得」之前說過的話。

但從工程實作的角度來看，上下文的維持並非簡單的「讀取」，而是一場關於顯示記憶體、計算量與頻寬的殘酷博弈。其核心機制正是 **KV Cache（Key-Value Cache）**。

為什麼需要 KV Cache？

要理解 KV Cache，首先要看 Transformer 的自回歸生成模式。LLM 生成文字是逐個 token 進行的：生成第 $N$ 個 token 時，模型需要回顧前 $N-1$ 個 token 的所有資訊。

在標準的 Attention 機制中，每個 token 都會被轉化為三個向量：**Query (Q)**, **Key (K)**, 和 **Value (V)**。

- **Q** 代表「我想找什麼」。

- **K** 代表「我包含什麼資訊」。

- **V** 代表「如果我被選中，我提供什麼內容」。

如果沒有快取，每生成一個新的 token，模型都要重新計算之前所有 token 的 K 和 V 向量。這意味著計算量隨序列長度呈**平方級增長** $\mathcal{O}(n^2)$。對於一個 1000 字的對話，生成最後一個字時，模型要重複計算前面 999 個字的 K 和 V。這在生產環境下是不可接受的延遲。

**KV Cache 的核心邏輯很簡單：既然之前的 token 不會改變，那麼它們的 K 和 V 向量在後續生成中永遠是一樣的。我們將它們存起來，下次直接用。**

KV Cache 的代價：顯示記憶體的「吞噬者」

KV Cache 雖然解決了計算冗餘（將複雜度降至 $\mathcal{O}(n)$），但它引入了另一個巨大的挑戰：**顯示記憶體佔用**。

KV Cache 需要儲存每一層、每個注意力頭的所有 token 的 K 和 V 向量。我們可以簡單算一筆帳：

假設一個模型有 $L$ 層，隱藏層維度為 $d$，序列長度為 $s$，使用 FP16（2位元組）精度：

$\text{Memory} = 2 \times L \times s \times d \times 2 \text{ bytes}$

對於一個中型模型（如 Llama-3-8B），當上下文達到 32K 時，僅 KV Cache 這一項就可能佔用數 GB 的顯示記憶體。這意味著即使你的 GPU 有足夠的算力去跑推論，也可能因為顯示記憶體被 KV Cache 填滿而導致 `Out of Memory (OOM)`。

工程優化：如何讓「記憶」更高效？

為了緩解顯示記憶體壓力並提升吞吐量，工業界演進出了幾種關鍵技術：

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

傳統的 Multi-Head Attention 為每個頭都維護一套 KV 對。

- **MQA (多查詢注意力)**：所有頭共享同一套 KV 對。極大地降低了顯示記憶體佔用，但可能會損失一定的表達能力。

- **GQA (分組查詢注意力)**：折衷方案。將頭分組，每組共享一套 KV 對（如 Llama-3 所採用）。它在效能和顯示記憶體之間取得了極佳的平衡。

2. PagedAttention (vLLM 的核心)

傳統的 KV Cache 在記憶體中是連續儲存的，這會導致嚴重的**記憶體碎片化**（類似於作業系統的記憶體管理問題）。

vLLM 引入了 **PagedAttention**，將 KV Cache 分塊儲存在不連續的實體記憶體頁中，透過一個頁表進行映射。這使得顯示記憶體利用率接近 100%，允許系統在同一時間處理更多併發請求（Batch Size 大幅提升）。

3. FlashAttention

雖然 FlashAttention 主要優化的是計算過程中的 IO 開銷（減少 HBM 與 SRAM 之間的資料搬運），但它與 KV Cache 的高效管理相輔相成，共同降低了長文本推論的整體延遲。

總結：上下文不是魔法，而是資源調度

當我們談論 AI 的「長文本能力」時，我們實際上是在討論該系統如何高效地管理其 KV Cache。一個支援百萬級上下文的模型，其背後的工程挑戰不在於演算法本身（Transformer 公式沒變），而在於如何透過 PagedAttention、GQA 等手段在有限的 HBM 中塞進更多的狀態資訊，並確保讀取速度跟得上推論速度。

理解了這一點，你就能明白為什麼增加上下文長度往往意味著推論成本的激增以及回應速度的下降——因為你正在與物理世界的記憶體頻寬和容量做鬥爭。