鍵值緩存 是什麼?

KV Cache:鍵值緩存 的完整解釋

在大型語言模型推理中,預先計算並存儲前面 token 的鍵和值向量,避免重複計算的優化技術。

核心概念

鍵值緩存(KV Cache)解決的是 Transformer 模型自回歸推理中的計算冗餘問題。在生成文本時,模型採用自回歸方式逐個生成 token:首先生成第一個 token,然後基於這個 token 生成第二個,再基於前兩個生成第三個,以此類推。

在 Transformer 的自注意力(Self-Attention)層中,注意力的計算涉及查詢(Query、Q)、鍵(Key、K)和值(Value、V)三個向量。對於序列中的每個位置,都需要計算該位置的 Q 與所有位置的 K 的相似度,然後用這些相似度權重和 V,得到該位置的輸出。

在生成第 N 個 token 時,序列長度為 N,模型需要為所有 N 個 token 計算 K 和 V。但關鍵觀察是:前 N-1 個 token 的 K 和 V 在生成第 N-1 個 token 時已經計算過了。在生成第 N 個 token 時,這些值不會改變,只有新 token 的 K 和 V 需要新計算。無優化的情況下,這些已經計算過的 K 和 V 會被重複計算,導致計算時間隨著生成序列的增長而二次方增長。

KV Cache 的核心思想是將前面 token 計算過的 K 和 V 向量存儲在緩存中。生成新 token 時,只需計算該新 token 的 K 和 V,然後將其附加到緩存中。這樣,計算量從 O(N²) 減少到 O(N)。

運作原理

KV Cache 的運作可以分為兩個階段:預填充階段和解碼階段。

預填充階段(Prefill Phase)發生在處理輸入的提示詞(prompt)時。例如,用戶輸入一個提示詞,包含 M 個 token。模型在這個階段計算所有 M 個 token 的查詢、鍵和值,執行標準的自注意力計算。這個階段的計算量較大,但通常是必要的。在這個階段結束後,前 M 個 token 的 K 和 V 被存儲到 KV Cache 中。

解碼階段(Decoding Phase)是逐個生成新 token 的過程。當模型需要生成第 (M+1) 個 token 時:

  1. 計算該新 token 的查詢向量 Q。
  2. 使用該新 token 的 Q 與 KV Cache 中的所有 K(共 M+1 個,包括新 token 自己的 K)進行注意力計算。
  3. 獲得該新 token 的輸出。
  4. 計算該新 token 的 K 和 V,將其添加到 KV Cache 中。
  5. 重複直至生成序列結束。

在解碼階段,每一步只需計算 1 個新 token 的 Q、K、V,然後執行注意力計算。計算量從預填充階段的 O(M²) 減少到每一解碼步的 O(M)。

KV Cache 的大小取決於:

  • 序列長度 N:越長,緩存越大。
  • 隱藏層大小(Hidden Dimension):越大,每個向量佔用空間越多。
  • 層數(Number of Layers):每一層都有自己的 KV Cache。
  • 批次大小(Batch Size):批量處理多個序列時,需要為每個序列維護自己的 KV Cache。

對於一個有 32 層、隱藏層大小 4096、批次大小 1 的模型,處理 2000 個 token 的序列,KV Cache 的大小約為 2000 × 4096 × 32 × 2(K 和 V)× 2(字節,FP16)≈ 500MB。

實際應用

KV Cache 在現代大型語言模型推理中是必不可少的。對於 OpenAI 的 GPT-4 或 Meta 的 Llama 等大型模型,使用 KV Cache 可以將推理速度提升 10-20 倍。

在聊天機器人應用中,用戶與模型進行多輪對話。每輪對話,系統需要考慮所有之前的對話歷史。KV Cache 使得系統可以高效地處理長對話歷史,而不需要重新計算之前 token 的注意力。

在文本生成服務中,如摘要或翻譯服務,KV Cache 大幅降低了每個請求的計算成本,使得系統可以支援更多的並發請求。

在流式文本生成中,當模型逐個生成 token 並將其實時流送到客戶端時,KV Cache 確保每新增一個 token,系統只需進行增量計算,而不是重新處理整個序列。

在多模態模型中(如視覺語言模型),KV Cache 用於高效地處理由視覺編碼器生成的視覺特徵和之前生成的文本 token。

常見誤區

一個常見的誤區是認為 KV Cache 總是有益的。實際上,KV Cache 會消耗大量的記憶體,特別是當處理很長的序列或大批次時。在某些情況下,記憶體成為瓶頸,反而限制了推理的速度。

另一個誤區是認為 KV Cache 對所有計算都有幫助。實際上,KV Cache 主要優化的是解碼階段。在預填充階段,KV Cache 並不能提高速度,反而可能增加開銷。對於推理長序列但生成短序列的任務(如分類、相似度計算),KV Cache 的收益可能不顯著。

還有人誤以為 KV Cache 的大小可以任意增長。實際上,隨著序列長度增加,KV Cache 的大小線性增長,最終會導致記憶體溢出。這限制了模型能處理的最大序列長度。

與相關技術的比較

KV Cache 與注意力機制的優化(如多頭注意力、分組查詢注意力)都是提高 Transformer 效率的方法。分組查詢注意力(Group Query Attention)減少了 KV Cache 的大小,同時保持性能。它通過讓多個查詢頭共享同一個鍵和值頭來實現。

與量化相比,KV Cache 是一種計算優化,而量化是精度優化。兩者可以結合使用,對 KV Cache 進行量化以進一步減少記憶體消耗。

與推測解碼(Speculative Decoding)相比,KV Cache 是一種基礎優化,而推測解碼是在 KV Cache 基礎上的進一步優化,通過並行生成多個候選 token 來加速。

常見問題