多轮对话为什么越聊越慢？FlashAttention 的 KV Cache 优化实践

这是一篇关于 FlashAttention 在多轮对话场景下 KV Cache 优化的深度技术解析文章，结合了生产环境痛点与昇腾 NPU 的适配实践。

多轮对话为什么越聊越慢？FlashAttention 的 KV Cache 优化实践

部署 Llama2-70B 做多轮对话，发现一个奇怪的现象：第一轮回复秒回，聊到第五轮开始卡，第十轮直接超时。查了一圈，问题出在 KV Cache 和 FlashAttention 的配合上。今天把这个问题彻底拆透。

一、先说现象：多轮对话的延迟曲线

假设你部署了一个 Llama2-70B 的对话服务，序列长度上限设 4096。
用户的对话过程是：

• 第 1 轮：用户输入 50 个 token → 模型生成 200 个 token
• 第 2 轮：用户输入 30 个 token → 模型生成 150 个 token
• 第 3 轮：用户输入 40 个 token → 模型生成 300 个 token
• …
• 第 N 轮：用户输入 X 个 token → 模型生成 Y 个 token

每轮结束后，模型的历史 KV（Key 和 Value）会缓存起来，下一轮复用。
问题在于：每新来一个 token，注意力要计算它和所有历史 token的相似度。

• 第 1 轮：算 50 个 token 的注意力
• 第 5 轮：算 50+200+30+150+40 = 470 个 token 的注意力
• 第 10 轮：可能超过 2000 个 token

序列越长，每 token 的推理延迟越高——这就是“越聊越慢”的根本原因。

二、标准注意力处理 KV Cache 的方式：低效

标准注意力的计算流程（第 N 轮推理时）：

• Q（当前新 token 的 Query）：形状[1, num_heads, head_dim]
• K_cache（历史所有 token 的 Key）：形状[seq_len, num_heads, head_dim]
• V_cache（历史所有 token 的 Value）：形状[seq_len, num_heads, head_dim]

计算过程如下代码所示：

scores=Q @ K_cache.T# [1, seq_len]attn=Softmax(scores)# [1, seq_len]output=attn @ V_cache# [1, num_heads, head_dim]

问题：

• K_cache 和 V_cache 存在显存里。
• 每轮推理都要把整个 K_cache 和 V_cache 读出来。
• 当seq_len = 2000时，读 2000 个 token 的 K 和 V，显存带宽瓶颈。

用个比喻——这就像炒菜时，每次要用到冰箱里的所有食材（不管用不用得上），全部搬出来，用完再放回去。食材越多（历史越长），搬来搬去越慢。

三、FlashAttention 怎么优化 KV Cache？

FlashAttention 的核心还是分块（Tiling），但针对 KV Cache 场景做了特化。

3.1 增量计算（Incremental Computation）

FlashAttention 不需要把整个 K_cache 和 V_cache 一次性读出来。
它把 K_cache 和 V_cache 按块读取：

1. 每次只读一小块 K_cache 和 V_cache（比如 128 个 token）。
2. 和 Q 算局部注意力。
3. 立刻和对应的 V_cache 小块相乘，累加到输出。
4. 中间结果留在片上（UB），不写回显存。

这意味着：历史 KV 越长，FlashAttention 的优势越明显——因为标准注意力是“全读”，FlashAttention 是“分块读”，后者的显存带宽需求低得多。

3.2 KV Cache 的复用

多轮对话时，K_cache 和 V_cache 只追加，不覆盖。

• 第 1 轮结束后，K_cache 里有 250 个 token（50 输入 + 200 输出）的 Key。
• 第 2 轮结束后，K_cache 里有 430 个 token（250 + 30 + 150）的 Key。

FlashAttention 可以只对新来的 Q 做分块计算，历史的 KV 已经在 Cache 里了，不需要重新计算。
标准注意力也能复用 KV Cache，但读取整个 Cache 的显存带宽开销还在。FlashAttention 把这个开销也大幅降低了。

3.3 PagedAttention 的协同优化

vLLM 框架里有个 PagedAttention，把 KV Cache 分成固定大小的“页”。
FlashAttention 可以和 PagedAttention 协同工作：

• PagedAttention：负责 KV Cache 的内存管理（避免碎片）。
• FlashAttention：负责每页内部的注意力计算（分块优化）。

CANN 8.0 的 FlashAttention 实现（在ops-transformer仓库里）已经适配了 PagedAttention 的内存布局，可以直接和 vLLM + 昇腾 NPU 的方案配合使用。

四、实测：多轮对话的延迟对比

在 Atlas 800（昇腾 NPU，64GB 显存）上部署 Llama2-70B，用 benchmarking 工具模拟多轮对话：

4.1 测试配置

• 模型：Llama2-70B，FP16
• 序列长度上限：4096
• 每轮输入：30-50 token
• 每轮输出：150-300 token
• 测试轮数：10 轮

4.2 延迟数据

关键观察：

• 第 1 轮差异不大：历史短，KV Cache 小，标准注意力也能应付。
• 第 5 轮开始差距拉大：历史超过 1000 token，标准注意力的显存带宽瓶颈显现。
• 第 10 轮差距最明显：标准注意力延迟 210ms/token（不可用），FlashAttention 还能维持在 55ms/token（可用）。

4.3 显存占用

第 10 轮标准注意力直接崩，因为要把整个 KV Cache 读出来，显存不够。FlashAttention 因为分块计算，显存占用低得多。

五、在昇腾 NPU 上开启 KV Cache 优化

5.1 确认 CANN 版本

需要 CANN 8.0+，因为 KV Cache 的优化是在这个版本里引入的。

cat/usr/local/Ascend/ascend-toolkit/latest/version.cfg

5.2 确认框架支持

如果用 vLLM + 昇腾 NPU：

fromvllmimportLLM,SamplingParams llm=LLM(model="meta-llama/Llama-2-70b-hf",tensor_parallel_size=4,# 4 卡模型并行max_model_len=4096,# 开启 FlashAttention + KV Cache 优化enable_flash_attn=True,# vLLM 0.3.0+ 支持)

如果用 PyTorch + torch_npu：

importtorchimporttorch_npu model=LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf",torch_dtype=torch.float16,device_map="npu",# 开启 KV Cache 复用use_cache=True,# 这个开关控制是否复用 KV Cache)

踩坑：use_cache=True是默认开启的，但如果你用了梯度检查点（gradient checkpointing），KV Cache 会被禁用。推理时别开梯度检查点。

5.3 监控 KV Cache 的使用情况

CANN 8.0 提供了 KV Cache 监控工具：

# 开启 KV Cache 监控torch_npu.set_kv_cache_monitor_enabled(True)# 推理outputs=model.generate(...)# 查看 KV Cache 占用stats=torch_npu.get_kv_cache_stats()print(f"KV Cache 占用:{stats['used_bytes']/1024**3:.1f}GB")print(f"KV Cache 命中率:{stats['hit_rate']:.2%}")

命中率是关键指标——如果命中率低于 90%，说明 KV Cache 没有正确复用，检查use_cache是否开启。

六、KV Cache 的进阶优化技巧

6.1 动态批处理（Dynamic Batching）

多用户并发时，每个用户的对话历史长度不一样，KV Cache 的大小也不一样。
动态批处理的思路：把历史长度相近的请求打包成一批，减少 padding 带来的显存浪费。

CANN 8.0 的 FlashAttention 实现支持可变长度序列的批处理（variable-length batching），可以直接用：

# 不同用户的输入长度不同inputs=[{"input_ids":[1,2,3],"kv_len":50},{"input_ids":[4,5,6,7],"kv_len":200},{"input_ids":[8,9],"kv_len":500},]# FlashAttention 会自动处理可变长度的 KV Cacheoutputs=model.generate(inputs)

6.2 KV Cache 的量化

KV Cache 的精度要求没有权重那么高，可以用 INT8 甚至 INT4 量化。
CANN 8.0 支持 KV Cache 量化（需要手动开启）：

# 开启 KV Cache 量化（INT8）model=model.half()# 先转 FP16model=torch_npu.quantize_kv_cache(model,dtype=torch.int8)

量化后：

• KV Cache 显存占用降低 50%（INT8）或 75%（INT4）。
• 精度损失很小（因为 KV Cache 的数值分布比较集中）。
• 吞吐提升 10-20%（因为显存带宽压力小了）。

踩坑：KV Cache 量化只支持 CANN 8.0+，且需要模型本身支持（Llama2 系列支持，部分小模型不支持）。

6.3 多轮对话的“遗忘”策略

如果用户对话历史太长（超过 4096 或 8192），可以主动遗忘最早的对话：

# 超过 4096 token 时，截断最早的 1024 个 tokenifkv_cache_len>4096:kv_cache=kv_cache[1024:]# 截断最早的

FlashAttention 支持稀疏 KV Cache（只保留最近的 N 个 token），可以在ops-transformer仓库里找到相关实现。

七、ops-transformer仓库里的 KV Cache 相关算子

除了 FlashAttention，ops-transformer仓库里还有几个和 KV Cache 相关的算子：

• GQA（Grouped Query Attention）：减少 KV Cache 的大小（多个 Query 头共享一组 KV）。
• MQA（Multi-Query Attention）：更激进的 KV Cache 共享（所有 Query 头共享一组 KV）。
• PagedAttention 适配层：和 vLLM 的 PagedAttention 协同工作。

这些算子和 FlashAttention 可以叠加使用，进一步降低 KV Cache 的显存占用。

八、总结

1. 跑一个多轮对话的 benchmark：用 ab（ApacheBench）或 locust 模拟多用户并发，观察延迟曲线。
2. 读ops-transformer里 FlashAttention 的 KV Cache 实现：重点看fa_kv_cache.cpp（KV Cache 复用逻辑）。
3. 试 KV Cache 量化：用 INT8 量化，观察显存占用和精度的变化。
4. 调动态批处理：模拟不同长度对话请求的并发，观察吞吐变化。
5. 看 vLLM + 昇腾 NPU 的方案：vLLM 官方文档里有昇腾 NPU 的适配指南。

仓库地址（纯文本，直接粘浏览器打开）：
https://atomgit.com/cann/ops-transformer