大語言模型（LLM）在生成文本時，通常是一個 token 一個 token 地進行。每當模型生成一個新的 token，它就會把這個 token 加入輸入序列，作為下一步預測下一個 token 的依據。這一過程不斷重復，直到完成整個輸出。
然而，這種逐詞生成的方式帶來了一個問題：每一步的輸入幾乎與前一步相同，只是多了一個新 token。如果不對計算過程進行優化，模型就不得不反復重復大量相同的計算，造成資源浪費和效率低下。
為了解決這個問題，KV-Cache（Key Value-Cache）應運而生，它是提升 LLM 推理性能的關鍵技術之一。
KV-Cache 的核心思想是緩存中間計算結果。在 Transformer 架構中，每個 token 在生成下一個 token 時都需要通過自注意力機制計算一系列 Key 和 Value 向量。這些向量描述了當前 token 與前面所有 token 的關系。
1. 關于自注意力機制
要理解KV緩存，首先需要掌握注意力機制的基本原理。在著名的論文《Attention Is All You Need》中，提出了用于Transformer模型的注意力機制公式，該公式幫助模型在生成每個新token時確定應關注哪些先前的token。
注意力機制的核心在于通過一系列計算來量化不同token之間的關聯強度。具體來說，當模型生成一個新的token時，它會根據輸入序列中的所有token計算出一組Query（查詢）、Key（鍵）和Value（值）向量。這些向量通過特定的公式相互作用，以決定當前上下文中每個token的重要性。
讓我們一步一步來看看這個方程是怎么在工程上實現的。
2. 提示詞階段（預填充階段）
我們從一個示例輸入開始，稱為“提示詞”（prompt）。這個輸入首先會被分詞器（tokenizer）切分為一個個的 token，也即模型可處理的基本單位。例如，短語 "The quick brown fox" 在使用 OpenAI 的 o200kbase 分詞器時會被拆分為 4 個 token。隨后，每個 token 被轉換為一個嵌入向量，記作 x? 到 x?，這些向量的維度由模型決定，通常表示為 d_model。在原始 Transformer 論文中，d_model 設定為 512。
為了高效計算，我們可以將所有 n 個嵌入向量堆疊成一個矩陣 X，其形狀為 [n × d_model]。接下來，為了執行注意力機制，我們需要通過三個可學習的投影矩陣 Wq、Wk 和 Wv 來分別生成查詢（Query）、鍵（Key）和值（Value）向量。
具體來說：
- Wq 的形狀為 [
d_model × d_k]，用于將嵌入向量映射到查詢空間，得到 q 矩陣，其形狀為 [n × d_k]； - Wk 的形狀也為 [
d_model × d_k]，生成 k 矩陣，形狀為 [n × d_k]； - Wv 的形狀為 [
d_model × d_v]，生成 v 矩陣，形狀為 [n × d_v]。
這三個矩陣在訓練過程中不斷優化，以捕捉不同 token 之間的依賴關系。其中，dk 和 dv 是設計模型結構時設定的超參數，在最初的 Transformer 模型中被設為 64。
這種線性變換不僅將高維嵌入壓縮到更易操作的空間，還保留了關鍵語義信息，為后續的注意力計算奠定了基礎。
圖片
接下來，我們通過計算查詢矩陣 q 與其對應鍵矩陣 k 的轉置的乘積，得到一個大小為 [n × n] 的自注意力分數矩陣。這個矩陣中的每個元素代表了某個查詢向量與所有鍵向量之間的相似性得分，反映了在生成當前 token 時，模型應將多少注意力分配給前面每一個 token。
對于解碼器類型的大型語言模型（LLM），為了避免模型在生成過程中“偷看”未來的信息，我們采用一種稱為“掩碼自注意力”的機制——即將該矩陣的上三角部分設為負無窮（-inf）。這樣一來，在后續的 softmax 操作中，這些位置的值會趨近于零，從而確保每個 token 在預測時只能關注到它之前的歷史 token，而不會看到未來的輸入。
處理后的自注意力分數矩陣因此成為一個下三角結構，體現了嚴格的因果關系。由于這一操作，模型在生成序列的過程中能夠保持邏輯連貫性和時間順序的正確性。
為了得到最終的注意力輸出，我們首先對這個帶有掩碼的注意力分數矩陣進行縮放，除以 sqrt（d_k)以防止點積結果過大導致梯度飽和；然后應用 softmax 函數將其轉化為概率分布；最后將該分布與值矩陣 v 相乘，得到加權聚合后的上下文信息。這一過程即完成了對輸入序列 "The quick brown fox" 的注意力輸出計算，為下一步的 token 預測提供了基礎。
圖片
3. 自回歸生成階段（解碼階段）
現在，當我們生成下一個 token（例如 "jumps"）并計算其對應的注意力輸出時，在自回歸生成階段（即解碼過程中），模型是如何工作的呢？
在生成“jumps”這一新 token 時，模型會為它生成新的查詢向量 q?、鍵向量 k? 和值向量 v?。但值得注意的是，并不需要重新計算之前所有 token 的 k 和 v，因為這些中間結果已經保存在 KV-Cache 中。
圖片
如圖所示，我們只需關注注意力矩陣中新增的最后一行，其余部分可以從緩存中直接復用。圖中以灰色突出顯示的 k? 到 k? 和 v? 到 v? 都是之前步驟中已計算并緩存的結果。這意味著，只有與當前 token 相關的 q?、k? 和 v? 是新生成的。
最終，第 n 個 token 的注意力輸出只需以下計算（忽略 softmax 和縮放因子以便理解）：
圖片
這正是 KV-Cache 的價值所在：通過緩存先前 token 的 Key 和 Value 向量，模型無需重復計算整個歷史序列，從而顯著減少冗余運算，提升推理效率。
隨著生成序列變長，KV-Cache 的作用愈發重要——它不僅減少了計算負擔，也降低了內存帶寬的壓力，使得 LLM 在實際應用中能夠實現高效、流暢的文本生成。
4. 速度與內存的權衡
KV-Cache 的引入顯著提升了大型語言模型（LLM）的推理速度。其核心思想在于緩存每一步計算生成的 Key 和 Value 向量，使得在生成新 token 時，模型無需重復計算歷史上下文中的 K 和 V 值，從而大幅減少冗余計算，加快響應生成。
然而，這種性能提升也帶來了內存上的代價。由于 KV-Cache 需要為每個已生成的 token 保存對應的 K 和 V 向量，它會持續占用 GPU 顯存。對于本身就需要大量資源的 LLM 來說，這進一步加劇了顯存壓力，尤其是在處理長序列時更為明顯。
因此，在實際應用中存在一個權衡：一方面，使用 KV-Cache 可以加速生成過程，但另一方面，它也會增加內存消耗。當顯存緊張時，開發者可能需要選擇犧牲部分生成速度來節省內存，或接受更高的硬件資源開銷以換取更快的推理表現。這種性能與資源之間的平衡，是部署 LLM 時必須仔細考量的關鍵因素之一。
5. KV-Cache 實踐
KV-Cache 是現代大型語言模型（LLM）推理引擎中至關重要的一項優化技術。在 Hugging Face 的 Transformers 庫中，當你調用 model.generate() 函數生成文本時，默認會啟用 use_cache=True 參數，也就是自動使用 KV-Cache 來提升生成效率。
類似的技術細節在 Hugging Face 官方博客的文章《KV Caching Explained: Optimizing Transformer Inference Efficiency》中也有深入解析。該文通過實驗展示了啟用 KV-Cache 帶來的顯著加速效果：在 T4 GPU 上對模型 HuggingFaceTB/SmoLLM2-1.7B 進行測試，使用 KV-Cache 相比不使用時，推理速度提升了 5.21 倍。
這一性能提升充分說明了 KV-Cache 在實際應用中的重要性，測試的參考代碼如下：
from transformers import AutoTokenizer, AutoModelForCausaLLM
import torch
import time
# Choose model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausaLLM.from_pretrained("gpt2")
# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
# Prepare input
input_sentence = "The red cat was"
inputs = tokenizer(input_sentence, return_tensors="pt").to(device)
# Function to measure generation time
def generate_and_time(use_cache: bool):
torch.cuda.empty_cache()
start_time = time.time()
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=300,
use_cache=use_cache
)
end_time = time.time()
duration = end_time - start_time
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
return duration, output_text
# Measure with KV-cache enabled
time_with_cache, text_with_cache = generate_and_time(use_cache=True)
print(f"\n[use_cache=True] Time taken: {time_with_cache:.4f} seconds")
print(f"Generated Text:\n{text_with_cache}\n")
# Measure with KV-cache disabled
time_without_cache, text_without_cache = generate_and_time(use_cache=False)
print(f"[use_cache=False] Time taken: {time_without_cache:.4f} seconds")
print(f"Generated Text:\n{text_without_cache}\n")
# Speedup factor
speedup = time_without_cache / time_with_cache
print(f"Speedup from using KV-cache: {speedup:.2f}×")KV-Cache的運行速度實際上受到多種因素的綜合影響，其中包括模型的規模（具體體現在注意力層數的多少）、輸入文本的長度n、所使用的硬件設備以及具體的實現細節等。
6. 小結
KV-cache作為一種極為強大的性能優化手段，能夠顯著提升語言模型（LLM）生成文本的速度。其核心機制在于，在生成文本的過程中，通過重用前面步驟中的注意力計算結果，避免重復計算，從而實現更高效的文本生成。具體而言，當計算下一個標記的注意力輸出時，系統會緩存并重用之前步驟中所產生的鍵和值。
不過，需要指出的是，這種性能上的提升并非沒有代價。由于需要存儲這些向量，會占用大量的GPU內存資源，而這些被占用的內存就無法再用于其他任務了。
【參考閱讀與關聯閱讀】
