大模型推理性能优化之KV Cache解读

最新推荐文章于 2024-08-26 23:20:51 发布

Linux基金会AI&Data基金会

最新推荐文章于 2024-08-26 23:20:51 发布

阅读量5.6k

点赞数 11

文章标签：性能优化 pytorch 人工智能 python 深度学习

本文链接：https://blog.csdn.net/lf_ai/article/details/130838524

版权

本文介绍了KV Cache在大模型推理性能优化中的作用，解释了其原理和实现细节，指出KV Cache能减少Self-Attention层的计算量，不影响MLP层，并降低了block间的数据传输。通过对Transformer模型的分析，展示了KV Cache如何通过缓存重复计算结果提升推理速度。

摘要由CSDN通过智能技术生成

0. 引言

做大模型性能优化的一定对KV Cache不陌生，那么我们对这个技术了解到什么程度呢？请尝试回答如下问题：

KV Cache节省了Self-Attention层中哪部分的计算？
KV Cache对MLP层的计算量有影响吗？
KV Cache对block间的数据传输量有影响吗？本文打算剖析该技术并给出上面问题的答案。

1. KV Cache是啥

大模型推理性能优化的一个常用技术是KV Cache，该技术可以在不影响任何计算精度的前提下，通过空间换时间思想，提高推理性能。网上有一些关于该技术的分析博客，但读过后仍然会很迷糊，甚至可能会被带偏，认为这个Cache过程和数据库读取或CPU Cache加速类似的荒谬结论。刚开始我也有类似误解，直到逐行查阅并运行源码，才清楚了解到其Cache了啥，以及如何节省计算的。

2. 背景

生成式generative模型的推理过程很有特点，我们给一个输入文本，模型会输出一个回答（长度为N），其实该过程中执行了N次推理过程。即GPT类模型一次推理只输出一个token，输出token会与输入tokens 拼接在一起，然后作为下一次推理的输入，这样不断反复直到遇到终止符。

如上描述是我们通常认知的GPT推理过程。代码描述如下：

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer


model = GPT2LMHeadModel.from_pretrained("gpt2", torchscript=True).eval()

# tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
in_text = "Lionel Messi is a"
in_tokens = torch.tensor(tokenizer.encode(in_text))

# inference
token_eos = torch.tensor([198]) # line break symbol
out_token = None
i = 0
with torch.no_grad():
    while out_token != token_eos:
        logits, _ = model(in_tokens)
        out_token = torch.argmax(logits[-1, :], dim=0, keepdim=True)
        in_tokens = torch.cat((in_tokens, out_token