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

最新推荐文章于 2025-05-13 20:48:56 发布

原创

最新推荐文章于 2025-05-13 20:48:56 发布 · 6.2k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#性能优化 #pytorch #人工智能 #python #深度学习

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

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), 0)
        text = tokenizer.decode(in_tokens)
        print(f'step {i} input: {text}', flush=True)