第 8 章 — 下一代 KV 快取管理
《LLM Primer VI:擴展 AI 系統》章節走讀的第八篇。這一章把作業系統的分頁洞見搬進推論引擎 — 並把 KV 快取從一段預留的位元組 slab 變成一個共享、可逐出、可做 prefix cache 的資源。
這一章為什麼存在
第 7 章讓連續批次留著一筆債。序列每次迭代進入與離開;天真的佈局給每個位子分一段最大尺寸的 slab,大部分是浪費;批次處理的吞吐量勝利被還了回去一部分。這筆債是內部碎裂化,剛好就是作業系統在 1960 年代用分頁解決掉的失敗模式。第 8 章的主題是把那個解法套用到 LLM 服務上:把 KV 快取切成小的實體區塊,用一個 page table 解耦邏輯 token 位置,並讓逐出與快取策略在序列之間回收或共享區塊。PagedAttention 是奠基性的動作;H2O 與 InfiniGen 是逐出策略;prefix caching 是讓一個生產叢集能用十幾顆 GPU 就服務數百萬個代理式請求的技術。
8.1 PagedAttention 是 KV 快取的虛擬記憶體
vLLM 的 PagedAttention 論文(2023)把作業系統的洞見直接拉進引擎裡。KV 快取被切成固定大小的區塊 — 通常一個 16 個 token — 放在一個扁平的實體池裡。一個序列由一個 block table 表示:一個把邏輯位置映射到實體區塊 ID 的陣列。注意力 kernel 把 block table 當成額外輸入,並靠 indirection 而不是連續切片來蒐集 keys 與 values;在 Hopper 上 L2 對隨機存取樣式吸收得夠好,kernel 跑起來跟 slab 版本相差不到幾個百分點。收益很大。內部碎裂從 60 到 80% 掉到大約 4%(每個序列一個部分填滿的尾巴區塊),這讓可用並行度提升 2 到 4 倍。以參照計數為基礎的區塊共享讓複雜取樣幾乎不用錢 — 一段 2,000 個 token 提示詞的 best-of-8 從 16,800 個 token 區塊掉到 2,800 個 — 而且它就是 prefix caching 底下的基底。
8.2 H2O 與 InfiniGen 逐出不重要的 token
PagedAttention 解決碎裂化,卻沒解決在任何佈局下都撐爆 VRAM 的長上下文。一段 20 萬 token 的 Llama-3-70B 上下文,除了權重之外還要 60 GB 的 KV。H2O(Heavy Hitter Oracle)觀察到 decoding 期間的注意力權重會集中在一小組來源位置 — 最近的 token、最開頭的注意力匯 token、以及一小組跟內容相關的位置 — 而大部分歷史位置得到的權重基本上是零。引擎為每個位置追蹤一個累積注意力分數;當序列的 KV 預算逼近上限時,它就逐出分數最低的位置,只保留一個「一定會留」的最近與匯 token 視窗。節省很大;代價是「回不去」— 如果後面查詢用到一個已經被逐出的位置,引擎救不回來。InfiniGen 用可回復的動態選擇把取捨修得更細:不直接丟掉 token,而是把它們的 KV 卸載到 CPU 記憶體,如果注意力又開始集中在它們身上就把它們分頁回來。對的逐出策略取決於工作量重新查詢自己歷史的程度;代理式工作量會懲罰永久逐出、獎勵 InfiniGen 式的回復。
8.3 Prefix cache 是 PagedAttention 解鎖的最大收益槓桿
在真實流量裡,大多數提示詞的前幾千個 token 在請求之間都是一樣的。一個聊天服務對每個對話重用同一段系統提示詞。一個 RAG 服務會把同一批擷取到的段落貼到每個搜過同一件事的使用者的提示詞裡。一個代理在每一步注入同一份工具描述與推理鷹架。PagedAttention 讓共享變得機械化:把提示詞以區塊大小做雜湊;如果雜湊出現在全域的「已運算完的區塊」快取裡,就把新請求的 block table 指向那個快取的區塊,並整段跳過 prefill;如果沒有,prefill 就跑,產生的區塊會被登記進去。生產環境的命中率很驚人 — 對聊天服務的系統提示詞可以超過 99%、對 RAG 依賴擷取的前綴可以 30 到 60%、對代理式鷹架接近 1.0。SGLang 的 RadixAttention 用 radix tree 把想法推得更遠,能替任何長度的共享前綴當索引鍵,而不侷限於區塊對齊。Prefix cache 是最常把「預算超支」的服務叢集變成「預算有餘」的單一技術。
第 8 章鋪陳出來的東西
分頁、逐出、prefix cache 收斂了每個 token 的 KV 佔用,並讓引擎的記憶體行為在高並行下變得可解。這三者都沒解決的是 decoding 的根本循序依賴:每次迭代、每個序列一個輸出 token,不管有多少個位子在作用。第 9 章對付這個限制,靠推測解碼 — 用便宜的草擬提前預測好幾個 token、用一次昂貴的前向傳遞驗證那份猜測,替使用者感受到的延遲最重要的那些序列打破「每步一個 token」的地板。
下一篇 — 第 9 章:推測解碼。草擬、驗證,以及推測何時值得的算術。
