https://medium.com/@prathamgrover777/kv-caching-attention-optimization-from-o-n%C2%B2-to-o-n-8b605f0d4072

我們見過 LLM 如何逐字逐句地敲出上千字的回答，仿佛“邊想邊說(shuō)”。表面順滑，背后卻低效得驚人。

在生成第 t 步時(shí)，模型必須確保下一個(gè)詞與之前所有內(nèi)容保持一致：你的提示、已生成的部分、系統(tǒng)指令，乃至任何隱藏上下文。實(shí)際上，模型會(huì)重新通過所有 Transformer 層重建之前全部的隱藏狀態(tài)，并再次計(jì)算 Query、Key 和 Value，即便前面的 token 絲毫未變。這種重復(fù)計(jì)算逐層、逐頭進(jìn)行，沒有任何復(fù)用。如圖 1 所示，由于之前步驟的結(jié)果未被復(fù)用，每個(gè) token 的計(jì)算量隨序列長(zhǎng)度持續(xù)增加。

圖 1. 解碼階段每個(gè) token 的計(jì)算成本（樸素 vs KV 緩存）。

為什么這種方式擴(kuò)展性這么差？想象你在寫一段文字，每添加一個(gè)新句子之前，都要從頭把整篇文檔重新讀一遍。然后寫下一個(gè)句子時(shí)，再?gòu)念^讀一遍。如此反復(fù)。這就是樸素解碼循環(huán)在做的事情。

現(xiàn)在把它放進(jìn) transformer：

• L= 層數(shù)

• H= 每層注意力頭數(shù)

• n= 當(dāng)前序列長(zhǎng)度

• 因此，每個(gè) token 的計(jì)算量不僅與序列長(zhǎng)度 n 成正比，內(nèi)部還要再乘以 L × H。這就是成本飆升的原因。

現(xiàn)在，把這個(gè)問題放到現(xiàn)代 LLMs 的規(guī)模下：

• 幾十層網(wǎng)絡(luò)層層堆疊

• 每一層有多個(gè)注意力頭，各自“角度”不同地回顧歷史。

• 長(zhǎng)提示（上千條），還要生成一長(zhǎng)段回答。

更致命的是，模型對(duì)過去的 Key 和 Value 毫無(wú)記憶；因此在第 1000 個(gè) token 時(shí)，它得把 token 1 到 999 全部重新算一遍。這就導(dǎo)致解碼的時(shí)間復(fù)雜度是 **O(n2)**。

結(jié)果就是巨大的冗余。

圖 2. 無(wú) KV 緩存的樸素解碼

如圖 2 所示，模型不會(huì)記住過去的 K/V 投影，而是對(duì)每個(gè)新 token 都重新計(jì)算它們。 當(dāng) t = 2 時(shí) → 重新計(jì)算 token **1**當(dāng) t = 3 時(shí) → 重新計(jì)算 token 1 和 **2**對(duì)于 t = 4 → 重新計(jì)算 tokens **1, 2, 3**對(duì)于 t = n → 重新計(jì)算 tokens 1 到 n?1

在每一層內(nèi)部，對(duì)于每個(gè) token，它執(zhí)行：

• K = X · W?

• V = X · W?

• Q · K?（與所有先前 key 的點(diǎn)積）

• 對(duì)所有先前 Value 的加權(quán)和。

這些結(jié)果從未被復(fù)用，因此到第 n 個(gè) token 時(shí)，你已經(jīng)重復(fù)了(n–1) × L × H次！這意味著計(jì)算量持續(xù)增長(zhǎng)， 永遠(yuǎn)不會(huì)穩(wěn)定 。

硬件瓶頸

在解決它之前，我們必須先理解瓶頸所在：搬運(yùn)數(shù)據(jù)代價(jià)高昂。GPU 做數(shù)學(xué)運(yùn)算極快，但真正的成本往往在于把正確的數(shù)據(jù)在正確的時(shí)間送到正確的位置。

舉個(gè)簡(jiǎn)單的例子，把它想象成廚房里的廚師：

• GPU 核心就是那位廚師，切菜和烹飪的速度快得驚人。

• VRAM（GPU 自有內(nèi)存）就像緊挨著廚師的小備餐臺(tái)，空間有限，但伸手就能拿到上面的東西。

• 系統(tǒng) RAM則是走廊盡頭的大儲(chǔ)藏室，空間充足，可每次都得停下、走過去再把食材搬回來(lái)。

用 ML 的話說(shuō)：

• GPU 核心每秒能執(zhí)行數(shù)萬(wàn)億次FLOPs，真正的瓶頸并不是算力。

• 但內(nèi)存帶寬（將 K/V 張量搬運(yùn)到高帶寬內(nèi)存）是有限的。

• 并且注意力機(jī)制需要反復(fù)讀取這些張量，這會(huì)把帶寬壓垮。

這就是為什么序列變長(zhǎng)時(shí)生成會(huì)變慢——不是因?yàn)橛?jì)算變復(fù)雜，而是因?yàn)?GPU等數(shù)據(jù)的時(shí)間比真正計(jì)算的時(shí)間還多。

模型的權(quán)重和對(duì)話的歷史 token都必須放在顯存里才能快速處理。但在樸素的解碼中，每生成一個(gè)新 token，廚師（GPU）就得折回儲(chǔ)藏室（內(nèi)存）再拿一遍同樣的食材（歷史 K/V）。儲(chǔ)藏室與灶臺(tái)之間的路越來(lái)越擠，走路時(shí)間越來(lái)越多，真正炒菜的時(shí)間越來(lái)越少。

那種對(duì)舊數(shù)據(jù)持續(xù)、重復(fù)的抓取，就是我們所說(shuō)的帶寬之痛。“帶寬”指的是這條通路的容量，“痛”則源于把剛才已經(jīng)存在的數(shù)據(jù)再次塞進(jìn)去所造成的嚴(yán)重?fù)矶�。帶著這個(gè)概念，我們來(lái)看看到底在哪個(gè)環(huán)節(jié)，這一過程在規(guī)模擴(kuò)大時(shí)會(huì)崩潰。更具體地說(shuō)，注意力機(jī)制會(huì)變成受內(nèi)存限制，而非計(jì)算限制。你的 GPU 空轉(zhuǎn)著，等待從顯存取數(shù)，盡管它的算力足以瞬間完成運(yùn)算。

樸素的注意力機(jī)制在何處因規(guī)模而崩潰？

• 延遲悄然上升：早期的 token 反應(yīng)迅捷，但隨著對(duì)話變長(zhǎng)，每一個(gè)新 token 的生成時(shí)間都比前一個(gè)更久。你正拖著一段越來(lái)越長(zhǎng)的歷史前行。這正是注意力計(jì)算復(fù)雜度帶來(lái)的直接后果——它是二次方，即 **O(n2)**，其中n為序列長(zhǎng)度。

• 帶寬之痛，而非算力之痛：現(xiàn)代 GPU 的數(shù)學(xué)運(yùn)算飛快，但數(shù)據(jù)在內(nèi)存之間搬來(lái)搬去才是“征稅員”；反復(fù)把整個(gè)“過去”拖過總線，會(huì)把帶寬壓垮。你更多時(shí)間是在等數(shù)據(jù)，而不是在計(jì)算。

• 推理崩潰：在線上，你不是給一個(gè)人生成，而是同時(shí)給成千上萬(wàn)人生成。如果每條 token 流都重新處理自己的完整歷史，系統(tǒng)立刻垮掉，成本飆升。

此刻，一個(gè)自然的問題浮現(xiàn)：**“等等……我們?yōu)槭裁疵看味贾貜?fù)同樣的投影？”****KV 緩存**就是你拒絕重讀過去的瞬間。簡(jiǎn)單說(shuō)，KV 緩存就是：一旦某 token 在某層里的 Key 和 Value 向量算完，我們把它存進(jìn) GPU 內(nèi)存，而不是直接扔掉。

在注意力機(jī)制中，每個(gè) token 被轉(zhuǎn)換為兩個(gè)緊湊向量：Key（K）和Value（V），它們描述_該 token 應(yīng)如何與后續(xù) token 交互_ 。生成方式是將 token 的嵌入（x）通過該層的權(quán)重矩陣：

Key = X @ W?
Value = X @ W?

這并非只計(jì)算一次，而是對(duì)模型的每一層、每一個(gè)注意力頭，都各自計(jì)算一套 K 和 V。
舉例：若模型有 32 層、每層 32 個(gè)頭，則每個(gè) token 要計(jì)算 32 × 32 = 1024 組 K/V 投影。

訣竅在這里：

• 在推理過程中，模型的權(quán)重（W?、W?）不會(huì)改變。

• 一旦計(jì)算完成，token 的嵌入向量X也不再改變。

• 這意味著它的Key和Value向量是確定性的。一旦算出，它們就像被刻在石頭上一樣，在整個(gè)序列中不會(huì)變化。

那為何每一步都要重新計(jì)算它們呢？

于是，不再 ：每生成一個(gè)新 token → 為所有舊 token 重新計(jì)算 K/V。 而是： 一次性算出 K/V → 存起來(lái) → 后續(xù)所有 token 直接復(fù)用。

這就是KV 緩存。

底層到底發(fā)生了什么變化？當(dāng)模型生成第 t 個(gè) token 時(shí)，常規(guī)的注意力操作會(huì)這樣執(zhí)行：

在樸素解碼中，每當(dāng)時(shí)間步 t 增加，模型會(huì)重新計(jì)算所有層、所有頭之前的 K 和 V 向量，只為再次把它們代入這個(gè)方程。

啟用 KV 緩存后，方程本身_并未_改變，但K 和 V 的來(lái)源變了。

• 我們停止重新計(jì)算K?…K??? 和 V?…V???

• 取而代之的是，我們從 GPU 內(nèi)存中一個(gè)名為KV Cache的張量里直接讀取它們

因此，同樣的公式變?yōu)椋?/p>

模型不再在每一步重新計(jì)算 K?…K??? 和 V?…V???，而是直接從 GPU 內(nèi)存中讀取。在 token t 時(shí)唯一需要的新工作是計(jì)算 K? 和 V?，并將它們追加到緩存中。

KV 緩存在 GPU 內(nèi)部到底是什么樣子？

它既不是列表，也不是 Python 字典。在實(shí)際的 LLM 實(shí)現(xiàn)中（如 vLLM、TensorRT-LLM、Hugging Face），緩存以張量形式存儲(chǔ)在 GPU 顯存中，維度固定。

• 
  

  num_layers= 模型中 Transformer 塊的總數(shù)（例如 LLaMA-7B 為 32）

• num_heads= 每層注意力頭數(shù)（例如 32）

• seq_len= 當(dāng)前已見的 token 數(shù)量

• head_dim= 每個(gè)注意力頭的維度（例如 64 或 128）

每當(dāng)生成一個(gè)新 token：
→ 我們計(jì)算K?和V?
→ 我們沿著seq_len維度將它們追加到末尾
→ 其他所有內(nèi)容保持不變

我們?yōu)槭裁床痪彺?Query（Q）？

因?yàn)?strong>Q（Query 向量）與K和V在本質(zhì)上不同。

• Key 和 Value 代表 之前 token 的記憶 。
  它們一旦計(jì)算完成就不會(huì)再改變。

• 但 Query 只依賴于_當(dāng)前正在生成的 token_，而非過去的。
  Q?用來(lái)提問：“鑒于我已看到的全部（所有已緩存的K/V），下一個(gè) token 應(yīng)該是什么？”

所以：

• K 和 V = 記憶 → 緩存一次，反復(fù)使用

• Q = 按步驟、臨時(shí)生成 → 無(wú)需存儲(chǔ)

如果我們緩存了Q，就永遠(yuǎn)不會(huì)再用到它，因?yàn)樗辉诋?dāng)前時(shí)間步使用。緩存它只會(huì)白白浪費(fèi)內(nèi)存，毫無(wú)收益。

可視化差異

為了真正理解 KV 緩存的神奇之處，讓我們跟隨動(dòng)畫，看看我們的 LLM 如何生成短語(yǔ)“I Love cats”。我們將重點(diǎn)觀察模型如何處理這些 token，以預(yù)測(cè)序列中的下一個(gè)詞。（GIF 可能加載較慢——稍等片刻）

圖 3. 有無(wú) KV 緩存的對(duì)比：過去的 K/V 是重新計(jì)算 vs. 復(fù)用。 1. 無(wú)緩存

原始而浪費(fèi)的做法。每一步都必須從頭重新處理全部?jī)?nèi)容。

步驟 1：預(yù)測(cè) “Love”（歷史：“I”）

• 模型接收第一個(gè) token “I”。

• 它計(jì)算其 Key（記為K?）及其 Value（記為V?）。

• 它還計(jì)算其 Query（Q?）。

• 它執(zhí)行注意力計(jì)算（Q?關(guān)注K?），以預(yù)測(cè)下一個(gè)詞：“Love”。

• 然后，它會(huì)丟棄掉K?和V?，所有這些計(jì)算成果都被浪費(fèi)。

• 在張量層面，這意味著 GPU 剛剛計(jì)算出形狀為[num_heads, head_dim]的矩陣K?和V?，卻立即將它們丟棄。當(dāng)模型處理下一個(gè) token 時(shí)，會(huì)毫無(wú)必要地重新構(gòu)建這些相同的矩陣。

步驟 2：預(yù)測(cè)“cats”（歷史：“I Love”）

• 模型現(xiàn)在需要處理新 token“Love”，但它對(duì)“I”沒有任何記憶。

• 它必須重新為“I”計(jì)算 Key 和 Value（生成K’?和V’?）。

• 它還會(huì)為“Love”計(jì)算 Key 和 Value（生成K?和V?），并為新 token“Love”計(jì)算 Query（Q?）。

• 它執(zhí)行注意力計(jì)算（Q?同時(shí)查看K’?和K?）來(lái)預(yù)測(cè)“cats”。

• 隨后，它把K’?、V’?、K?和V?全部丟棄。

• 你看出規(guī)律了嗎？為了預(yù)測(cè)第三個(gè)詞，我們不得不把第一個(gè)詞的所有計(jì)算重新做一遍。

這樣，模型只需計(jì)算每個(gè) token 的 Key 和 Value 一次 ，并將其保存下來(lái)。

步驟 1：預(yù)測(cè) “Love”（歷史：“I”）

• 模型接收第一個(gè) token “I”。

• 它計(jì)算自己的 Key（K?）和 Value（V?）。

• 執(zhí)行注意力計(jì)算以預(yù)測(cè) “Love”。

• 關(guān)鍵是，它將K?和V?存入一塊特殊內(nèi)存：KV 緩存。

• 緩存現(xiàn)在包含：{(K?, V?)}。

• 這意味著 GPU 現(xiàn)在為每一層、每一個(gè)注意力頭都保存著一個(gè)小張量，代表該詞元對(duì)所有未來(lái)注意力查詢的貢獻(xiàn)。無(wú)需重新計(jì)算，只需查表即可。

步驟 2：預(yù)測(cè)“cats”（歷史：“I Love”）

• 模型忽略 “I” 詞元。

• 它知道它已經(jīng)存在于緩存中。

• 它只處理新的 token “Love”。

• 它只為 “Love” 計(jì)算 Key（K?）和 Value（V?）。

• 它將這對(duì)新的鍵值追加到緩存中。

• 緩存現(xiàn)在包含：_{(K?, V?), (K?, V?)}_ 它為“Love”計(jì)算查詢（Q?），并通過查看整個(gè)緩存（K?和K?）進(jìn)行注意力計(jì)算。

• 它預(yù)測(cè)“cats”。沒有任何內(nèi)容被丟棄。

• 此時(shí)工作量恒定且最小。這就是每一步的工作量變?yōu)榫�性（**O(n)**）的方式。

模型依舊會(huì)“回顧”歷史，只是它通過查詢一張簡(jiǎn)單的查找表（緩存）來(lái)完成，而無(wú)需從頭重新計(jì)算。注意力計(jì)算仍會(huì)將Q?與所有過去的鍵進(jìn)行比較，但由于過去的 K/V 已從緩存中取出，重算成本降至 **O(1)**。每一步的工作現(xiàn)在僅僅是將一個(gè)新查詢與已有鍵進(jìn)行比較。

這種優(yōu)化是現(xiàn)代推理服務(wù)器（如 vLLM、TensorRT-LLM 或 Hugging Face 的transformers庫(kù)）能夠?qū)崿F(xiàn)實(shí)時(shí)文本生成的原因。沒有它，每生成幾百個(gè) token，延遲就會(huì)翻倍，讓聊天模型根本無(wú)法使用。

老實(shí)說(shuō)，這并不像發(fā)明一種新算法，你只是不再做重復(fù)勞動(dòng)。這就是 KV 緩存的作用：過去不再重新計(jì)算，而是被記住。

權(quán)衡：用內(nèi)存換計(jì)算

圖 4. KV 緩存作為時(shí)空權(quán)衡

KV 緩存并非“免費(fèi)的午餐”，而是一種典型的時(shí)空權(quán)衡：我們消除了冗余計(jì)算（時(shí)間），但必須把緩存塞進(jìn) GPU 顯存（空間），而且這塊緩存會(huì)變得非常非常龐大！

KV 緩存的大小由以下幾個(gè)因素決定：

對(duì)于一個(gè)上下文窗口長(zhǎng)達(dá) 32,000 token 的大模型，這塊緩存就能吃掉幾十 GB 的寶貴 GPU 顯存。如果服務(wù)器還要同時(shí)服務(wù)這么多并發(fā)用戶（大批量），內(nèi)存需求就會(huì)成為主要瓶頸，直接限制系統(tǒng)容量。

極簡(jiǎn)代碼示例（PyTorch 樸素實(shí)現(xiàn) vs KV 緩存）

到目前為止，我們只_討論_了 KV 緩存�，F(xiàn)在讓我們通過代碼來(lái)_親眼看看它是如何工作的_ 。

下面的示例清楚地展示了樸素解碼如何一次又一次地重復(fù)計(jì)算所有內(nèi)容，以及 KV 緩存如何通過存儲(chǔ)先前計(jì)算好的鍵（K）和值（V）來(lái)避免這種浪費(fèi)。

import torch  
import torch.nn as nn  

 # A single multi-head attention layer  
attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)  

 # Dummy input sequence  
tokens = torch.randn(1, 5, 512)  # [batch, seq_len, embedding_dim]  

 # -------------------------------  
# 1. Naive decoding (no caching)  
# Recomputes attention over full history at every step  
# -------------------------------  
for t in range(1, tokens.size(1)):  
    x = tokens[:, :t, :]         # tokens from 1 to t  
    out, _ = attn(x, x, x)       # recompute Q,K,V for all past tokens again  

 # -------------------------------  
# 2. KV Caching (compute K/V once → reuse forever)  
# -------------------------------  
past_k, past_v = None, None

 for t in range(tokens.size(1)):  
    x = tokens[:, t:t+1, :]      # only the new token  

     # Project to Q, K, V (like attention does internally)  
    q = attn.in_proj_q(x)  
    k = attn.in_proj_k(x)  
    v = attn.in_proj_v(x)  

     # Save (or append) K/V into cache  
    past_k = k if past_k isNoneelse torch.cat([past_k, k], dim=1)  
    past_v = v if past_v isNoneelse torch.cat([past_v, v], dim=1)  

     # Attention now only compares new query with cached keys  
    attn_scores = torch.matmul(q, past_k.transpose(-1, -2)) / (k.size(-1) ** 0.5)  
    attn_probs = attn_scores.softmax(dim=-1)  
    output = torch.matmul(attn_probs, past_v)

 

故事基本上就是這樣。

KV 緩存并沒有讓注意力機(jī)制變得更智能，它只是讓它不再愚蠢。與其在每個(gè) token 上重新計(jì)算過去的內(nèi)容，我們只需計(jì)算一次并記住它。這就是為什么即使上下文變長(zhǎng)，生成速度依然很快。