現代 AI 系統的「記憶擴容」：KV Cache 壓縮與量化深度解析

在 LLM（大型語言模型）的推論過程中，最昂貴的資源不是計算量（FLOPs），而是顯示記憶體頻寬和容量。當你與 AI 進行長對話時，模型需要記住之前所有的上下文。為了避免每次生成新 token 都重新計算一遍之前的所有 token，AI 系統引入了 KV Cache（Key-Value Cache）。

然而，KV Cache 是一個巨大的「記憶體吞噬者」。隨著上下文長度（Context Length）的增加，KV Cache 呈線性增長，迅速填滿 GPU 顯示記憶體，導致最大併發數（Batch Size）下降，甚至觸發 OOM（Out of Memory）。

本文將深入探討現代 AI 系統如何透過壓縮（Compression）與量化（Quantization）來優化 KV Cache，在不顯著犧牲精度的情況下，實現更長的上下文和更高的吞吐量。

1. KV Cache 的本質：為什麼它這麼大？

在 Transformer 的注意力機制中，每個 token 在每一層都會生成一個 Query (Q)、Key (K) 和 Value (V) 向量。
- Query：當前 token 「想要尋找什麼」。
- Key：歷史 token 「包含什麼資訊」。
- Value：歷史 token 「實際提供的內容」。

在自回歸生成時，歷史 token 的 K 和 V 在後續步驟中是不變的。將它們快取起來，只需計算當前新 token 的 QKV 即可。

顯示記憶體佔用計算公式：
$\text{Size} = 2 \times \text{layers} \times \text{heads} \times \text{dim_head} \times \text{precision} \times \text{seq_len}$

以 Llama-3-8B 為例（FP16 精度）：
- 層數 $\approx 32$
- 頭數 $\approx 32$
- 每個頭的維度 $= 128$
- 每個元素 $2$ bytes (FP16)
對於 $1024$ 個 token，單請求的 KV Cache 就需要約 $0.5\text{GB}$。當併發數達到 $100$ 或上下文達到 $32\text{k}$ 時，顯示記憶體壓力將變得不可承受。

2. KV Cache 量化：從 FP16 到 INT8/INT4

量化是最直接的減重手段——將高精度的浮點數轉換為低精度的整數。

FP16 $\rightarrow$ INT8 / FP8

透過將 $\text{FP16}$ 的值映射到 $\text{INT8}$ 或 $\text{FP8}$（如 NVIDIA H100 支援的 E4M3/E5M2），顯示記憶體佔用直接減半。為了減少精度損失，通常採用 Per-token Quantization 或 Per-channel Quantization，為每個向量維護一個縮放因子（Scale Factor）。

INT4 量化與離群值處理

進一步壓到 $\text{INT4}$ 可以節省 $75\%$ 的空間。但 KV Cache 中存在極少數數值巨大的「離群值」（Outliers），這些值承載了關鍵的語意資訊。簡單的線性量化會導致嚴重的精度崩塌。現代方案如 KIVI 透過對 K 和 V 分別進行低比特量化，並結合動態縮放因子來保留這些關鍵資訊。

3. KV Cache 壓縮：丟棄不重要的記憶

並非所有歷史 token 對當前預測都同等重要。壓縮的核心在於：識別並剔除冗餘的 KV 對。

H2O (Heavy Hitter Oracle)

研究發現，注意力權重分佈具有極強的稀疏性——只有少數幾個「重擊者」（Heavy Hitters）token 被頻繁存取。H2O 演算法透過維護一個固定大小的快取池，即時追蹤哪些 token 的累積注意力權重最高，僅保留這些核心 token，而將低貢獻的 token 從快取中剔除。

StreamingLLM (滑動視窗 + 注意力錨點)

在處理超長文本時，模型會出現「注意力崩潰」現象。StreamingLLM 發現只要保留 Attention Sink（即序列開頭的前幾個 token）加上最近的一個滑動視窗（Sliding Window），模型就能在不重新訓練的情況下維持穩定的推論能力且顯示記憶體佔用恆定。

4. 工程實踐中的權衡

在實際部署 AI 系統時，選擇哪種方案取決於業務場景：

方案	實作難度	顯示記憶體收益	精度影響	適用場景
FP8 量化	低 (硬體支援)	$2\times$	微乎其微	通用加速, H100 叢集
INT4 量化	中	$4\times$	有感知	端側設備, 超長上下文
H2O / 稀疏化	高 (需動態管理)	可自訂	中等	長文件分析, RAG 系統
StreamingLLM	中	常數級佔用	對遠端記憶遺失嚴重	無限流對話, Bot 服務

總結

KV Cache 的優化是 LLM 從「實驗室玩具」走向「工業級產品」的關鍵一步。透過量化降低單個元素的體積，透過壓縮減少元素的數量，AI 系統正在打破記憶體牆的限制。未來的趨勢將是量化與稀疏化的深度結合——讓模型像人類一樣，既能高效儲存海量資訊，又能精準地提取關鍵記憶。