第 2 章 — KV 快取的挑戰
《LLM Primer VI:擴展 AI 系統》章節走讀的第二篇。這一章替那個會先於權重吃光每個服務系統 VRAM 的資料結構命名。
這一章為什麼存在
第 1 章的自回歸迴圈靠記憶來避免平方級的工作量。在每一次 decoding 步驟裡,模型都需要對之前所有 token 做注意力,而如果每次都從頭做,等於每次迭代都要重跑一次 prefill — 對一段完成長度而言就是 N 倍的乘數。key-value 快取就是省下這些算術的工作區。它在每一層儲存每個 token 的 K 與 V 投影,好讓後面的 token 只算自己的 K 與 V 並讀回之前的。這個快取是 decoding 保持在大約 O(N) 而不是 O(N²) 的原因。它同時也是我看過的每一個生產系統裡,VRAM 的最大單一消耗者。第 2 章走過主導它大小的公式、把它縮小的架構變體,以及會毀掉天真配置方式的碎裂化問題。
2.1 主導快取大小的公式
單一序列的 KV 記憶體是 2 × L × H_kv × D × S × bytes:兩個張量(K 與 V)、L 層、H_kv 個 key/value 頭、頭維度 D、序列長度 S、以及每個元素的精度位元組數。乘上批次大小就是機群層級的數字。每一個維度都是相乘的。Llama-3-70B 有 L=80、H_kv=8(GQA)、D=128,在 S=8,192 與 BF16 之下,每個序列大約握著 40 GB 的快取 — 還沒乘上批次。在批次 32 的情況下,光是快取就逼近 1.3 TB,是權重本身的好幾倍,也是一顆 H100 VRAM 的好幾倍。Llama-2-70B 這個最後一個主要的 MHA 模型,H_kv 是 64 而不是 8,同樣配置下它的快取要大八倍。這個公式不是個小趣聞;它是決定一顆晶片能容納多少同時在線使用者的數字。機器上其他每一個記憶體預算 — 啟動值、框架開銷、暫存緩衝區 — 都得繞著它塞。
2.2 MHA、GQA、MQA 是快取大小的設計選擇
H_kv 這個項,是現代架構在工程上壓下來的項。多頭注意力(MHA)給每個 query 頭一個 K 與 V 頭;快取很肥、品質最高,而且每個頭能專門化它自己對應注意什麼的概念。多查詢注意力(MQA)塌陷成一個共享的 K 與 V 頭;快取縮小 H 倍,但頭失去專門化,較長的上下文上可量測到的品質就會流失。分組查詢注意力(GQA)在 2023 年被提出,現在是 Llama-3、Mistral、Mixtral、Qwen、DeepSeek 的主流選擇,它把 query 頭切成共享 K 與 V 的 G 組;典型的 G = H/8 會帶來八倍的快取縮減,而品質成本小到評估套件無法可靠地看出來。GQA 不是白吃的午餐 — 它也削減了 K/V 投影的參數量,而那個預算會被重新分配到模型別的地方 — 但從經驗上這種重新分配鮮少造成傷害。MLA(multi-latent attention,DeepSeek-V2)靠一個低秩的 latent cache 更進一步;它是研究方向,但 2026 年的生產部署裡以 GQA 為主。
2.3 天真的配置會浪費掉大部分預算
配置 KV 記憶體的顯然做法 — 為每個序列預留一段連續的 slab,大小抓在最大可能長度 — 接觸真實流量就會失敗。大多數請求完成的長度遠遠不到上限;預留的尾巴是死掉的記憶體。在 Llama-3-70B 上,以 32K 為上限,一個 1,000 個 token 的完成會浪費掉 31 個 slab 份量的 KV 快取。批次 32、上限 8K、平均長度 800 的情況,任何時刻大約 90% 的預留 KV 快取都沒被用到。更糟的是,這種佈局沒有彈性:新來的請求沒辦法用另一個序列沒用完的尾巴,因為那個尾巴已經被承諾出去了。並行度會塌陷成最壞情況預算所允許的水準,而不是平均流量原本能撐起的水準。這個問題是內部碎裂化 — 作業系統在 1960 年代用分頁解決掉的同一種失敗模式。第 8 章會展示 PagedAttention 如何把這個解法搬過來。
第 2 章鋪陳出來的東西
每當後面的章節得推理一顆 GPU 塞得下多少使用者的時候,這個快取公式與它的維度都會再次浮出來。第 3 章用它解釋為什麼挑一顆服務 GPU 時 VRAM 容量比 FLOPs 還重要。第 7 章用它解釋為什麼連續批次的上限是被 KV 快取而不是被模型鎖住。第 8 章用它替 PagedAttention、H2O 逐出、prefix caching 辯護 — 這三項技術整個目的就是讓快取更像分頁式虛擬記憶體、更不像一塊預留的位元組 slab。這個公式就是這本書後面所有記憶體論點的文法。
下一篇 — 第 3 章:生成式 AI 的資料中心 GPU。那些必須同時裝下權重與 KV 快取的矽晶片,用機制而非規格表來讀。
