現代 AI 的「記憶體管理」：KV Cache 與 PagedAttention 的工程真相

在 LLM 的推論過程中，最昂貴的資源不是計算量（FLOPs），而是記憶體頻寬。當你與 AI 對話時，模型需要回顧之前所有的對話歷史。如果每次生成新 Token 都要重新計算一遍之前的所有 Key 和 Value 向量，推論速度將呈平方級下降。

為了解決這個問題，工業界引入了 KV Cache（鍵值快取）。但 KV Cache 帶來了另一個噩夢：記憶體碎片化。

KV Cache：用空間換時間的權衡

在 Transformer 的解碼階段，每個 Token 的生成依賴於之前所有 Token 的注意力權重。這意味著對於第 $n$ 個 Token，我們需要它與前 $n-1$ 個 Token 的 $\text{Key}$ 和 $\text{Value}$ 向量進行點積運算。

由於前 $n-1$ 個 Token 在之前的步驟中已經計算過了，且其結果在推論過程中保持不變，我們可以將這些 $\text{K}$ 和 $\text{V}$ 向量儲存在顯示記憶體中。這樣，每一步只需要計算當前 Token 的 $\text{K}$ 和 $\text{V}$ 並將其追加到快取中即可。

代價是顯示記憶體的劇增。對於一個 Llama-3-70B 模型，在 FP16 精度下，單請求的 KV Cache 佔用量隨序列長度線性增長。當併發請求增加時，顯示記憶體會被迅速填滿，導致 OOM（Out of Memory）。

靜態分配的困境：記憶體碎片化

早期的推論框架（如 HuggingFace Transformers）採用的是連續記憶體分配策略。為了保證效能，系統會為每個請求預先分配一塊足夠大的連續顯示記憶體空間（例如最大支援 4096 個 Token）。

這種做法導致了嚴重的記憶體浪費：
1. 內部碎片：如果使用者只輸入了 10 個 Token，剩下的 4086 個位置依然被佔用且無法給他人使用。
2. 外部碎片：隨著請求的動態增加和釋放，顯示記憶體中會出現大量不連續的小空隙，無法容納新的大請求。
3. 預留浪費：為了防止序列增長導致崩潰，開發者不得不保守地設定最大長度，進一步降低了吞吐量。

PagedAttention：借鑑作業系統的虛擬記憶體

vLLM 團隊提出的 PagedAttention 從根本上改變了這一現狀。其核心思想是將 LLM 的 KV Cache 類比為作業系統的虛擬記憶體頁（Virtual Memory Pages）。

核心機制

PagedAttention 不再要求 KV Cache 在實體顯示記憶體中連續儲存，而是將其劃分為固定大小的區塊（Blocks）。
- 邏輯區塊 $\rightarrow$ 實體區塊：模型在邏輯上看到的是連續的序列，但底層透過一個「區塊表」（Block Table）將邏輯索引映射到不連續的實體顯示記憶體區塊中。
- 動態按需分配：只有當目前的實體區塊填滿時，系統才會為該請求分配一個新的實體區塊。
- 高效共享：在平行採樣（Parallel Sampling）或 Beam Search 時，多個輸出序列可以共享同一個前綴（Prompt）的實體區塊，無需重複儲存相同的 KV 資料。

工程效果

透過 PagedAttention， KV Cache 的記憶體利用率從之前的 $60\% \sim 80\%$ 提升到了接近 $96\%$ 以上。這意味著在同樣的硬體條件下，單機能夠承載的併發請求數（Batch Size）提升了數倍，直接降低了推論成本並提高了回應速度。

總結：從演算法到工程的閉環

KV Cache 是對 Transformer 計算冗餘的優化，而 PagedAttention 是對 KV Cache 儲存冗餘的優化。這揭示了現代 AI 系統的一個關鍵趨勢：模型能力的上限由演算法決定，但商業落地的下限由工程實作決定。當我們討論「推論加速」時，關注點已從單純的算子優化轉向了更深层的資源調度與記憶體管理方案。