KV Cache 原理与优化LLM 推理成本压降实战在 LLM 推理中KV Cache 是性能与成本的核心变量。理解它的工作原理才能从根本上降低延迟和费用。本文从底层原理出发结合 vLLM、SGLang 等主流推理框架给出 KV Cache 的生产级优化方案。一、为什么需要 KV Cache自回归生成的核心流程是每生成一个 token需要对整个历史序列做一次 Self-Attention。对于一个有nnn个历史 token 的序列生成第n1n1n1个 token 时需要计算Attention(Qn,K1:n,V1:n)\text{Attention}(Q_n, K_{1:n}, V_{1:n})Attention(Qn,K1:n,V1:n)其中K1:nK_{1:n}K1:n和V1:nV_{1:n}V1:n是所有历史 token 的 Key 和 Value 向量。没有 KV Cache 的情况每次生成新 token 都要重新计算所有历史 token 的 K、V对于长度nnn的序列时间复杂度是O(n2)O(n^2)O(n2)有 KV Cache 的情况把已经计算过的 K、V 缓存起来每次只计算新 token 的 K、V再用缓存的做 Attention时间复杂度降至O(n)O(n)O(n)增量计算二、KV Cache 的内存计算精确理解 KV Cache 的内存占用是做优化的前提。KV Cache (Bytes)2×nseq×nlayers×nheads×dhead×dtype_bytes\text{KV Cache (Bytes)} 2 \times n_\text{seq} \times n_\text{layers} \times n_\text{heads} \times d_\text{head} \times \text{dtype\_bytes}KV Cache (Bytes)2×nseq×nlayers×nheads×dhead×dtype_bytes其中2分别缓存 K 和 Vnseqn_\text{seq}nseq序列长度输入 输出nlayersn_\text{layers}nlayers模型层数nheadsn_\text{heads}nheadsKV Head 数量GQA 下比 Q Head 少dheadd_\text{head}dhead每个 head 的维度实例计算LLaMA-3 8B配置 - n_layers 32 - n_kv_heads 8GQA8 个 KV head - d_head 128 - dtype fp162 bytes 每个 token 的 KV Cache 2 × 32 × 8 × 128 × 2 131,072 bytes ≈ 128 KB 对于 batch_size1seq_len4096 的请求 128 KB × 4096 ≈ 512 MB 对于 batch_size32seq_len4096 512 MB × 32 ≈ 16 GB三、KV Cache 的三大优化方向3.1 显存分配优化PagedAttention传统 KV Cache 的问题内存碎片化严重。每个请求在开始时就预分配了最大长度max_seq_len的 KV Cache即使实际只用了一半剩余空间也浪费了。vLLM 的 PagedAttention借鉴操作系统的分页内存管理传统方案 [request_A: ████████████__________ ] - 预分配 2048实际用 1024浪费 50% [request_B: ████____________________] - 预分配 2048实际用 512浪费 75% [空闲碎片: ━━━━━━━━━━━━━━━━━━━━━━━━] PagedAttention Physical Memory: [Block 0][Block 1][Block 2][Block 3][Block 4]... request_A → Block 0, Block 2, Block 4不连续动态分配 request_B → Block 1, Block 3不连续动态分配 利用率接近 100%实际效果vLLM 的 PagedAttention 相比传统方案吞吐量提升 2-4x。# 使用 vLLM 的示例fromvllmimportLLM,SamplingParams# vLLM 自动使用 PagedAttentionllmLLM(modelmeta-llama/Llama-3-8b-instruct,max_model_len8192,# 最大序列长度gpu_memory_utilization0.85,# 85% 显存用于 KV Cachemax_num_seqs256,# 最大并发请求数enable_chunked_prefillTrue,# 分块 prefill减少首 token 延迟)# 批量推理KV Cache 自动管理outputsllm.generate(prompts[请介绍一下大模型的工作原理,什么是注意力机制],sampling_paramsSamplingParams(temperature0.7,max_tokens512))3.2 减少 KV 大小量化与压缩KV Cache 量化INT8/INT4# vLLM 开启 KV Cache 量化llmLLM(modelmeta-llama/Llama-3-8b-instruct,kv_cache_dtypefp8,# FP8 量化显存减半# kv_cache_dtypeauto, # 默认与模型一致fp16/bf16)各量化方案的显存/精度权衡KV Cache 格式每 token 显存相对 fp16精度损失fp16128 KB8B模型1x无fp864 KB0.5x极小0.5% benchmarkint864 KB0.5x小建议验证int432 KB0.25x中复杂任务注意SnapKV重要 token 保留不重要的淘汰核心思路不是所有历史 token 的 KV 都同等重要可以根据 attention score 淘汰低重要性的 token# SnapKV 核心逻辑简化版defsnap_kv_compress(key_states,value_states,attention_scores,keep_ratio0.5,window_size32): 压缩 KV Cache保留最近 window_size 个 token 历史中 attention score 最高的 keep_ratio 个 token seq_lenkey_states.shape[-2]n_keepint((seq_len-window_size)*keep_ratio)# 计算每个历史 token 的重要性平均 attention scoreimportanceattention_scores[...,:-window_size].mean(dim-1)# 选出最重要的 n_keep 个历史 token_,top_indicesimportance.topk(n_keep,dim-1,sortedTrue)# 合并选出的历史 token 最近 window_size 个 tokenrecent_indicestorch.arange(seq_len-window_size,seq_len)keep_indicestorch.cat([top_indices.sort().values,recent_indices])compressed_kkey_states[...,keep_indices,:]compressed_vvalue_states[...,keep_indices,:]returncompressed_k,compressed_v实测效果保留 50% KV 的情况下多数任务性能损失 ❤️%显存减少 40%。3.3 跨请求复用Prefix Caching核心思路如果多个请求有相同的前缀如相同的 System Prompt把这部分的 KV Cache 缓存起来复用。典型场景所有请求都用同一个 System Prompt如企业知识库 RAG多轮对话的前几轮内容相同RAG 中相同文档被多次引用# vLLM 开启 Prefix CachingllmLLM(modelmeta-llama/Llama-3-8b-instruct,enable_prefix_cachingTrue,# 开启前缀缓存)# 所有请求都有相同的 system prompt首次计算后缓存system_prompt你是一位专业的客服助手为用户提供技术支持。 公司名称TechCorp 产品CloudDB 数据库服务 [以下是详细的产品文档共 2000 字...]# 首次请求计算 system_prompt 的 KV缓存response1llm.generate([f{system_prompt}\n\n用户问题如何创建数据库],sampling_paramsSamplingParams(max_tokens512))# 后续请求直接复用 system_prompt 的 KV CacheTTFT首 token 延迟大幅降低response2llm.generate([f{system_prompt}\n\n用户问题如何备份数据],sampling_paramsSamplingParams(max_tokens512))# 命中 Prefix Cache首 token 延迟可降低 50-80%Prefix Caching 适用场景与收益场景Prefix 比例TTFT 降低固定 System Prompt2K token短请求的 50%40-60%RAG5K token 知识库背景变化请求的 60%50-70%多轮对话10轮每轮 200 token历史轮次30-50%代码补全文件内容不变80%70-85%四、Radix AttentionSGLang 的创新SGLang斯坦福 伯克利的 Radix Attention 在 Prefix Caching 基础上更进一步用前缀树Radix Tree管理 KV Cache支持任意公共前缀的精细复用传统 Prefix Caching只能复用完全相同的前缀 Radix Attention 请求 A: [System][Doc1][Query1] 请求 B: [System][Doc1][Query2] 请求 C: [System][Doc2][Query3] Radix Tree: ├─ [System] ← 三个请求共享 │ ├─ [Doc1] ← A、B 共享 │ │ ├─ [Query1] ← 仅 A │ │ └─ [Query2] ← 仅 B │ └─ [Doc2] ← 仅 C │ └─ [Query3] ← 仅 CSGLang 实测在 ShareGPT 数据集上相比 vLLM 吞吐量提升1.5-4x在高 Prefix 重用场景下效果更显著。# 使用 SGLangimportsglangassglsgl.functiondefmulti_turn_qa(s,system_prompt,turns):ssgl.system(system_prompt)foruser_msg,_inturns:ssgl.user(user_msg)ssgl.assistant(sgl.gen(response,max_tokens256))returns# SGLang 自动识别公共前缀复用 KV Cacheruntimesgl.Runtime(model_pathmeta-llama/Llama-3-8b-instruct)sgl.set_default_backend(runtime)五、生产环境配置指南5.1 vLLM 推荐配置# 生产环境 vLLM 启动命令python-mvllm.entrypoints.openai.api_server\--modelmeta-llama/Llama-3-8b-instruct\--host0.0.0.0\--port8000\--max-model-len8192\--gpu-memory-utilization0.85\--max-num-seqs256\--enable-prefix-caching\# 开启前缀缓存--enable-chunked-prefill\# 分块 prefill--kv-cache-dtype fp8\# KV Cache FP8 量化--max-num-batched-tokens8192\# 批处理 token 上限--num-scheduler-steps1\--disable-log-requests# 生产环境关闭请求日志5.2 监控 KV Cache 使用率# 通过 vLLM metrics 监控 KV Cache 健康状态importrequestsdefcheck_kv_cache_usage():检查 KV Cache 使用率metricsrequests.get(http://localhost:8000/metrics).textforlineinmetrics.split(\n):ifvllm:gpu_cache_usage_percinlineandnotline.startswith(#):usage_pctfloat(line.split()[-1])*100print(fKV Cache 使用率{usage_pct:.1f}%)ifusage_pct90:print(⚠️ KV Cache 使用率过高考虑减少 max_num_seqs 或增加 GPU)elifusage_pct50:print(ℹ️ KV Cache 使用率偏低可以增加并发请求数)returnusage_pct六、成本优化效果对比在实际生产部署中4x A100 80GLLaMA-3 8B混合请求平均序列长度 2048优化方案吞吐量tokens/s显存使用相对基准提升基准HuggingFace generate1,20065 GB-vLLMPagedAttention4,80058 GB300% FP8 KV Cache5,50038 GB358% Prefix Caching7,20038 GB500%SGLangRadix Attention8,50036 GB608%七、避坑总结不要在 CPU 上用 KV CacheCPU 内存带宽远低于 HBM卸载 KV 到 CPU 只适合特定低延迟要求不高的场景Prefix Caching 需要相同 tokenizer不同版本的 tokenizer 可能导致相同文本 token ID 不同缓存失效GQA 是必选项新模型如果不用 GQAKV Cache 会是 MHA 的 N 倍生产部署极度不划算序列长度对 KV Cache 的影响是线性的用户请求包含的历史轮次越多KV Cache 越大要考虑截断策略batch size 和 KV Cache 是跷跷板增大 batch 可以提高吞吐量但 KV Cache 显存占用也线性增加参考文献Kwon W, et al. “Efficient Memory Management for Large Language Model Serving with PagedAttention.” SOSP, 2023. https://arxiv.org/abs/2309.06180Li Y, et al. “SGLang: Efficient Execution of Structured Language Model Programs.” arXiv, 2023. https://arxiv.org/abs/2312.07104Li Y, et al. “SnapKV: LLM Knows What You are Looking for Before Generation.” arXiv, 2024. https://arxiv.org/abs/2404.14469vLLM Documentation. https://docs.vllm.ai/Ainslie J, et al. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” EMNLP, 2023. https://arxiv.org/abs/2305.13245
KV Cache 原理与优化：LLM 推理成本压降实战
KV Cache 原理与优化LLM 推理成本压降实战在 LLM 推理中KV Cache 是性能与成本的核心变量。理解它的工作原理才能从根本上降低延迟和费用。本文从底层原理出发结合 vLLM、SGLang 等主流推理框架给出 KV Cache 的生产级优化方案。一、为什么需要 KV Cache自回归生成的核心流程是每生成一个 token需要对整个历史序列做一次 Self-Attention。对于一个有nnn个历史 token 的序列生成第n1n1n1个 token 时需要计算Attention(Qn,K1:n,V1:n)\text{Attention}(Q_n, K_{1:n}, V_{1:n})Attention(Qn,K1:n,V1:n)其中K1:nK_{1:n}K1:n和V1:nV_{1:n}V1:n是所有历史 token 的 Key 和 Value 向量。没有 KV Cache 的情况每次生成新 token 都要重新计算所有历史 token 的 K、V对于长度nnn的序列时间复杂度是O(n2)O(n^2)O(n2)有 KV Cache 的情况把已经计算过的 K、V 缓存起来每次只计算新 token 的 K、V再用缓存的做 Attention时间复杂度降至O(n)O(n)O(n)增量计算二、KV Cache 的内存计算精确理解 KV Cache 的内存占用是做优化的前提。KV Cache (Bytes)2×nseq×nlayers×nheads×dhead×dtype_bytes\text{KV Cache (Bytes)} 2 \times n_\text{seq} \times n_\text{layers} \times n_\text{heads} \times d_\text{head} \times \text{dtype\_bytes}KV Cache (Bytes)2×nseq×nlayers×nheads×dhead×dtype_bytes其中2分别缓存 K 和 Vnseqn_\text{seq}nseq序列长度输入 输出nlayersn_\text{layers}nlayers模型层数nheadsn_\text{heads}nheadsKV Head 数量GQA 下比 Q Head 少dheadd_\text{head}dhead每个 head 的维度实例计算LLaMA-3 8B配置 - n_layers 32 - n_kv_heads 8GQA8 个 KV head - d_head 128 - dtype fp162 bytes 每个 token 的 KV Cache 2 × 32 × 8 × 128 × 2 131,072 bytes ≈ 128 KB 对于 batch_size1seq_len4096 的请求 128 KB × 4096 ≈ 512 MB 对于 batch_size32seq_len4096 512 MB × 32 ≈ 16 GB三、KV Cache 的三大优化方向3.1 显存分配优化PagedAttention传统 KV Cache 的问题内存碎片化严重。每个请求在开始时就预分配了最大长度max_seq_len的 KV Cache即使实际只用了一半剩余空间也浪费了。vLLM 的 PagedAttention借鉴操作系统的分页内存管理传统方案 [request_A: ████████████__________ ] - 预分配 2048实际用 1024浪费 50% [request_B: ████____________________] - 预分配 2048实际用 512浪费 75% [空闲碎片: ━━━━━━━━━━━━━━━━━━━━━━━━] PagedAttention Physical Memory: [Block 0][Block 1][Block 2][Block 3][Block 4]... request_A → Block 0, Block 2, Block 4不连续动态分配 request_B → Block 1, Block 3不连续动态分配 利用率接近 100%实际效果vLLM 的 PagedAttention 相比传统方案吞吐量提升 2-4x。# 使用 vLLM 的示例fromvllmimportLLM,SamplingParams# vLLM 自动使用 PagedAttentionllmLLM(modelmeta-llama/Llama-3-8b-instruct,max_model_len8192,# 最大序列长度gpu_memory_utilization0.85,# 85% 显存用于 KV Cachemax_num_seqs256,# 最大并发请求数enable_chunked_prefillTrue,# 分块 prefill减少首 token 延迟)# 批量推理KV Cache 自动管理outputsllm.generate(prompts[请介绍一下大模型的工作原理,什么是注意力机制],sampling_paramsSamplingParams(temperature0.7,max_tokens512))3.2 减少 KV 大小量化与压缩KV Cache 量化INT8/INT4# vLLM 开启 KV Cache 量化llmLLM(modelmeta-llama/Llama-3-8b-instruct,kv_cache_dtypefp8,# FP8 量化显存减半# kv_cache_dtypeauto, # 默认与模型一致fp16/bf16)各量化方案的显存/精度权衡KV Cache 格式每 token 显存相对 fp16精度损失fp16128 KB8B模型1x无fp864 KB0.5x极小0.5% benchmarkint864 KB0.5x小建议验证int432 KB0.25x中复杂任务注意SnapKV重要 token 保留不重要的淘汰核心思路不是所有历史 token 的 KV 都同等重要可以根据 attention score 淘汰低重要性的 token# SnapKV 核心逻辑简化版defsnap_kv_compress(key_states,value_states,attention_scores,keep_ratio0.5,window_size32): 压缩 KV Cache保留最近 window_size 个 token 历史中 attention score 最高的 keep_ratio 个 token seq_lenkey_states.shape[-2]n_keepint((seq_len-window_size)*keep_ratio)# 计算每个历史 token 的重要性平均 attention scoreimportanceattention_scores[...,:-window_size].mean(dim-1)# 选出最重要的 n_keep 个历史 token_,top_indicesimportance.topk(n_keep,dim-1,sortedTrue)# 合并选出的历史 token 最近 window_size 个 tokenrecent_indicestorch.arange(seq_len-window_size,seq_len)keep_indicestorch.cat([top_indices.sort().values,recent_indices])compressed_kkey_states[...,keep_indices,:]compressed_vvalue_states[...,keep_indices,:]returncompressed_k,compressed_v实测效果保留 50% KV 的情况下多数任务性能损失 ❤️%显存减少 40%。3.3 跨请求复用Prefix Caching核心思路如果多个请求有相同的前缀如相同的 System Prompt把这部分的 KV Cache 缓存起来复用。典型场景所有请求都用同一个 System Prompt如企业知识库 RAG多轮对话的前几轮内容相同RAG 中相同文档被多次引用# vLLM 开启 Prefix CachingllmLLM(modelmeta-llama/Llama-3-8b-instruct,enable_prefix_cachingTrue,# 开启前缀缓存)# 所有请求都有相同的 system prompt首次计算后缓存system_prompt你是一位专业的客服助手为用户提供技术支持。 公司名称TechCorp 产品CloudDB 数据库服务 [以下是详细的产品文档共 2000 字...]# 首次请求计算 system_prompt 的 KV缓存response1llm.generate([f{system_prompt}\n\n用户问题如何创建数据库],sampling_paramsSamplingParams(max_tokens512))# 后续请求直接复用 system_prompt 的 KV CacheTTFT首 token 延迟大幅降低response2llm.generate([f{system_prompt}\n\n用户问题如何备份数据],sampling_paramsSamplingParams(max_tokens512))# 命中 Prefix Cache首 token 延迟可降低 50-80%Prefix Caching 适用场景与收益场景Prefix 比例TTFT 降低固定 System Prompt2K token短请求的 50%40-60%RAG5K token 知识库背景变化请求的 60%50-70%多轮对话10轮每轮 200 token历史轮次30-50%代码补全文件内容不变80%70-85%四、Radix AttentionSGLang 的创新SGLang斯坦福 伯克利的 Radix Attention 在 Prefix Caching 基础上更进一步用前缀树Radix Tree管理 KV Cache支持任意公共前缀的精细复用传统 Prefix Caching只能复用完全相同的前缀 Radix Attention 请求 A: [System][Doc1][Query1] 请求 B: [System][Doc1][Query2] 请求 C: [System][Doc2][Query3] Radix Tree: ├─ [System] ← 三个请求共享 │ ├─ [Doc1] ← A、B 共享 │ │ ├─ [Query1] ← 仅 A │ │ └─ [Query2] ← 仅 B │ └─ [Doc2] ← 仅 C │ └─ [Query3] ← 仅 CSGLang 实测在 ShareGPT 数据集上相比 vLLM 吞吐量提升1.5-4x在高 Prefix 重用场景下效果更显著。# 使用 SGLangimportsglangassglsgl.functiondefmulti_turn_qa(s,system_prompt,turns):ssgl.system(system_prompt)foruser_msg,_inturns:ssgl.user(user_msg)ssgl.assistant(sgl.gen(response,max_tokens256))returns# SGLang 自动识别公共前缀复用 KV Cacheruntimesgl.Runtime(model_pathmeta-llama/Llama-3-8b-instruct)sgl.set_default_backend(runtime)五、生产环境配置指南5.1 vLLM 推荐配置# 生产环境 vLLM 启动命令python-mvllm.entrypoints.openai.api_server\--modelmeta-llama/Llama-3-8b-instruct\--host0.0.0.0\--port8000\--max-model-len8192\--gpu-memory-utilization0.85\--max-num-seqs256\--enable-prefix-caching\# 开启前缀缓存--enable-chunked-prefill\# 分块 prefill--kv-cache-dtype fp8\# KV Cache FP8 量化--max-num-batched-tokens8192\# 批处理 token 上限--num-scheduler-steps1\--disable-log-requests# 生产环境关闭请求日志5.2 监控 KV Cache 使用率# 通过 vLLM metrics 监控 KV Cache 健康状态importrequestsdefcheck_kv_cache_usage():检查 KV Cache 使用率metricsrequests.get(http://localhost:8000/metrics).textforlineinmetrics.split(\n):ifvllm:gpu_cache_usage_percinlineandnotline.startswith(#):usage_pctfloat(line.split()[-1])*100print(fKV Cache 使用率{usage_pct:.1f}%)ifusage_pct90:print(⚠️ KV Cache 使用率过高考虑减少 max_num_seqs 或增加 GPU)elifusage_pct50:print(ℹ️ KV Cache 使用率偏低可以增加并发请求数)returnusage_pct六、成本优化效果对比在实际生产部署中4x A100 80GLLaMA-3 8B混合请求平均序列长度 2048优化方案吞吐量tokens/s显存使用相对基准提升基准HuggingFace generate1,20065 GB-vLLMPagedAttention4,80058 GB300% FP8 KV Cache5,50038 GB358% Prefix Caching7,20038 GB500%SGLangRadix Attention8,50036 GB608%七、避坑总结不要在 CPU 上用 KV CacheCPU 内存带宽远低于 HBM卸载 KV 到 CPU 只适合特定低延迟要求不高的场景Prefix Caching 需要相同 tokenizer不同版本的 tokenizer 可能导致相同文本 token ID 不同缓存失效GQA 是必选项新模型如果不用 GQAKV Cache 会是 MHA 的 N 倍生产部署极度不划算序列长度对 KV Cache 的影响是线性的用户请求包含的历史轮次越多KV Cache 越大要考虑截断策略batch size 和 KV Cache 是跷跷板增大 batch 可以提高吞吐量但 KV Cache 显存占用也线性增加参考文献Kwon W, et al. “Efficient Memory Management for Large Language Model Serving with PagedAttention.” SOSP, 2023. https://arxiv.org/abs/2309.06180Li Y, et al. “SGLang: Efficient Execution of Structured Language Model Programs.” arXiv, 2023. https://arxiv.org/abs/2312.07104Li Y, et al. “SnapKV: LLM Knows What You are Looking for Before Generation.” arXiv, 2024. https://arxiv.org/abs/2404.14469vLLM Documentation. https://docs.vllm.ai/Ainslie J, et al. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” EMNLP, 2023. https://arxiv.org/abs/2305.13245