DeepSeek-V3的推理缓存优化：如何复用历史对话的KV缓存减少计算

在智能对话系统中，用户经常需要与模型进行多轮交互。每次对话时，模型都需要重新处理整个对话历史，这不仅浪费计算资源，还会导致响应延迟。DeepSeek-V3通过KV缓存（Key-Value Cache）技术，将历史对话中的注意力计算结果缓存起来，显著提升了多轮对话场景下的推理效率。本文将详细介绍这一技术的实现原理和使用方法。## KV缓存的工作原理### 传统推理的性能瓶颈在标准的Tra...

成冠冠Quinby

555人浏览 · 2025-10-06 01:14:31

成冠冠Quinby · 2025-10-06 01:14:31 发布

DeepSeek-V3的推理缓存优化：如何复用历史对话的KV缓存减少计算

【免费下载链接】DeepSeek-V3 项目地址: https://gitcode.com/gh_mirrors/de/DeepSeek-V3

KV缓存的工作原理

传统推理的性能瓶颈

在标准的Transformer模型推理过程中，每生成一个新token都需要处理整个对话历史序列。例如，在一个包含100个token的对话中，生成第101个token时，模型需要重新计算所有101个token的注意力分数，其中前100个token的计算与上一轮完全重复。随着对话轮次增加，这种重复计算会导致推理时间呈二次增长。

KV缓存的优化思路

KV缓存通过存储每一层注意力模块中计算得到的Key和Value矩阵，避免在后续生成过程中重复计算历史序列的注意力结果。具体来说：

缓存存储：在处理提示词（Prompt）阶段，将每层注意力计算的Key和Value矩阵存入缓存
增量计算：生成新token时，只计算当前token的Key和Value，并与缓存中的历史KV矩阵拼接后进行注意力计算
内存管理：通过限制缓存的最大序列长度（由max_seq_len参数控制），防止内存溢出

图1：KV缓存技术对推理速度的提升效果（数据来源：figures/benchmark.png）

DeepSeek-V3中的KV缓存实现

缓存数据结构

在DeepSeek-V3的model.py中，KV缓存通过两种方式实现，由attn_impl参数控制：

# 代码片段来自[model.py](https://link.gitcode.com/i/8b1538a418d5b45977181630463b56bd#L439-L443)
if attn_impl == "naive":
    self.register_buffer("k_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim), persistent=False)
    self.register_buffer("v_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim), persistent=False)
else:
    self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank), persistent=False)
    self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim), persistent=False)

naive模式：使用独立的k_cache和v_cache存储Key和Value矩阵
absorb模式：将Key和Value合并存储在kv_cache中，并额外维护位置编码缓存pe_cache

缓存写入机制

在注意力计算过程中，新生成token的KV矩阵会被追加到缓存中：

# 代码片段来自[model.py](https://link.gitcode.com/i/8b1538a418d5b45977181630463b56bd#L477-L478)
self.k_cache[:bsz, start_pos:end_pos] = k
self.v_cache[:bsz, start_pos:end_pos] = v

其中start_pos和end_pos标记当前批次在缓存中的位置，确保多批次处理时缓存的正确管理。

缓存读取与注意力计算

生成新token时，模型从缓存中读取历史KV矩阵，并与当前token的KV矩阵拼接后进行注意力分数计算：

# 代码片段来自[model.py](https://link.gitcode.com/i/8b1538a418d5b45977181630463b56bd#L479)
scores = torch.einsum("bshd,bthd->bsht", q, self.k_cache[:bsz, :end_pos]) * self.softmax_scale

缓存配置与调优

关键配置参数

KV缓存的行为由configs目录下的JSON配置文件控制，主要参数包括：

参数	含义	推荐值	配置文件示例
`max_seq_len`	最大缓存序列长度	4096-16384	config_16B.json
`dim`	模型隐藏层维度	2048-8192	config_671B.json
`n_heads`	注意力头数	16-64	config_v3.1.json

内存与性能权衡

增大max_seq_len可以缓存更长的对话历史，减少缓存失效导致的性能下降，但会增加显存占用。不同模型规模的显存占用参考：

16B模型（config_16B.json）：每个batch约占用8GB显存
236B模型（config_236B.json）：每个batch约占用40GB显存

建议根据可用GPU内存和典型对话长度合理设置这一参数。

使用KV缓存进行推理

命令行推理

通过generate.py脚本启动推理时，KV缓存会自动启用。以下是典型的启动命令：

python inference/generate.py \
    --ckpt-path /path/to/checkpoints \
    --config inference/configs/config_v3.1.json \
    --interactive \
    --max-new-tokens 2048

缓存行为控制

在generate.py的生成循环中，缓存通过prev_pos变量控制增量计算：

# 代码片段来自[generate.py](https://link.gitcode.com/i/63cc2b2ea5e43a38c768494cf113d9a5#L60-L61)
for cur_pos in range(min(prompt_lens), total_len):
    logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos)

prev_pos标记当前序列在缓存中的起始位置
每次迭代只计算prev_pos:cur_pos区间的新token，实现增量推理

多轮对话示例

以下是使用交互式推理进行多轮对话的示例，KV缓存在幕后自动管理对话历史：

>>> 什么是人工智能？
人工智能是计算机科学的一个分支，致力于创建能够模拟人类智能的系统...

>>> 它有哪些应用领域？
人工智能的应用领域非常广泛，主要包括以下几个方面：
1. 自然语言处理...

每次输入新问题时，模型仅处理当前问题的token，而对话历史通过KV缓存复用，大幅提升响应速度。

缓存优化进阶

缓存压缩技术

DeepSeek-V3提供了FP8量化选项（通过dtype: "fp8"配置），可以在fp8_cast_bf16.py中找到相关实现，进一步减少KV缓存的内存占用。

动态缓存管理

对于超长对话场景，DeepSeek-V3实现了基于YARN（Yet Another Rotary Position Embedding）的动态位置编码校正，在model.py的precompute_freqs_cis函数中：

# 代码片段来自[model.py](https://link.gitcode.com/i/8b1538a418d5b45977181630463b56bd#L368-L370)
if seqlen > args.original_seq_len:
    low, high = find_correction_range(beta_fast, beta_slow, dim, base, args.original_seq_len)
    smooth = 1 - linear_ramp_factor(low, high, dim // 2)

这一技术允许模型在有限缓存空间下处理更长的对话序列。