什么是鍵值緩存？讓 LLM 閃電般快速

一、為什么 LLMs 需要 KV 緩存？

大語言模型（LLMs）的文本生成遵循 “自回歸” 模式 —— 每次僅輸出一個 token（如詞語、字符或子詞），再將該 token 與歷史序列拼接，作為下一輪輸入，直到生成完整文本。這種模式的核心計算成本集中在注意力機制上：每個 token 的輸出都依賴于它與所有歷史 token 的關聯，而注意力機制的計算復雜度會隨序列長度增長而急劇上升。

以生成一個長度為 n 的序列為例，若不做優化，每生成第 m 個 token 時，模型需要重新計算前 m 個 token 的 “查詢（Q）、鍵（K）、值（V）” 矩陣，導致重復計算量隨 m 的增長呈平方級增加（時間復雜度 O (n2)）。當 n 達到數千（如長文本生成），這種重復計算會讓推理速度變得極慢。KV 緩存（Key-Value Caching）正是為解決這一問題而生 —— 通過 “緩存” 歷史計算的 K 和 V，避免重復計算，將推理效率提升數倍，成為 LLMs 實現實時交互的核心技術之一。

二、注意力機制：KV 緩存優化的 “靶心”

要理解 KV 緩存的作用，需先明確注意力機制的計算邏輯。在 Transformer 架構中，注意力機制的核心公式為：

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

其中：

Q（查詢矩陣）：維度為 $(n \times d_k)$ ，代表當前 token 對 “需要關注什么” 的查詢；
K（鍵矩陣）：維度為 $(n \times d_k)$ ，代表歷史 token 的 “特征標識”；
V（值矩陣）：維度為 $(n \times d_v)$ ，代表歷史 token 的 “特征值”（通常 $d_v = d_k$ ）；
$d_k$ 是Q和K的維度（由模型維度 $d_{\text{model}}$ 和注意力頭數決定，如 $d_k = \frac{d_{\text{model}}}{\text{num\_heads}}$ ）；
$QK^T$ 會生成一個 $(n \times n)$ 的注意力分數矩陣，描述每個 token 與其他所有 token 的關聯強度；
經過 softmax 歸一化后與V相乘，最終得到每個 token 的注意力輸出（維度 $(n \times d_v)$ ）。

三、KV 緩存的核心原理：“記住” 歷史，避免重復計算

自回歸生成的痛點在于：每輪生成新 token 時，歷史 token 的 K 和 V 會被重復計算。例如：

生成第 3 個 token 時，輸入序列是 $[t_1, t_2]$ ，已計算過 $t_1$ 和 $t_2$ 的 $K_1, K_2$ 與 $V_1, V_2$ ；
生成第 4 個 token 時，輸入序列變為 $[t_1, t_2, t_3]$ ，若不優化，模型會重新計算 $t_1, t_2, t_3$ 的K和V—— 其中 $t_1, t_2$ 的K、V與上一輪完全相同，屬于無效重復。

KV 緩存的解決方案極其直接：

緩存歷史 K 和 V：每生成一個新 token 后，將其K和V存入緩存，與歷史緩存的K、V拼接；
僅計算新 token 的 K 和 V：下一輪生成時，無需重新計算所有 token 的K、V，只需為新 token 計算 $K_{\text{new}}$ 和 $V_{\text{new}}$ ，再與緩存拼接，直接用于注意力計算。

這一過程將每輪迭代的計算量從 “重新計算 n 個 token 的 K、V” 減少到 “計算 1 個新 token 的 K、V”，時間復雜度從O(n2)優化為接近O(n)，尤其在生成長文本時，效率提升會非常顯著。

四、代碼實現：從 “無緩存” 到 “有緩存” 的對比

以下用 PyTorch 代碼模擬單頭注意力機制，直觀展示 KV 緩存的作用（假設模型維度 $d_{\text{model}}=64$ ， $d_k=64$ ）：

import torch
import torch.nn.functional as F# 1. 定義基礎參數與注意力函數
d_model = 64  # 模型維度
d_k = d_model  # 單頭注意力中Q、K的維度
batch_size = 1  # 批量大小def scaled_dot_product_attention(Q, K, V):"""計算縮放點積注意力"""# 步驟1：計算注意力分數 (n×d_k) @ (d_k×n) → (n×n)scores = torch.matmul(Q, K.transpose(-2, -1))  # 轉置K的最后兩維，實現矩陣乘法scores = scores / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))  # 縮放# 步驟2：softmax歸一化，得到注意力權重 (n×n)attn_weights = F.softmax(scores, dim=-1)  # 沿最后一維歸一化# 步驟3：加權求和 (n×n) @ (n×d_k) → (n×d_k)output = torch.matmul(attn_weights, V)return output, attn_weights# 2. 模擬輸入數據：歷史序列與新token
# 歷史序列（已生成3個token）的嵌入向量：shape=(batch_size, seq_len, d_model)
prev_embeds = torch.randn(batch_size, 3, d_model)  # 1×3×64
# 新生成的第4個token的嵌入向量：shape=(1, 1, 64)
new_embed = torch.randn(batch_size, 1, d_model)# 3. 模型中用于計算K、V的權重矩陣（假設已訓練好）
Wk = torch.randn(d_model, d_k)  # 用于從嵌入向量映射到K：64×64
Wv = torch.randn(d_model, d_k)  # 用于從嵌入向量映射到V：64×64# 場景1：無KV緩存——重復計算所有token的K、V
full_embeds_no_cache = torch.cat([prev_embeds, new_embed], dim=1)  # 拼接為1×4×64
# 重新計算4個token的K和V（包含前3個的重復計算）
K_no_cache = torch.matmul(full_embeds_no_cache, Wk)  # 1×4×64（前3個與歷史重復）
V_no_cache = torch.matmul(full_embeds_no_cache, Wv)  # 1×4×64（前3個與歷史重復）
# 計算注意力（Q使用當前序列的嵌入向量，此處簡化為與K相同）
output_no_cache, _ = scaled_dot_product_attention(K_no_cache, K_no_cache, V_no_cache)# 場景2：有KV緩存——僅計算新token的K、V，復用歷史緩存
# 緩存前3個token的K、V（上一輪已計算，無需重復）
K_cache = torch.matmul(prev_embeds, Wk)  # 1×3×64（歷史緩存）
V_cache = torch.matmul(prev_embeds, Wv)  # 1×3×64（歷史緩存）# 僅計算新token的K、V
new_K = torch.matmul(new_embed, Wk)  # 1×1×64（新計算）
new_V = torch.matmul(new_embed, Wv)  # 1×1×64（新計算）# 拼接緩存與新K、V，得到完整的K、V矩陣（與無緩存時結果一致）
K_with_cache = torch.cat([K_cache, new_K], dim=1)  # 1×4×64
V_with_cache = torch.cat([V_cache, new_V], dim=1)  # 1×4×64# 計算注意力（結果與無緩存完全相同，但計算量減少）
output_with_cache, _ = scaled_dot_product_attention(K_with_cache, K_with_cache, V_with_cache)# 驗證：兩種方式的輸出是否一致（誤差在浮點精度范圍內）
print(torch.allclose(output_no_cache, output_with_cache, atol=1e-6))  # 輸出：True

代碼中，“有緩存” 模式通過復用前 3 個 token 的 K、V，僅計算新 token 的 K、V，就得到了與 “無緩存” 模式完全一致的結果，但計算量減少了 3/4（對于 4 個 token 的序列）。當序列長度增至 1000，這種優化會讓每輪迭代的計算量從 1000 次矩陣乘法減少到 1 次，效率提升極其顯著。

五、權衡：內存與性能的平衡

KV 緩存雖能提升速度，但需面對 “內存占用隨序列長度線性增長” 的問題：

緩存的 K 和 V 矩陣維度為 $(n \times d_k)$ ，當序列長度 n 達到 10000，且 $d_k=64$ 時，單頭注意力的緩存大小約為 $10000 \times 64 \times 2$ （K 和 V 各一份） $=1.28 \times 10^6$ 個參數，若模型有 12 個注意力頭，總緩存會增至約 150 萬參數，對顯存（尤其是 GPU）是不小的壓力。

為解決這一問題，實際應用中會采用以下優化策略：

滑動窗口緩存：僅保留最近的k個 token 的 K、V（如 k=2048），超過長度則丟棄最早的緩存，適用于對長距離依賴要求不高的場景；
動態緩存管理：根據輸入序列長度自動調整緩存策略，在短序列時全量緩存，長序列時啟用滑動窗口；
量化緩存：將 K、V 從 32 位浮點（float32）量化為 16 位（float16）或 8 位（int8），以犧牲少量精度換取內存節省，目前主流 LLMs（如 GPT-3、LLaMA）均采用此方案。

六、實際應用：KV 緩存如何支撐 LLMs 的實時交互？

在實際部署中，KV 緩存是 LLMs 實現 “秒級響應” 的關鍵。例如：

聊天機器人（如 ChatGPT）生成每句話時，通過 KV 緩存避免重復計算歷史對話的 K、V，讓長對話仍能保持流暢響應；
代碼生成工具（如 GitHub Copilot）在補全長代碼時，緩存已輸入的代碼 token 的 K、V，確保補全速度與輸入長度無關；
語音轉文本實時生成（如實時字幕）中，KV 緩存能讓模型隨語音輸入逐詞生成文本，延遲控制在數百毫秒內。

可以說，沒有 KV 緩存，當前 LLMs 的 “實時交互” 體驗幾乎無法實現 —— 它是平衡模型性能與推理效率的 “隱形支柱”。

總結

KV 緩存通過復用歷史 token 的 K 和 V 矩陣，從根本上解決了 LLMs 自回歸生成中的重復計算問題，將時間復雜度從O(n2)優化為接近O(n)。其核心邏輯簡單卻高效：“記住已經算過的，只算新的”。盡管需要在內存與性能間做權衡，但通過滑動窗口、量化等策略，KV 緩存已成為現代 LLMs 推理不可或缺的技術，支撐著從聊天機器人到代碼生成的各類實時交互場景。