現代 AI 系統中的 KV Cache 最佳化：從記憶體牆到 PagedAttention

在 LLM 推論過程中，最核心的效能瓶頸之一並非計算量（Compute-bound），而是記憶體頻寬（Memory-bound）。當我們討論生成式 AI 的推論速度時，實際上是在討論如何有效率地管理 KV Cache（Key-Value Cache）。

什麼是 KV Cache？

在 Transformer 的解碼階段，模型每生成一個新 Token，都需要與之前所有已生成的 Token 進行 Attention 計算。如果每次都重新計算所有歷史 Token 的 Key 和 Value 向量，計算複雜度將隨序列長度呈平方級增長。

為了避免重複計算，我們將之前步驟產生的 K 和 V 向量儲存在記憶體中，這就是 KV Cache。這意味著在生成第 $n$ 個 Token 時，我們只需要計算當前 Token 的 K 和 V，並將其追加到快取中。

記憶體牆與碎片化問題

儘管 KV Cache 解決了計算冗餘，但它引入了巨大的記憶體壓力。一個典型的 Llama-3-70B 模型，其 KV Cache 在 FP16 精度下，每個 Token 消耗的記憶體量驚人。隨著併發請求（Batch Size）和上下文長度（Context Length）的增加，顯示記憶體會被迅速填滿。

傳統的 KV Cache 管理方式是預先配置一塊連續的記憶體空間。這種做法存在兩個致命缺陷：
1. 內部碎片化：為了應對最大可能長度（Max Sequence Length），系統必須為每個請求預留最大空間。如果實際生成的長度遠低於最大值，大量顯示記憶體被浪費。
2. 外部碎片化：由於請求的生命週期不同且長度不一，記憶體空間會變得支離破碎，導致無法配置大塊連續空間給新請求。

PagedAttention：借鑑虛擬記憶體機制

vLLM 提出的 PagedAttention 徹底改變了這一現狀。其核心思想是將 KV Cache 的儲存方式從「連續陣列」改為「分頁儲存」，類似於作業系統的虛擬記憶體管理。

工作原理

PagedAttention 將 KV Cache 分割為固定大小的「區塊」（Blocks）。每個區塊可以儲存固定數量的 Token（例如 16 個）。
- 邏輯區塊 $\rightarrow$ 實體區塊：模型在邏輯上認為 KV Cache 是連續的，但實體上它們分佈在不連續的顯示記憶體區塊中。
- 區塊表（Block Table）：系統維護一張映射表，記錄每個邏輯區塊對應的實體位址。
- 動態配置：只有當目前的實體區塊被填滿時，系統才會為該請求配置一個新的實體區塊。

帶來的提升

1. 近乎零浪費：除了最後一個區塊可能存在少量空隙外，所有顯示記憶體都被高效利用。
2. 高效共享（Copy-on-Write）：在平行採樣（Parallel Sampling）或多輪對話中，多個請求可以共享同一組前綴（Prefix）的實體區塊。只有當某個請求需要修改內容時，才觸發 Copy-on-Write 操作複製該區塊。

工程實踐中的權衡

雖然 PagedAttention 大幅提升了吞吐量，但在實際部署中仍需關注以下維度：
- Block Size 的選擇：過小的 Block 會增加映射表的開銷；過大的 Block 則會重新引入內部碎片化。通常 16 或 32 是平衡點。
- 量化壓縮：為了進一步降低壓力，業界採用了 FP8 或 INT8 量化 KV Cache $\text{KV_quant}$，甚至使用 GQA（Grouped Query Attention）來減少 K 和 V 頭數。

總結

KV Cache 的演進路徑是從「簡單快取」 $\rightarrow$ 「靜態預配置」 $\rightarrow$ 「動態分頁管理」。PagedAttention 不僅是演算法的最佳化，更是對底層硬體資源管理的一次深刻重構。對於建構高效能 AI 系統而言，理解並最佳化記憶體佈局比單純追求算力堆疊更為關鍵。