大型語言模型在推論階段的真正瓶頸不是算力,而是記憶體。每產生一個 token,模型都要讀取先前所有 token 的 Key 與 Value 向量,也就是 KV Cache;上下文越長、批次越大,這塊快取就膨脹得越誇張,常常在你還沒撞到 GPU 算力上限前,就先把顯存塞爆。今天把推論階段最關鍵的三條優化路線一次講清楚:讓注意力運算貼著硬體跑的 FlashAttention、把顯存當作業系統分頁管理的 PagedAttention/vLLM,以及用量化把 KV Cache 直接壓小的做法。
📖 學(核心)
KV Cache 為什麼是長上下文的真正天花板
自迴歸生成的每一步,理論上都要重算整段序列的注意力,但重算太浪費,所以我們把每個 token 算過的 Key/Value 存起來重複使用,這就是 KV Cache。它的大小是「層數 × 2(K 與 V)× 序列長度 × 注意力頭數 × 每頭維度 × 每個元素的位元組數 × 批次大小」的乘積,關鍵是它跟序列長度成正比。以一個中型模型、單一請求跑到 128K 上下文為例,光是 KV Cache 就可能吃掉數十 GB 顯存,遠遠超過模型權重本身在推論時佔的空間。
這帶來兩個直接後果。第一,顯存決定了你能同時服務多少請求(批次大小)與多長的上下文,而不是浮點運算能力。第二,解碼階段是「記憶體頻寬受限」的:每產生一個 token 都要把整份 KV Cache 從顯存搬到運算單元,搬運時間主導了延遲,算術單元反而常常在等資料。理解這兩點,才會明白後面所有優化其實都在做同一件事——減少要搬的資料量,或讓搬運與運算重疊得更漂亮。
架構層面已經先做了一輪壓縮。MQA(Multi-Query Attention)讓所有注意力頭共用一組 K/V;GQA(Grouped-Query Attention)則折衷成幾組共享,是目前多數開源模型的標配;DeepSeek 的 MLA(Multi-head Latent Attention)更進一步,把 K/V 壓成一個低維潛在向量再還原,等於在快取層動了外科手術。這些都是在「模型設計時」就把 KV Cache 縮小,和後面講的推論引擎優化互為表裡。
FlashAttention:讓注意力貼著記憶體階層跑
FlashAttention 的核心洞見是:標準注意力會把 N×N 的注意力分數矩陣完整寫進 GPU 的高頻寬記憶體(HBM)再讀回來,而這來回搬運才是真正的瓶頸,不是矩陣乘法本身。它改用 tiling(分塊)搭配 online softmax,把序列切成小塊,在晶片上速度極快但容量極小的 SRAM 裡邊算邊累加,全程不把完整分數矩陣落地到 HBM。結果是記憶體用量從跟序列長度平方成正比降成線性,長序列訓練與推論的速度和顯存都大幅改善,而且數值上是精確的,不是近似。
這條線一路演進到硬體特化。FlashAttention-3 針對 NVIDIA Hopper 架構,約可達 740 TFLOPs/s、75% 硬體利用率;2026 年的 FlashAttention-4 則專為 Blackwell(B200/GB200)重寫執行模型,官方數字約 1,613 TFLOPs/s、71% 利用率,相較 cuDNN 9.13 有最高約 1.3 倍、相較 Triton 最高約 2.7 倍的加速,增益在 4K 以上的長序列最明顯。要留意的是,這種極致優化和特定 GPU 世代深度綁定,換硬體就得換 kernel;NVIDIA 也已把部分 FA4 技巧併入新版 cuDNN,讓不同實作的效能差距慢慢收斂。
PagedAttention 與 vLLM:把顯存當作業系統來管
PagedAttention 借用了作業系統虛擬記憶體的分頁概念。傳統做法會為每個請求預留一段連續顯存來放 KV Cache,但你事先不知道會生成多長,只能抓最壞情況,於是產生大量「內部碎片」與「保留浪費」——傳統推論常有 60% 到 80% 的 KV 顯存被浪費掉。PagedAttention 改成把 KV Cache 切成固定大小的區塊(block),用一張區塊表把邏輯位置對應到實際不連續的顯存頁,用多少配多少,vLLM 官方數據把浪費壓到 4% 以下,直接換來 2 到 4 倍的吞吐量。
分頁還解鎖了「共享」這件事。多個請求若有相同前綴(例如同一段 system prompt,或束搜尋的多條分支),可以共用同一批 KV 區塊,這就是 vLLM 的 Automatic Prefix Caching(APC):自動偵測跨請求的共同前綴並共享區塊,對固定系統提示、重複脈絡的應用能省下可觀顯存與首 token 延遲。這套做法已經是產業標配,vLLM、TGI、TensorRT-LLM 都採用。不過它也有代價:注意力 kernel 必須改寫成能處理不連續記憶體,軟體複雜度上升,因此近期出現像 vAttention 這類主張「保留連續虛擬位址、把分頁交給底層驅動」的替代路線,想在效能與工程簡潔間找新的平衡點。
KV Cache 量化:直接把每個數字變小
前面兩條路都沒有動快取裡數字的精度,量化則是正面處理這件事:把原本 FP16 的 K/V 存成更低位元的格式。FP8/INT8 相較 FP16 直接省一半顯存,通常是第一個該試的壓縮格式,因為數值行為夠接近、多數工作負載幾乎無感;INT4 再往下砍到約 75% 的節省。NVIDIA 在 Blackwell 上推的 NVFP4 KV Cache 又比 FP8 再省一半,等於在相同顯存下把可用上下文長度與批次大小翻倍,官方宣稱在程式生成與長上下文任務上精度損失低於 1%。
量化的難點在於「不是每個值都一樣好壓」。像 MLA 這種架構會產生數值異質性:被壓縮的內容分量相對耐壓,但和位置編碼(RoPE)綁定的分量對精度極度敏感,得用感知 RoPE 的分 token 量化策略分開處理。研究前沿已經推到 2-bit 甚至 1-bit(如 MiniKV、CSR),搭配 query-aware 的混合精度(重要的 token 保留高精度、其餘壓到底)。實務心法很單純:先上 FP8 幾乎穩賺,要更激進就分層、分 token 做混合精度,並務必用你自己的長上下文任務實測精度,不要只看 perplexity。
🧠 記
- 解碼階段是記憶體頻寬受限,不是算力受限;每個 token 都要把整份 KV Cache 從 HBM 搬一次,搬運主導延遲。
- KV Cache 大小與序列長度成正比,長上下文下常比模型權重更吃顯存,是能同時服務多少請求的真正上限。
- FlashAttention 靠 tiling + online softmax,不把 N×N 分數矩陣落地 HBM,記憶體從平方降成線性,且結果精確非近似。
- FlashAttention 高度綁硬體世代:FA3 對 Hopper、FA4 對 Blackwell,約 1,613 TFLOPs/s,長序列增益最大。
- PagedAttention 用分頁把 KV 顯存浪費從 60–80% 壓到 4% 以下,吞吐量提升 2–4 倍;APC 讓相同前綴跨請求共享區塊。
- 架構層壓縮 MQA/GQA/MLA 在「設計時」縮小 KV,和引擎層優化互補;GQA 是開源標配,MLA 壓成低維潛在向量。
- KV 量化:FP8 省一半幾乎無痛先上,INT4/NVFP4 更激進但要分層、分 token 混合精度並實測精度。
✍️ 實踐
- 用 vLLM 起一個服務,分別設定不同的
gpu_memory_utilization與max_model_len,觀察 KV Cache 區塊數與可容納的並行請求數如何變動,親手感受顯存即上限。 - 對同一模型開關 Automatic Prefix Caching,用一批共享同一段長 system prompt 的請求測 TTFT(首 token 延遲)與吞吐量,量化前綴共享的實際效益。
- 手算一次你目標模型在 32K / 128K 上下文、批次 1 與 16 下的 KV Cache 顯存,對照 GPU 容量,找出你真正的瓶頸落在哪。
- 把 KV Cache 精度從 FP16 切到 FP8,跑你自己的長上下文檢索或程式生成任務,記錄顯存節省與精度變化,驗證「先上 FP8」這條心法。
- 比較 MHA、GQA、MLA 三種架構的模型在相同上下文長度下的顯存佔用,理解架構層壓縮和引擎層優化如何疊加。
🔗 延伸學習
- We reverse-engineered Flash Attention 4
- Efficient Memory Management for LLM Serving with PagedAttention(原始論文)
- Optimizing Inference for Long Context with NVFP4 KV Cache(NVIDIA)
- Introduction to vLLM and PagedAttention
💬 問 AI
我要在 {GPU 型號與顯存} 上用 vLLM 部署 {模型名稱},服務 {並行請求數} 個並行請求、上下文長度 {上下文長度}。請幫我:
1. 估算此設定下的 KV Cache 顯存需求,判斷瓶頸是顯存還是算力;
2. 建議 gpu_memory_utilization、max_model_len、以及是否開啟 Automatic Prefix Caching;
3. 評估把 KV Cache 量化成 FP8 或 INT4 的取捨,並針對我的任務 {任務類型} 給出精度風險提醒與實測方式。