ARTICLE DETAIL

资讯详情

深耕网站SEO优化与搜索引擎排名提升的一线实战洞察。

大模型推理加速实战:从 KV Cache 到 vLLM 调度机制的深度解析

大模型推理加速实战:从 KV Cache 到 vLLM 调度机制的深度解析 一、为什么大模型推理越来越慢大模型生成文本的过程并不是一次完成的而是典型的自回归生成输入人工智能正在改变 生成人工智能正在改变 - 世界 生成人工智能正在改变世界 - 的 生成人工智能正在改变世界的 - 运行方式每生成一个 Token模型都需要执行一次完整的前向计算。对于长度为T的序列如果每次都重新计算全部历史 Token注意力计算会产生大量重复工作。理想情况下我们希望历史 Token 的 Key 和 Value 只计算一次 后续生成过程直接复用历史计算结果这就是 KV Cache 的核心思想。大模型推理通常包含两个阶段阶段计算对象主要瓶颈Prefill一次处理完整输入提示词GPU 计算能力Decode每次生成一个新 Token显存带宽和 KV CachePrefill 阶段适合批量矩阵计算通常 GPU 利用率较高。Decode 阶段每次只生成一个 Token但需要读取模型权重和历史 KV Cache因此往往是典型的显存带宽受限任务。二、KV Cache 到底缓存了什么Transformer 的自注意力计算可以表示为Q XWq K XWk V XWv Attention(Q, K, V) softmax(QK^T / sqrt(d))V在生成新 Token 时新 Token 只会产生新的 Query、Key 和 Value。历史 Token 的 Key 和 Value 不会发生变化因此可以缓存K_cache [K1, K2, K3, ..., Kt] V_cache [V1, V2, V3, ..., Vt]生成第t 1个 Token 时只需要计算Q(t1), K(t1), V(t1)然后将新的 Key 和 Value 追加到缓存中。如果没有 KV Cache每一步都要重新计算历史 Token 的 K 和 V推理速度会随着上下文增长快速下降。KV Cache 的显存计算KV Cache 的理论显存占用可以通过以下公式估算显存 2 × 层数 × Batch Size × 序列长度 × KV Head 数量 × Head Dim × 每个元素字节数其中2表示 Key 和 ValueKV Head 数量在 GQA 或 MQA 模型中通常小于 Query Head 数量Head Dim是每个注意力头的维度FP16 和 BF16 通常占用 2 字节FP32 占用 4 字节下面的代码可以估算一个模型的 KV Cache 显存。fromtransformersimportAutoConfigdefestimate_kv_cache_memory(model_name:str,sequence_length:int,batch_size:int1,dtype_bytes:int2,):configAutoConfig.from_pretrained(model_name)num_layersconfig.num_hidden_layers num_attention_headsconfig.num_attention_heads# GQA/MQA 模型存在 num_key_value_headsnum_kv_headsgetattr(config,num_key_value_heads,num_attention_heads,)head_dimgetattr(config,head_dim,config.hidden_size//num_attention_heads,)total_bytes(2*num_layers*batch_size*sequence_length*num_kv_heads*head_dim*dtype_bytes)gibtotal_bytes/1024/1024/1024print(f模型{model_name})print(f层数{num_layers})print(fKV Head{num_kv_heads})print(fHead Dim{head_dim})print(f序列长度{sequence_length})print(fBatch Size{batch_size})print(fKV Cache{gib:.3f}GiB)estimate_kv_cache_memory(model_nameQwen/Qwen2.5-7B-Instruct,sequence_length8192,batch_size1,)需要特别注意模型权重经过 INT4 量化后显存占用会明显下降但这并不意味着 KV Cache 也自动变成 INT4。很多推理服务的显存主要不是被模型权重占满而是被长上下文、大 Batch 的 KV Cache 占满。三、GQA 为什么能够降低推理成本传统多头注意力中Query、Key、Value 的头数相同Query Heads Key Heads Value Heads而在 GQA 中Query Heads Key/Value Heads例如Query Heads 32 KV Heads 8这样可以在保留较强表达能力的同时将 KV Cache 降低到原来的四分之一。KV Cache 的内存大小与KV Head 数量成正比KV Cache ∝ KV Head 数量因此模型结构本身就会直接影响推理成本。这也是为什么在部署阶段不能只关注参数量。两个同样是 7B 参数的模型由于层数、KV Head 数量和上下文长度不同实际推理显存可能存在显著差异。四、Transformers 中使用 KV CacheHugging Face Transformers 默认通常会启用 KV Cache但建议在代码中显式指定。importtimeimporttorchfromtransformersimportAutoTokenizer,AutoModelForCausalLM MODEL_IDQwen/Qwen2.5-7B-InstructtokenizerAutoTokenizer.from_pretrained(MODEL_ID,trust_remote_codeTrue,)modelAutoModelForCausalLM.from_pretrained(MODEL_ID,torch_dtypeauto,device_mapauto,trust_remote_codeTrue,)model.eval()prompt请解释大模型推理中的 KV Cache并分析它对显存和延迟的影响。inputstokenizer(prompt,return_tensorspt,).to(model.device)torch.inference_mode()defgenerate(use_cache:bool):returnmodel.generate(**inputs,max_new_tokens128,do_sampleFalse,use_cacheuse_cache,)# 预热generate(use_cacheTrue)torch.cuda.synchronize()starttime.perf_counter()outputgenerate(use_cacheTrue)torch.cuda.synchronize()elapsedtime.perf_counter()-start texttokenizer.decode(output[0],skip_special_tokensTrue,)print(f耗时{elapsed:.3f}秒)print(text)可以使用下面的代码对比启用和禁用 KV Cache 的差异。defbenchmark(use_cache:bool,repeat:int3):times[]for_inrange(repeat):torch.cuda.synchronize()starttime.perf_counter()generate(use_cacheuse_cache)torch.cuda.synchronize()times.append(time.perf_counter()-start)returnsum(times)/len(times)with_cachebenchmark(use_cacheTrue)without_cachebenchmark(use_cacheFalse)print(f启用 KV Cache{with_cache:.3f}秒)print(f禁用 KV Cache{without_cache:.3f}秒)print(f加速比{without_cache/with_cache:.2f}x)这个实验通常能够说明两个问题上下文越长KV Cache 的收益越明显。KV Cache 用显存换取计算量和延迟。但 Transformers 的默认生成方式并不适合高并发生产服务因为它通常需要为每个请求维护一套连续的缓存并且缺乏高效的请求调度机制。五、传统推理服务的三个问题1. KV Cache 连续分配导致显存浪费假设最大上下文长度设置为 8192请求 A实际使用 512 Token 请求 B实际使用 2048 Token 请求 C实际使用 7000 Token如果服务按照最大长度提前分配空间大量显存会处于空闲状态。此外不同请求的结束时间不同显存会产生碎片已分配 | 空闲 | 已分配 | 空闲 | 已分配即使剩余显存总量足够也可能因为缺乏连续空间而无法接收新请求。2. 静态 Batch 无法适应请求变化传统静态 Batch 通常需要等待一批请求全部完成请求 A生成 20 Token 请求 B生成 200 Token 请求 C生成 50 Token如果按照最长请求等待A 和 C 完成后GPU 仍然需要为它们保留 Batch 位置。这会造成GPU 计算资源浪费短请求延迟增加吞吐量下降3. 长提示词会阻塞短请求如果一个请求包含几十万 Token 的长文档另一个请求只有一句话那么两者共用一个调度队列时长 Prefill 可能长时间占用 GPU。因此推理优化不仅是 Kernel 优化更是显存管理 请求调度 批处理策略六、vLLM 的核心设计vLLM 主要通过以下机制提升推理性能PagedAttention Continuous Batching 高效 KV Cache 管理 Prefix Caching Chunked Prefill其中最关键的是 PagedAttention 和 Continuous Batching。安装 vLLMpipinstallvllm启动 OpenAI 兼容服务vllm serve Qwen/Qwen2.5-7B-Instruct--host0.0.0.0--port8000--dtypeauto--gpu-memory-utilization0.90--max-model-len8192--enable-prefix-cachingWindows PowerShell 中使用反引号换行Linux 或 macOS 中使用反斜杠。发送请求curlhttp://localhost:8000/v1/chat/completions-HContent-Type: application/json-d{ model: Qwen/Qwen2.5-7B-Instruct, messages: [ { role: user, content: 解释 PagedAttention 和普通 Attention 的区别。 } ], temperature: 0.2, max_tokens: 256, stream: false }Python 客户端代码如下fromopenaiimportOpenAI clientOpenAI(base_urlhttp://localhost:8000/v1,api_keyEMPTY,)responseclient.chat.completions.create(modelQwen/Qwen2.5-7B-Instruct,messages[{role:user,content:请从显存管理角度解释 PagedAttention。,}],temperature0.2,max_tokens256,)print(response.choices[0].message.content)七、PagedAttention 如何解决显存碎片PagedAttention 的设计思想类似于操作系统的虚拟内存。传统方式通常将一个请求的 KV Cache 存储在一块连续显存中Request A - 连续物理内存PagedAttention 则将 KV Cache 切分为固定大小的 Block逻辑 Block 0 - 物理 Block 17 逻辑 Block 1 - 物理 Block 4 逻辑 Block 2 - 物理 Block 29逻辑上仍然是一段连续序列但物理显存可以分散存放。每个请求维护一个 Block TableRequest A: [17, 4, 29, 11] Request B: [8, 12, 31]当请求新增 Token 时只需要申请新的物理 Block而不需要重新申请一整块连续空间。一个简化版的 Block 管理器如下classKVBlockManager:def__init__(self,total_blocks:int,block_size:int):self.block_sizeblock_size self.free_blockslist(range(total_blocks))self.block_tables{}defallocate(self,request_id:str,token_count:int):required(token_countself.block_size-1)//self.block_sizeiflen(self.free_blocks)required:raiseRuntimeError(KV Cache 显存不足)physical_blocks[self.free_blocks.pop()for_inrange(required)]self.block_tables[request_id]physical_blocksreturnphysical_blocksdefappend_token(self,request_id:str,token_count:int):blocksself.block_tables[request_id]used_tokenstoken_count capacitylen(blocks)*self.block_sizeifused_tokenscapacity:ifnotself.free_blocks:raiseRuntimeError(没有可用 KV Block)blocks.append(self.free_blocks.pop())defrelease(self,request_id:str):blocksself.block_tables.pop(request_id,[])self.free_blocks.extend(blocks)这段代码只展示了基本思想真实 vLLM 还需要处理Block 引用计数Prefix Cache 共享Copy-on-Write请求结束后的回收多 GPU 下的缓存管理不同请求的 Block Table 映射PagedAttention 的主要收益是减少了预分配浪费和外部碎片。如果 Block Size 过大最后一个 Block 可能浪费更多空间。如果 Block Size 过小Block Table 和调度管理开销会增加因此实际系统需要在显存利用率和管理开销之间取平衡。八、Continuous Batching 的工作方式静态 Batch 的执行方式类似Batch 1请求 A、B、C 等待 A、B、C 全部完成 Batch 2请求 D、EContinuous Batching 则在每一步重新调整 Batch第 1 步A、B、C 第 2 步A、B、C、D 第 3 步A、C、D 第 4 步A、D、E请求 B 完成后可以立即释放它的 KV Cache并将新请求加入正在运行的 Batch。简化调度逻辑如下running_requests[]waiting_requests[]whileTrue:# 将等待队列中的请求加入运行队列whilecan_admit_new_request(running_requests):requestwaiting_requests.pop(0)running_requests.append(request)# 每个请求只生成一个或一小组 Tokenresultsmodel.decode_step(running_requests)finished[]forrequest,resultinresults:request.append(result)ifrequest.is_finished():finished.append(request)# 释放已经完成请求的 KV Cacheforrequestinfinished:running_requests.remove(request)kv_cache_manager.release(request.id)真实系统还需要考虑调度优先级等待时间请求长度当前 KV Cache 占用最大 Batch Token 数Prefill 和 Decode 的公平性是否启用 Prefix CacheContinuous Batching 的关键并不是简单地“把更多请求放进 Batch”而是在每一个调度周期内尽可能提高 GPU 的有效工作量。九、Prefill 和 Decode 为什么需要不同的调度策略PrefillPrefill 一次处理用户输入的全部 Token例如输入长度4096 Token它主要执行大规模矩阵运算计算密度较高适合充分利用 GPU。DecodeDecode 一次通常只生成一个 Token输入长度4096 Token 生成长度1 Token此时模型仍然需要读取大量权重和 KV Cache但实际计算量相对较小通常受显存带宽限制。如果一个超长 Prefill 请求长时间占用 GPUDecode 请求就会出现明显的首 Token 延迟。因此高性能推理引擎通常会对 Prefill 进行切分这就是 Chunked Prefill 的基本思想4096 Token Prefill 拆分为 1024 1024 1024 1024这样可以让系统在处理长输入的同时穿插执行已有请求的 Decode改善整体延迟。需要区分两个指标TTFTTime To First Token首 Token 延迟 TPOTTime Per Output Token后续 Token 平均延迟Chunked Prefill 通常有助于改善多请求场景下的 TTFT 公平性但也可能增加调度复杂度需要通过压测确定最佳参数。十、Prefix Caching 如何进一步减少重复计算很多业务请求具有相同的前缀系统提示词 公司知识库说明 固定格式约束 统一安全策略例如下面两个请求请求 A [相同系统提示词] 用户问题 A 请求 B [相同系统提示词] 用户问题 B如果每次都重新执行相同前缀的 Prefill会产生重复计算。Prefix Caching 可以按照前缀 Token 序列计算哈希hash(prefix_tokens) - KV Cache Blocks后续请求发现相同前缀后可以直接复用已经计算好的 KV Block。适合使用 Prefix Caching 的场景长系统提示词多轮对话代码仓库分析固定文档模板批量处理同一份上下文不适合的场景每个请求前缀都完全不同前缀非常短请求生命周期很短GPU 显存非常紧张Prefix Caching 主要减少 Prefill 计算不会消除用户问题部分和新生成 Token 的 Decode 计算。十一、使用 vLLM 进行离线批量推理如果不需要 HTTP 服务也可以直接使用 vLLM 的离线接口。fromvllmimportLLM,SamplingParams model_nameQwen/Qwen2.5-7B-InstructllmLLM(modelmodel_name,dtypeauto,gpu_memory_utilization0.90,max_model_len8192,enable_prefix_cachingTrue,)sampling_paramsSamplingParams(temperature0.2,top_p0.9,max_tokens256,)prompts[解释 KV Cache 的工作原理。,解释 PagedAttention 如何减少显存碎片。,解释 Continuous Batching 的调度过程。,]outputsllm.generate(prompts,sampling_params,)foroutputinoutputs:print(*60)print(f输入{output.prompt})print(output.outputs[0].text)离线推理适合数据集批量生成自动摘要离线评测文档分类合成训练数据在线服务更关注TTFT、P95 延迟、并发数离线推理更关注总吞吐、平均 Token 成本、GPU 利用率二者的最优配置并不完全相同。十二、如何正确进行性能测试不能只执行一次请求然后用总耗时判断性能。一个有效的推理压测至少需要记录首 Token 延迟 TTFT每个输出 Token 延迟 TPOT完整请求延迟P50、P95、P99 延迟输入 Token 数输出 Token 数每秒生成 Token 数GPU 显存使用量GPU 利用率并发请求数下面是一个简化的流式压测脚本importjsonimporttimeimportstatisticsfromconcurrent.futuresimportThreadPoolExecutorimportrequests URLhttp://127.0.0.1:8000/v1/chat/completionsMODELQwen/Qwen2.5-7B-InstructPROMPT请从工程角度解释大模型推理优化要求包含 KV Cache 和批处理调度。defone_request(_):body{model:MODEL,messages:[{role:user,content:PROMPT,}],temperature:0,max_tokens:128,stream:True,}starttime.perf_counter()first_token_timeNonewithrequests.post(URL,jsonbody,streamTrue,timeout300,)asresponse:response.raise_for_status()forlineinresponse.iter_lines():ifnotlineornotline.startswith(bdata:):continuepayloadline[5:].strip()ifpayloadb[DONE]:breakjson.loads(payload)iffirst_token_timeisNone:first_token_timetime.perf_counter()endtime.perf_counter()return{ttft_ms:((first_token_time-start)*1000iffirst_token_timeelseNone),e2e_ms:(end-start)*1000,}defpercentile(values,p):valuessorted(values)indexint((len(values)-1)*p)returnvalues[index]defmain():request_count20concurrency4withThreadPoolExecutor(max_workersconcurrency)aspool:resultslist(pool.map(one_request,range(request_count)))ttft[item[ttft_ms]foriteminresultsifitem[ttft_ms]isnotNone]e2e[item[e2e_ms]foriteminresults]print(f请求数{request_count})print(f并发数{concurrency})print(fTTFT 平均值{statistics.mean(ttft):.2f}ms)print(fTTFT P95{percentile(ttft,0.95):.2f}ms)print(fE2E 平均值{statistics.mean(e2e):.2f}ms)print(fE2E P95{percentile(e2e,0.95):.2f}ms)if__name____main__:main()压测时必须保证以下条件一致相同模型 相同量化方式 相同输入长度 相同输出长度 相同采样参数 相同 GPU 相同并发数否则对比结果没有实际意义。十三、常用参数如何调整1.gpu_memory_utilization--gpu-memory-utilization0.90该参数控制 vLLM 可以使用的 GPU 显存比例。过低会导致 KV Cache 容量不足过高可能影响其他 CUDA 操作或导致显存不足。一般可以从0.85到0.92之间逐步测试。2.max-model-len--max-model-len8192该参数越大理论上支持的上下文越长但 KV Cache 占用也会随之增加。如果业务实际只需要 4096 Token就没有必要设置为 32768。3.max-num-seqs--max-num-seqs64它限制并发序列数量。并发并不是越高越好。当显存、带宽或调度开销达到瓶颈后继续增加并发可能导致 P95 延迟恶化。4.max-num-batched-tokens该参数影响单次调度中允许处理的 Token 总量。较大值通常有利于提高吞吐但可能增加短请求的等待时间。在线对话场景应该同时观察吞吐和 TTFT。5. Tensor Parallel当模型无法放入单张 GPU 时可以使用张量并行vllm serve Qwen/Qwen2.5-72B-Instruct --tensor-parallel-size4但多卡并行会引入 GPU 间通信开销。模型规模较小时盲目增加 GPU 数量可能反而降低单请求性能。十四、FlashAttention 和 PagedAttention 不是同一个东西这两个概念经常被混淆。FlashAttentionFlashAttention 主要优化注意力 Kernel减少 HBM 与片上 SRAM 之间的数据读写 降低 Attention 中间矩阵的显存占用它解决的是计算 Kernel 的访存效率问题。PagedAttentionPagedAttention 主要优化 KV Cache 的存储和分配减少连续显存分配要求 降低 KV Cache 内部碎片 支持动态请求调度它解决的是推理服务中的缓存管理问题。二者可以同时使用FlashAttention提升单次 Attention 计算效率 PagedAttention提升多请求 KV Cache 利用率一个偏 Kernel一个偏系统调度。真正的高性能推理系统需要两者协同。十五、常见误区误区一模型量化后KV Cache 也会自动降低不一定。权重量化主要降低模型参数显存。KV Cache 是否量化取决于推理框架和具体配置。误区二增加 Batch Size 一定提高吞吐Batch 增大后GPU 利用率可能提高但也会带来KV Cache 占用增加请求排队时间增加P95 延迟上升显存不足风险增加应该通过压测寻找平衡点。误区三流式输出会减少推理耗时流式输出主要改善用户体验让用户更早看到第一个 Token。它不会减少模型实际计算量。真正影响推理计算的因素包括模型结构 KV Cache Batch 调度 Kernel 量化 显存带宽误区四设置更大的上下文窗口更保险更大的最大上下文长度会预留或占用更多缓存资源。正确做法是根据真实业务分布设置P50 输入长度 P95 输入长度 最大允许输入长度而不是简单地把最大上下文设置到模型理论上限。十六、优化思路总结大模型推理优化可以抽象成四个层次第一层减少计算启用 KV Cache使用 GQA 或 MQA使用更小模型使用量化使用投机采样第二层减少显存占用PagedAttentionKV Cache Block 管理Prefix CachingKV Cache 量化合理限制最大上下文第三层提高 GPU 利用率Continuous Batching合理增加并发Chunked Prefill高效 Attention Kernel合理设置 Batch Token 上限第四层优化服务质量监控 TTFT监控 TPOT监控 P95 和 P99区分短请求和长请求控制排队时间进行动态限流最终的优化目标不是单纯追求某一个指标而是在以下目标之间取得平衡吞吐量 响应延迟 显存占用 服务稳定性 单 Token 成本结语KV Cache 解决的是重复计算问题PagedAttention 解决的是 KV Cache 的高效存储问题Continuous Batching 解决的是多请求调度问题。三者之间的关系可以概括为KV Cache 避免重复计算历史 Token PagedAttention 提高 KV Cache 的显存利用率 Continuous Batching 让不同生命周期的请求共享 GPU 计算资源如果只是本地验证模型效果Transformers 已经足够使用。如果需要面向真实业务提供高并发推理服务就必须从模型结构、缓存管理、请求调度和硬件资源四个维度进行整体优化。只有完成完整压测才能知道系统究竟提升了多少性能。不同模型、GPU、上下文长度和并发规模下最终结果可能存在数量级差异。
返回列表