KVCache feat. 单向注意力VS双向注意力

Wed, 12 Jun 2024 00:00:00 +0000

前言

在使用大模型进行推理时我们都会觉得模型的推理速度比较慢，当前也有很多推理加速的方法，像是模型量化、FlashAttention、使用更好的GPU、KVCache等，今天我们就来聊聊KVCache是怎么回事，我会从transformer decoder的推理过程开始，到为什么KVCache可以加速模型推理，以及具体的效果进行介绍。

Transformer的self-attention机制

目前的大模型都是基于Transformer的decoder架构进行设计的，所以先把著名的self-attention公式抛出来。

由于大模型使用的是decoder，所以在计算时会mask掉当前token的后边的token，例如当前为token_n，mask矩阵会mask掉n之后的所有token值，目的是防止解码过程中看到答案。

下图是GPT2的解码过程，给定输入，模型预测下一个token，然后在下一步中使用上一步预测的token作为输入再次进行预测。图来源：https://jalammar.github.io/illustrated-gpt2/

我们来逐步计算一下解码的过程。

因为decoder中的self-attention是masked self-attention，所以在计算时需要注意进行mask。

我们将q和k的计算整理成矩阵，可以得到以下公式。

展开的话就是这样的一个矩阵。

对于token_1，Attention_1(Q,K,V)=softmax(\frac{Q_1K^T_1}{\sqrt{d_k}})\vec{V_1}

对于token_2，Attention_2(Q,K,V)=softmax(\frac{Q_2K^T_1}{\sqrt{d_k}})\vec{V_1}+softmax(\frac{Q_2K^T_2}{\sqrt{d_k}})\vec{V_2}

对于token_3，Attention_3(Q,K,V)=softmax(\frac{Q_3K^T_1}{\sqrt{d_k}})\vec{V_1}+softmax(\frac{Q_3K^T_2}{\sqrt{d_k}})\vec{V_2}+softmax(\frac{Q_3K^T_3}{\sqrt{d_k}})\vec{V_3}

可以看到在token_2推理的过程中，K_1、V_1是重复使用的；可以看到在token_3推理的过程中，K_1、V_1、K_2、V_2是重复使用的。

还可以看出来，每次计算attention只需要使用当前的Q即可，并不需要之前的向量。

所以我们就可以把之前计算过的K、V缓存起来，这就是今天我们要介绍的KVCache。

KVCache的作用

现在我们就可以给KVCache进行定义了，在decoder-only架构中，通过缓存解码过程中的K、V，来避免重复计算，从而达到推理加速的效果。

下图清晰了对比了使用KVCache和不使用KVCache的区别，图来源https://medium.com/@joaolages/kv-caching-explained-276520203249

实现细节

在huggingface的transformers中，通过modelling_gpt2.py可以看到具体的实现细节。

        query = self._split_heads(query, self.num_heads, self.head_dim)
        key = self._split_heads(key, self.num_heads, self.head_dim)
        value = self._split_heads(value, self.num_heads, self.head_dim)

				# 通过layer_past来传递上一次计算的key value
        if layer_past is not None:
            past_key, past_value = layer_past
            # 通过拼接来得到最新的key 和 value
            key = torch.cat((past_key, key), dim=-2)
            value = torch.cat((past_value, value), dim=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

在transformers中使用KVCache对比

在Transformers中进行对比开启KVCache后的效果。

LLM Inference on 厨师海森堡的厨房

KVCache feat. 单向注意力VS双向注意力

前言

Transformer的self-attention机制

KVCache的作用

实现细节

在transformers中使用KVCache对比