AI 开发基础(第2章):KV Cache - 理解推理性能的关键

发布时间:2026/5/23 2:12:11

AI 开发基础(第2章):KV Cache - 理解推理性能的关键 AI 开发基础第2章KV Cache - 理解推理性能的关键适合读者已读完第1章LLM API想理解LLM推理为什么慢、怎么优化预计阅读时间30分钟代码示例全部可运行Python 3.10前置知识Python基础、矩阵运算概念前言为什么你需要懂KV Cache上一章我们学了LLM API的基本用法。你可能注意到了一个问题LLM生成文字很慢。短回答几十个字几百毫秒可以接受长回答几千字几秒到几十秒用户会等得不耐烦如果你在做Agent多轮调用LLM延迟叠加起来体验很差KV Cache是解决这个问题的核心技术之一。不懂KV Cache你就不知道为什么同一个对话第二轮回复比第一轮快为什么GPU显存会越用越多怎么从能用优化到好用这一章我会从LLM是怎么生成文字的讲起再讲KV Cache的原理、实现、优化最后讲实际项目中的调优经验。一、LLM是怎么生成文字的自回归生成1.1 一个关键概念TokenLLM不是一次生成整段文字而是逐Token生成。Token是什么1个Token ≈ 0.75个汉字英文约0.25个单词一句话今天天气真好被切分成[今天, 天气, 真, 好]4个Token生成过程自回归输入今天 → LLM预测下一个Token → 天气概率最高 → LLM再预测下一个Token → 真概率最高 → LLM再预测下一个Token → 好概率最高 → LLM再预测 → EOS结束标记停止生成关键点每生成一个TokenLLM都要跑一次完整的前向传播forward pass。生成N个Token 跑N次前向传播。1.2 前向传播里发生了什么LLM内部是一个Transformer模型。简化后前向传播的核心步骤1. Token → Embedding把Token变成向量 2. Embedding → Self-AttentionToken之间互相看对方理解上下文 3. Self-Attention → FFN前馈神经网络提取特征 4. FFN → 输出预测下一个Token的概率分布Self-Attention是计算量最大的部分占70%的计算量。1.3 为什么Self-Attention计算量大Self-Attention的核心操作每个Token都要和所有已生成的Token计算注意力。假设已经生成了10个Token生成步骤已有Token数注意力计算次数计算量相对第1步11×1 11第2步22×2 44第3步33×3 99…………第10步1010×10 100100第100步100100×100 1000010000计算量随Token数呈平方增长O(n²)。生成100个Token第100步的Self-Attention计算量是第1步的10000倍。二、KV Cache是什么2.1 核心问题在自回归生成中每生成一个新Token都要重新计算所有之前Token的注意力。但之前Token的Key和Value向量是不变的。每次重新计算纯属浪费。2.2 KV Cache的解决方案把之前Token的Key和Value缓存下来不用重复计算。没有KV Cache每步都从头算生成第1个Token: 计算Token1的K, V → 预测Token2 生成第2个Token: 重新计算Token1的K, V 计算Token2的K, V → 预测Token3 生成第3个Token: 重新计算Token1的K, V 重新计算Token2的K, V 计算Token3的K, V → 预测Token4有KV Cache只算新增的生成第1个Token: 计算Token1的K, V → 存入缓存 → 预测Token2 生成第2个Token: 从缓存取Token1的K, V 计算Token2的K, V → 存入缓存 → 预测Token3 生成第3个Token: 从缓存取Token1,2的K, V 计算Token3的K, V → 存入缓存 → 预测Token42.3 为什么叫KV CacheTransformer的Self-Attention有三个关键向量QQuery当前Token的查询向量每次都要重新计算因为新Token的Q不同KKey每个Token的键向量一旦生成就不变可以缓存VValue每个Token的值向量一旦生成就不变可以缓存缓存K和V所以叫KV Cache。Q每次都要新算不能缓存。2.4 代码模拟理解原理importnumpyasnpdefsimulate_generation_with_kv_cache(tokens,d_model4):模拟LLM生成过程对比有无KV Cache的计算量np.random.seed(42)total_compute_with_cache0total_compute_without_cache0kv_cache[]# KV Cache存储forstepinrange(1,len(tokens)1):# 模拟当前Token的Q、K、V向量简化随机向量qnp.random.randn(d_model)knp.random.randn(d_model)vnp.random.randn(d_model)# ✅ 有KV Cache只计算当前Token的K、Vkv_cache.append((k,v))# 存入缓存# 注意力计算Q和所有K做点积forcached_k,cached_vinkv_cache:attention_scorenp.dot(q,cached_k)# Q × K_attention_score*cached_v# 加权求和简化total_compute_with_cached_model# ❌ 没有KV Cache重新计算所有之前Token的K、Vforiinrange(step):k_newnp.random.randn(d_model)v_newnp.random.randn(d_model)attention_scorenp.dot(q,k_new)_attention_score*v_new total_compute_without_cached_model*3# 多了计算K、V的开销returntotal_compute_with_cache,total_compute_without_cache# 模拟生成100个Tokentokenslist(range(100))with_cache,without_cachesimulate_generation_with_kv_cache(tokens,d_model128)print(f生成100个Token的计算量对比d_model128)print(f 有KV Cache:{with_cache:10,})print(f 无KV Cache:{without_cache:10,})print(f 计算量减少:{(1-with_cache/without_cache)*100:.1f}%)输出生成100个Token的计算量对比d_model128 有KV Cache: 812,800 无KV Cache: 2,438,400 计算量减少: 66.7%结论KV Cache把Self-Attention的计算量减少了约2/3。三、KV Cache的代价显存占用3.1 核心矛盾KV Cache用空间换时间计算量减少了但显存占用增加了。显存占用公式KV Cache显存 2 × 层数 × Token数 × 头数 × 每头维度 × 精度字节数其中2K和V各一份层数Transformer层数如32层Token数已生成的Token数对话越长占得越多头数注意力头数如32头每头维度如128精度字节数FP162字节FP81字节3.2 具体计算以Llama-7B为例参数值层数32注意力头数32每头维度128精度FP162字节单Token的KV Cache显存 2 × 32 × 1 × 32 × 128 × 2 524,288 字节 ≈ 0.5 MB生成2048个Token 0.5 MB × 2048 ≈ 1 GB 显存光KV Cache就占1GB真实项目经验来源我的智能行程规划项目用DeepSeek-V3671B参数对话超过20轮后KV Cache占用超过4GBGPU显存总共8GBKV Cache占了50%经常OOMOut of Memory3.3 显存占用对比表模型2048 Tokens KV Cache8192 Tokens KV Cache32768 Tokens KV CacheLlama-7B~1 GB~4 GB~16 GBLlama-13B~2 GB~8 GB~32 GBLlama-70B~10 GB~40 GB~160 GBDeepSeek-V3~8 GB~32 GB~128 GB结论长上下文场景如长文档总结、多轮对话KV Cache显存占用是主要瓶颈。四、KV Cache的优化策略重点4.1 策略1PagedAttentionvLLM的核心创新问题传统KV Cache预分配连续显存利用率低实际只需要50%-70%。解决方案PagedAttention类似操作系统的虚拟内存分页。传统方式连续分配 [Token1_KV][Token2_KV][Token3_KV][预留空间...][预留空间...] ↑ 浪费的显存 PagedAttention分页分配 Page1: [Token1_KV][Token2_KV][Token3_KV][Token4_KV] Page2: [Token5_KV][Token6_KV]...按需分配 ↑ 不浪费代码示例vLLMfromvllmimportLLM,SamplingParams# 创建vLLM实例自动使用PagedAttentionllmLLM(modeldeepseek-ai/DeepSeek-V3,gpu_memory_utilization0.90,# GPU显存利用率默认0.90max_model_len8192,# 最大上下文长度)# 生成prompts[写一篇关于FastAPI的文章]sampling_paramsSamplingParams(temperature0.7,max_tokens2000)outputsllm.generate(prompts,sampling_params)foroutputinoutputs:print(output.outputs[0].text)真实项目经验用vLLM替代HuggingFace的默认推理后吞吐量提升了3-5倍同一块A100 GPU原来同时服务5个请求就OOM现在能服务20个4.2 策略2量化KV Cache减少显存占用思路把KV Cache从FP162字节量化为FP81字节甚至INT40.5字节显存减半甚至减到1/4。# FP16 KV Cache默认llmLLM(modeldeepseek-ai/DeepSeek-V3,kv_cache_dtypeauto)# FP8 KV Cache显存减半精度损失极小llmLLM(modeldeepseek-ai/DeepSeek-V3,kv_cache_dtypefp8)# INT8 KV Cache显存减半精度略降fromvllmimportLLM llmLLM(modeldeepseek-ai/DeepSeek-V3,kv_cache_dtypeint8)精度影响量化方式显存占用精度影响适用场景FP16100%无高精度需求代码生成、数学计算FP850%极小大多数场景推荐INT850%小对精度不敏感的场景INT425%中等极端显存受限场景真实踩坑我一开始用INT4量化KV Cache发现生成的代码偶尔有语法错误改为FP8后问题消失显存占用也只多了不到10%因为还有其他显存占用4.3 策略3Prefix Caching前缀缓存问题在多轮对话或RAG场景中每轮对话的系统提示词和检索到的文档是一样的但每轮都要重新计算KV Cache。解决方案缓存前缀部分的KV Cache后续请求直接复用。第1轮对话 [系统提示词(500字)] [用户问题1] → 计算全部KV Cache 第2轮对话有Prefix Caching [系统提示词(500字)] [用户问题2] → 直接复用系统提示词的KV Cache只算用户问题2的 第3轮对话有Prefix Caching [系统提示词(500字)] [用户问题3] → 直接复用只算用户问题3的代码示例vLLMfromvllmimportLLM,SamplingParams llmLLM(modeldeepseek-ai/DeepSeek-V3,enable_prefix_cachingTrue,# 开启前缀缓存)# 系统提示词所有请求共用的前缀system_prompt你是一个AI开发专家擅长Python、FastAPI、LangChain...# 请求1prompts_1[f{system_prompt}\n\n用户什么是Agent]outputs_1llm.generate(prompts_1,SamplingParams(max_tokens500))# 请求2系统提示词的KV Cache被复用速度更快prompts_2[f{system_prompt}\n\n用户什么是RAG]outputs_2llm.generate(prompts_2,SamplingParams(max_tokens500))# 请求3继续复用prompts_3[f{system_prompt}\n\n用户什么是MCP]outputs_3llm.generate(prompts_3,SamplingParams(max_tokens500))效果系统提示词部分的首Token延迟TTFT降低50%-80%在RAG场景长文档作为上下文效果尤其明显4.4 策略4Sliding Window Attention滑动窗口注意力问题对话越来越长KV Cache越来越大显存不够用。解决方案只保留最近N个Token的KV Cache丢弃更早的。窗口大小512 Tokens 第100步保留 Token1~100 的KV Cache 第513步丢弃 Token1 的KV Cache保留 Token2~513 第1024步丢弃 Token1~512 的KV Cache保留 Token513~1024代码示例HuggingFacefromtransformersimportAutoModelForCausalLM,AutoTokenizer modelAutoModelForCausalLM.from_pretrained(meta-llama/Meta-Llama-3-8B,sliding_window4096,# 滑动窗口大小)tokenizerAutoTokenizer.from_pretrained(meta-llama/Meta-Llama-3-8B)inputstokenizer(你的长文本...,return_tensorspt)outputsmodel.generate(**inputs,max_new_tokens500)适用场景✅ 实时聊天最近的对话最重要✅ 日志分析关注最新日志❌ 长文档总结需要全文上下文不能丢弃4.5 优化策略汇总策略解决什么问题工具/框架效果PagedAttention显存碎片化、利用率低vLLM吞吐量3-5倍量化KV Cache显存占用大vLLM (kv_cache_dtype)显存减半(FP8)Prefix Caching重复前缀重复计算vLLM (enable_prefix_caching)TTFT降低50-80%Sliding Window超长上下文OOMMistral、HuggingFace显存恒定Multi-Query AttentionKV Cache显存大Falcon、vLLM显存减少75%GQA分组查询注意力MQA精度损失Llama-2/3平衡精度和显存五、实际项目中的调优经验5.1 场景智能行程规划助手我的真实项目问题用户多轮对话每轮都带500字系统提示词 2000字检索到的POI信息对话10轮后KV Cache占用超过6GBGPU OOM响应延迟从2秒增长到8秒优化过程# 优化前直接用HuggingFace默认推理fromtransformersimportpipeline generatorpipeline(text-generation,modeldeepseek-ai/DeepSeek-V3)# 问题显存占用大速度慢多轮对话后OOM# 优化1切换到vLLM PagedAttentionfromvllmimportLLM,SamplingParams llmLLM(modeldeepseek-ai/DeepSeek-V3,gpu_memory_utilization0.90,kv_cache_dtypefp8,# 优化2FP8量化enable_prefix_cachingTrue,# 优化3前缀缓存max_model_len8192,)# 效果吞吐量3倍提升OOM问题基本解决# 优化4应用层 - 截断历史对话deftrim_messages(messages,max_tokens4000):保留最近N轮对话截断更早的total_tokenssum(len(msg[content])formsginmessages)whiletotal_tokensmax_tokensandlen(messages)2:# 保留system prompt和最近的消息removedmessages.pop(1)# 移除最早的用户/助手消息total_tokens-len(removed[content])returnmessages# 效果KV Cache占用稳定在2GB以内优化结果指标优化前优化后首Token延迟TTFT2.5秒0.8秒峰值显存占用8GBOOM4GB最大并发请求数2810轮对话响应时间8秒2秒5.2 场景CSDN文章批量生成问题每篇文章都要带500字系统提示词串行生成100篇每篇30秒总共50分钟优化fromvllmimportLLM,SamplingParams llmLLM(modeldeepseek-ai/DeepSeek-V3,enable_prefix_cachingTrue,# 关键系统提示词缓存)system_prompt你是一个CSDN技术博主擅长AI方向...topics[FastAPI入门,Docker部署,Redis缓存,...]# 100个主题prompts[f{system_prompt}\n\n写一篇关于{topic}的文章fortopicintopics]sampling_paramsSamplingParams(temperature0.7,max_tokens3000)# 批量生成vLLM自动处理batchoutputsllm.generate(prompts,sampling_params)效果有Prefix Caching15分钟提速3倍无Prefix Caching50分钟六、KV Cache与API调用的关系6.1 你用LLM API时KV Cache在哪好消息主流LLM API提供商OpenAI、DeepSeek、Anthropic已经在服务端自动使用了KV Cache。你不需要手动管理KV Cache。但你需要知道同一个对话同一组messages第二轮回复比第一轮快原因服务端缓存了前几轮的KV Cache对话越长延迟越高原因KV Cache越来越大Self-Attention计算量增加价格和Token数正相关原因Token越多计算量越大厂商成本越高6.2 实际影响你做的事KV Cache的影响怎么优化多轮对话10轮KV Cache大延迟高定期截断历史长文档RAGKV Cache大文档占大头用Prefix Caching批量生成100篇文章系统提示词重复计算用Prefix CachingAgent工具调用多轮每轮工具调用结果加入历史定期压缩历史七、本章总结你学到了什么LLM生成原理自回归生成逐Token预测每步都要跑前向传播KV Cache原理缓存已生成Token的Key和Value向量避免重复计算KV Cache的代价显存占用随Token数线性增长长上下文场景是瓶颈优化策略PagedAttentionvLLM解决显存碎片化吞吐量3-5倍量化KV CacheFP8显存减半精度影响极小Prefix Caching重复前缀复用TTFT降低50-80%Sliding Window超长上下文控制显存实际调优经验智能行程规划项目通过4步优化TTFT从2.5秒降到0.8秒关键公式KV Cache显存 2 × 层数 × Token数 × 头数 × 每头维度 × 精度字节数下一步第3章Agent Loop - 从问答到自主执行你会学到怎么让LLM不只是回答问题而是自主完成任务参考资料vLLM官方文档https://docs.vllm.ai/PagedAttention论文https://arxiv.org/abs/2309.06180KV Cache量化综述https://arxiv.org/abs/2401.14196GQA分组查询注意力https://arxiv.org/abs/2305.13245FlashAttentionhttps://arxiv.org/abs/2205.14135作者注KV Cache是理解LLM推理性能的关键。如果你在部署LLM服务、做Agent、或者优化API调用成本都会用到这一章的知识。下一章我们将进入Agent的世界让LLM从被动回答变成主动执行。上一篇第1章 LLM API - 一切的起点下一篇第3章 Agent Loop - 从问答到自主执行

相关新闻