LLM Inference

前言 在使用大模型进行推理时我们都会觉得模型的推理速度比较慢，当前也有很多推理加速的方法，像是模型量化、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后的效果。 ...