現代 AI 系統的「記憶體牆」：KV Cache 的壓力與最佳化路徑

在 LLM（大型語言模型）的推論過程中，最昂貴的資源往往不是計算量（FLOPs），而是記憶體頻寬。當我們談論 AI 系統的性能瓶頸時，一個核心概念是 KV Cache（Key-Value Cache）。本文將深入探討 KV Cache 的本質、它如何造成「記憶體牆」，以及業界目前主流的優化方案。

什麼是 KV Cache？

Transformer 架構的核心是注意力機制（Attention）。在生成文字時，模型採用自迴歸方式：每產生一個新的 Token，都需要回顧之前所有已生成的 Token。

如果沒有快取，每生成第 $N$ 個 Token，模型都要重新計算前 $N-1$ 個 Token 的 Key 和 Value 向量。這意味著計算量隨序列長度呈平方級增長。為了避免重複計算，我們將之前步驟產生的 Key 和 Value 向量儲存在顯示記憶體中，這就是 KV Cache。

「記憶體牆」的成因

雖然 KV Cache 解決了計算冗餘，但它引入了巨大的記憶體壓力：

1. 空間佔用驚人：KV Cache 的大小與模型層數、隱藏層維度、序列長度以及 Batch Size 成正比。對於一個 Llama-3-70B 模型，在 FP16 精度下，單一請求的 KV Cache 在長文本情境下可達數 GB。
2. 頻寬受限 (Memory Bound)：推論過程是典型的存取記憶體密集型任務。GPU 計算核心的速度遠快於顯示記憶體讀取速度。在解碼階段（Decoding phase），模型需要將巨大的 KV Cache 從 HBM（高頻寬記憶體）搬運到 SRAM 中進行計算，導致 GPU 大部分時間在「等待資料」，而非執行計算。
3. 碎片化問題：傳統的顯示記憶體配置方式會導致大量碎片（External Fragmentation），使得實際可用顯示記憶體低於物理上限，限制了最大併發數（Throughput）。

業界的最佳化路徑

為了打破這堵「牆」，目前的最佳化方向主要集中在三個維度：減少儲存量、提高存取效率、最佳化記憶體管理。

1. 減少儲存量：MQA 與 GQA

早期的 Multi-Head Attention (MHA) 為每個 Query Head 分配一對 KV Head。
- MQA (Multi-Query Attention)：所有 Query Head 共用一組 KV Head。極大地降低了快取大小，但可能損失部分精度。
- GQA (Grouped-Query Attention)：折衷方案。將 Query 分組，每組共用一組 KV Head。Llama-3 等主流模型均採用 GQA，在性能和精度之間取得了平衡。

2. 提高存取效率：PagedAttention (vLLM)

借鑑作業系統虛擬記憶體的概念，vLLM 提出了 PagedAttention。它不再為每個請求分配連續的顯示記憶體區塊，而是將 KV Cache 分頁儲存於不連續的實體區塊中。
- 消除碎片：幾乎消除了內部和外部碎片。
- 動態共享：允許不同的請求共享相同的前綴快取（如相同的 System Prompt），極大提升了吞吐量。

3. 量化與壓縮

將 KV Cache 從 FP16/BF16 量化為 INT8 或 FP8，甚至更低位元寬度（如 4-bit）。這直接將記憶體佔用減半或更多，並減輕了頻寬壓力，而對模型生成品質的影響在可控範圍內。

總結

KV Cache 是 LLM 推論的必然產物，也是性能最佳化的主戰場。從演算法層面的 GQA 到系統層面的 PagedAttention，再到硬體友善的量化技術，其核心目標只有一個：讓資料搬運的速度跟上計算的速度。對於開發者而言，理解這一機制有助於在選擇部署框架和配置併發參數時做出更科學的決策。