昨天討論了 KV Cache 的目的以及怎麼優化，盡可能壓縮想保存下來的資訊，讓 KV Cache 需要的 GPU 顯存空間可以越少越好。所以我們昨天總歸一句是在討論 「要存什麼？」，探討的主要都是模型架構層面，資料結構與演算法可以怎麼優化。接著我們今天要探討的是，那我們在硬體的 GPU 顯存裡，具體 「要怎麼存？」

一、什麼是「推論引擎」？

簡單來說，「模型」只是一份訓練好的靜態神經網路權重；而「推論引擎」則是負責驅動模型運算、管理硬體資源的執行系統（Runtime）（例如：vLLM，或是在本地跑 Ollama 時底層運作的 llama.cpp）。它負責接收並發請求、調度 GPU 計算核心，以及管理 GPU 顯存。

二、連續預分配引發的碎片問題

在 vLLM 出現之前，早期許多推論框架在管理 KV Cache 時，做法非常直覺：為每一個進來的請求，在 GPU 顯存裡劃出一整塊連續的記憶體空間。

但大語言模型在生成文字時，很難預測使用者這次請求會聊多長，模型會吐多少字。 使用者可能只是打個招呼，模型回了 10 個字就結束了；也可能丟了一篇論文，要求模型輸出 2000 字的深度摘要。

為了解決這種長度未知、隨時間動態增加的問題，傳統系統為了防止生成到一半顯存不夠導致 OOM，通常會採用保守的策略：直接按照模型支援的最大長度，預先替每個請求分配一整塊連續空間。

這就引發了幾個顯存空間的浪費問題，主要來自兩種碎片（Fragmentation，我沒想到會在這裡看到這幾個詞）：

1. 內部碎片（Internal Fragmentation）：
   假設最大長度是 2048 個 Token，伺服器會預先幫請求申請 2,048 個 Token 的空間，結果模型生成了 80 個字就吐出結束符號。剩下的 1,960 多個 Token 空間全部空在那裡，別的請求沒辦法使用，因此白白佔用一塊空間。
2. 外部碎片（External Fragmentation）：
   伺服器同時跑著很多長短不一的請求，有些剛結束、有些剛進來，顯存被切得很散。雖然把零碎的空閒空間加總起來可能夠用，但因為找不出一塊「連續」且夠大的區塊，所以新的請求只能卡在那裡乾等。

根據加州大學柏克萊分校團隊在 vLLM 論文 中的實測統計，在傳統的連續預分配機制下，推論系統中實際真正拿來存有效 Token 的顯存，往往只有 20% ~ 40%，其餘高達 60% ~ 80% 的顯存全部浪費在預留與碎片上。

三、PagedAttention

為了解決顯存的碎片問題，柏克萊團隊提出了解法：PagedAttention。到這裡大概也有點感覺，這跟作業系統 RAM 的碎片問題很像，因此這個方法就是參考了 Virtual Memory 與 Paging 的機制。

在作業系統中，程式的角度會以為自己拿到的是一段連續的虛擬記憶體，但實際上底層硬體會把空間切成一塊塊固定大小的 Page（例如 4 KB），散落在實體 RAM 的各處，中間靠一張 Page Table 進行映射。因此程式不需依賴「連續實體記憶體的分配」，而記憶體也只需要以 Page 為單位配置空間就好。

PagedAttention 就是把這個概念直接搬到 GPU 顯存裡：

• 將 KV Cache 切成固定大小的 Block：
  它不再為每個請求保留一整塊連續記憶體，而是把 KV Cache 劃分成固定大小的「物理區塊（Physical Blocks）」，通常一個 Block 裡面放固定數量（例如 16 或 32 個）的 Token。
• 邏輯連續，物理離散：
  對模型來說，在進行 Attention 運算時，看到的是邏輯（Logical Blocks）上連續的 Token 序列（Token 0, 1, 2, ...）；但在實體記憶體上，這些 Block 可能被離散存放在 GPU 顯存中任何有空位的地方。
• Block Table 動態映射：
  推論引擎內部維護一張 Block Table。請求剛進來時，只需要分配剛好放得下輸入的幾個 Block；Decode 階段每多生成一個 Block 能承載的 Token 數，系統才去顯存池裡要一個新的 Block 來用，用多少給多少。

四、PagedAttention 附帶好處

在很多應用場景裡，我們會想要模型對同一個輸入給出多種結果，例如 Antigravity 的 /fork 這個 Slash Command；或者是很多對話共同使用同一份 System Prompt。PagedAttention 可以容易做到這種跨請求的顯存共享（Memory Sharing）。

在傳統的連續記憶體做法下，因為每個請求都有自己專屬的一整條連續顯存，哪怕前面這 N 個分支的輸入 Prompt 完全長得一模一樣，也會硬生生把這份 Prompt 的 KV Cache 在顯存裡複製 N 份。但在 PagedAttention 裡，因為有了 Block Table 這一層間接引用，這件事變得格外漂亮：

1. Physical Block 共享

當使用者要求對同一個 Prompt 生成多個分支（例如下圖中的 Sample A1 與 Sample A2）時，系統在初始階段只需要存一份物理 Block。兩個分支的邏輯 Block Table，通通指向相同的實體 Block，並在系統內部把這些 Block 的 Reference Count 標記為 2。

2. Copy-on-Write (CoW)

在 Decode 階段，對於前面那些已經裝滿 Prompt 的歷史 Block（如圖中的 Block 7），因為大家都在讀取它，所以這些已滿的 Block 會繼續保持共享，引用計數也完全不會變。

但在 Prompt 最後那個 還沒裝滿的 Block 要寫入新字時，就會觸發 Copy-on-Write 的機制。就像下圖中的 Block 1 原本裝了 years ago our，還剩最後一格：

• 分支 A2 寫入了 mothers（綠色）；
• 分支 A1 想要寫入不同的字 fathers（黃色）；
• 為了避免兩邊打架，系統替分支 A1 分配一個全新獨立的 Block 3，把原本那幾個字複製過去並寫入 fathers。
• 這時，只有那個舊 Block 1 的引用計數減 1（Ref count 從 2 變成 1），因為分支 A1 換去指向新 Block 3 了。
• 而前面早已填滿的歷史 Block 7，依然共同指向同一塊實體顯存，完全不受影響。

到這裡可以感覺到，前人們針對 KV Cache 做了非常多努力，雖然可以想像 KV Cache 還是會佔用不少 GPU 顯存空間，但這幾套組合拳搭在一起後，現實我們所使用的 AI 服務的 KV Cache 機制已經優化非常多了，所需空間也已經小非常多了。這就是為什麼現在的 AI 服務商有能力處理這麼大量的需求，不會 GPU 顯存撐不住的原因（當然也是因為鈔能力 XDD）。明天繼續來聊聊在實際應用中，KV Cache 對我們帶來的影響吧。

參考資料

• Woosuk Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023.
• vLLM Team. (2023). vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention.