理解KV cache的作用及優(yōu)化方法

作者丨紫氣東來

在 Transformer 的 Encoder-base 的模型（如 BERT系列）中，推理和訓練過程保持了高度的統(tǒng)一性（差異僅僅在于是否存在反向過程）。而在 Decoder-base 的生成式模型（如 GPT系列）中，推理和訓練存在相當大的差異性，主要體現(xiàn)在推理過程具有以下3點特征：

自回歸

兩階段（第一階段輸入 prompt，第二階段輸入上一個生成的token）

KV cache

以上三點實際上也是相輔相成、不可分割的，其中自回歸的生成模式是根本原因，兩階段是外在的體現(xiàn)形式，KV cache是優(yōu)化手段。

下面將通過梳理整個推理過程，來理解 KV cache 的作用及優(yōu)化方法。

一、KV cache 的由來與基本矛盾

第一階段（prompt 輸入）：

KV cache 作用過程

第二階段（token by token）：

KV cache的顯存占用分析

可見隨著 batch size 和 長度的增大，KV cache 占用的顯存開銷快速增大，甚至會超過模型本身。

而 LLM 的窗口長度也在不斷增大，因此就出現(xiàn)一組主要矛盾，即：對不斷增長的 LLM 的窗口長度的需要與有限的 GPU 顯存之間的矛盾。因此優(yōu)化 KV cache 就顯得非常必要。

二、KV cache 優(yōu)化的典型方法

2.1 共用 KV cache：MQA，GQA

MQA (Multi Query Attention，多查詢注意力) 是多頭注意力的一種變體。其主要區(qū)別在于，在 MQA 中不同的注意力頭共享一個K和V的集合，每個頭只單獨保留了一份查詢參數(shù)。因此K和V的矩陣僅有一份，這大幅度減少了顯存占用，使其更高效。由于MQA改變了注意力機制的結(jié)構(gòu)，因此模型通常需要從訓練開始就支持 MQA 。也可以通過對已經(jīng)訓練好的模型進行微調(diào)來添加多查詢注意力支持，僅需要約 5% 的原始訓練數(shù)據(jù)量 就可以達到不錯的效果。包括 Falcon、SantaCoder、StarCoder 等在內(nèi)很多模型都采用了 MQA 機制。

# Multi Head Attention
self.Wqkv = nn.Linear(     # Multi-Head Attention 的創(chuàng)建方法
    self.d_model,
    3 * self.d_model,     # Q、K和V 3 個矩陣, 所以是 3 * d_model
    device=device
)
query, key, value = qkv.chunk(3, dim=2)      # 每個 tensor 都是 (1, 512, 768)

# Multi Query Attention
self.Wqkv = nn.Linear(       # Multi-Query Attention 的創(chuàng)建方法
    d_model,
    d_model + 2 * self.head_dim,    # 只創(chuàng)建Q的頭向量，所以是 1* d_model, 而K和V不再具備單獨的頭向量, 所以是 2 * self.head_dim
    device=device,
)
query, key, value = qkv.split(
    [self.d_model, self.head_dim, self.head_dim],    # query -> (1, 512, 768), key   -> (1, 512, 96), value -> (1, 512, 96)
    dim=2
)

MHA v.s. GQA v.s. MQA

GQA（Grouped Query Attention，分組查詢注意力）是一種介于多頭注意力和 MQA 之間的折中方案。它將查詢頭（Query Heads）分組，并在每組中共享一個鍵頭（Key Head）和一個值頭（Value Head）。表達能力與推理速度：GQA既保留了多頭注意力的一定表達能力，又通過減少內(nèi)存訪問壓力來加速推理速度。

MHA, GQA, MQA 性能比較

2.2 窗口優(yōu)化

3）箭型 attention 窗口，在LM-Infinit中就已經(jīng)被提出了，其基本原理和StreamingLLM是一致的。

2.3 量化與稀疏

該類方法是基于壓縮的思想，通過量化與稀疏壓縮 KV cache 的 顯存消耗。

當前主流推理框架都在逐步支持 KV cache 量化，一個典型的案例是lmdeploy，下圖展示了其在TurboMind框架下 KV INT8 的支持情況。

lmdeploy 的推理特性

稀疏的方法也比較簡單，其做法無外乎以下幾種方式：

這里最值得一提的是H2O。簡單來說就是通過動態(tài)的評價方式來判斷需要保留和廢棄的KV值，其評估的算法如下所示：

結(jié)果顯示，在 KV cache 稀疏到只有原來的 20% 時仍然可以保持很高的精度。

2.4 存儲與計算優(yōu)化

該方法的典型代表即vLLM的 PagedAttention，簡單來說就是允許在非連續(xù)的內(nèi)存空間中存儲連續(xù)的 K 和 V。詳情可參考筆者之前的文章，在此不予贅述

FlashDecoding 是在 FlashAttention 的基礎(chǔ)上針對 inference 的優(yōu)化主要分為三步：

長文本下將KV分成更小且方便并行的chunk

對每個chunk的KV，Q和他們進行之前一樣的FlashAttention獲取這個chunk的結(jié)果

對每個chunk的結(jié)果進行reduce

三、StreamingLLM：簡潔高效的“無限長度”

StreamingLLM 的基本思想同樣是來源于上述的窗口思想，其最大的創(chuàng)新在于提出了識別并保存模型固有的「注意力池」（attention sinks）錨定其推理的初始 token。下面將詳細討論其工作的原理。

3.1 精度是如何保證的？

核心的發(fā)現(xiàn)：Lost in the Middle。

多個研究都發(fā)現(xiàn)，self-attention 的注意力比較集中于頭部和尾部，對文本中段的注意力相對較弱，如下圖所示：

繪制出 self-attention 的熱力圖也能看到這一點，由此當文本長度超過額定長度時，頭部的 token 就會被遺棄掉，這就會在 softmax 階段產(chǎn)生很大的問題。

3.2 “無限長度”是如何做到的？

該問實際上可以換種表述為：如何在文本長度不斷增加的情況下，保證GPU顯存不會溢出。由于該方案主要應用于多輪對話的場景，那么有必要回顧一下當前多輪對話生成的主流做法，概括起來就以下幾點：

將用戶輸入與模型輸出拼接，中間做必要分割；

多個輪次之間倒序排列，并拼接；

如果前邊所有輪次長度之和超過最大長度，則截斷到最大長度；

上述過程可以用代碼描述如下：

  history = ["
[|Human|]{}
[|AI|]{}".format(x[0], x[1]) for x in history]
  history.append("
[|Human|]{}
[|AI|]".format(text))
  history_text = ""
  flag = False
  for x in history[::-1]:
    if tokenizer(prompt + history_text + x, return_tensors="pt")["input_ids"].size(-1) <= max_length:
 ? ? ? ? ? ?history_text = x + history_text
 ? ? ? ? ? ?flag = True
 ? ? ? ?else:
 ? ? ? ? ? ?break
 ? ?if flag:
 ? ? ? ?inputs = tokenizer(prompt + history_text, return_tensors="pt")
 ? ? ? ?input_ids = inputs["input_ids"][:, -max_length:].to(device)
 ? ? ? ?torch.cuda.empty_cache()
 ? ? ? ?return input_ids, text
 ? ?else:
 ? ? ? ?return None

實際上這就是典型的滑動窗口的做法，滑窗?的存在保證了 GPU 的顯存不會溢出，但是由于上節(jié)的討論，會存在精度損失。

審核編輯：黃飛