首页 / 资讯中心 / 文章详情

大模型推理加速:从KV Cache到vLLM的工程优化实践

大模型推理加速:从KV Cache到vLLM的工程优化实践 ★ FEATURED ARTICLE
1. 为什么大模型推理这么慢自回归生成中的算力陷阱前阵子一个做后端的朋友跟我吐槽他把一个7B量级的开源模型部署到内部测试服务上一开始觉得响应还挺快可一旦上下文拉长、多轮对话聊下去每个字的生成速度肉眼可见地变慢。他一度怀疑是不是显卡散热不行还琢磨着要不要加一张卡。我告诉他先别急着掏钱你撞上的大概率不是算力不足而是推理框架和KV Cache处理方式的问题。这里的关键在于目前主流大模型都是自回归结构一次推理只生成一个token可以粗略理解成一个字词片段生成完之后把它拼回输入再预测下一个token。这个过程串行且不可跳过模型没法一口气把一句话的所有token全部并发算完。你喂给它的Prompt越长它在每一步要回忆的信息就越多如果实现得不好每多生成一个token都要把前面所有token的中间状态从头算一遍计算量随序列长度平方级增长。这也就是为什么很多人在长对话场景下感觉模型越用越慢。我见过不少刚接触推理优化的同学一上来就盯着算子优化、量化这些词却忽略了一个朴素但影响极大的事实模型推理时消耗的计算量很大一部分是被重复计算历史token浪费掉的。想真正把推理速度提上去第一件事不是换卡而是把注意力机制里每一步都在重复计算的KV部分缓存起来。这也是本篇文章的核心主线——先讲清KV Cache的动机与原理再展开聊vLLM这类推理框架如何利用它做工程化优化。1.1 一次只产出一个token的逐字游戏大模型生成文本的模式本质上是在逐字接龙。比如你输入中国的首都是模型内部会先把这个Prompt切分成token序列然后一次性做一轮前向计算预测出第一个输出token最可能是北。接着模型把北追加到Prompt末尾再整体做一次前向预测下一个token京。你可能会想第二次前向时中国的首都是这些token的中间计算结果和第一次应该完全一样为什么还要重算一遍这就是自回归推理的低效根源。如果推理框架把注意力层里每个历史token的Key键和Value值向量重新计算一遍相当于每次生成一个新token都要把整个对话历史重读一次。以序列长度L为例一次解码的时间复杂度是O(L²)整个生成过程的复杂度则是O(L³)级别的累积。换成工程上的说法当你把上下文从1K翻到4K时如果框架没有缓存机制推理速度不会慢4倍而是可能慢一个数量级。这就是我反复跟朋友强调先看KV Cache再谈推理加速的原因。1.2 每一步都在重读历史的数学账为了更直观地理解这个重复计算的代价可以把Transformer每一层注意力简化为一个公式Attention(Q, K, V) softmax(Q · K^T / sqrt(d)) · V这里Q是当前token的查询向量K和V是序列中每个token对应的键向量和值向量。模型要预测下一个token时当前token充当Q它需要跟上下文里所有token的K逐一做点积算完注意力权重后再加权V。如果「当前token之前的那些K和V」每次都被重新计算一遍那么序列越长单步开销就越大。我在实际调服务时测过一个对比6B规模模型、无KV Cache缓存时生成长对话到第200个token左右速度衰减已经非常明显而到了400个token以后基本处于打字机卡壳的状态。同样一台机器、同一个模型接入KV Cache之后单token生成速度在长上下文下趋于平稳。这个差距不是靠推理框架微调能弥补的它是算法结构决定的。1.3 prefill阶段与decode阶段两种节奏的差别理解了逐字接龙之后还有一个工程上很重要的区分注意力计算其实分两个阶段业界一般叫prefill预填充和decode解码。prefill阶段整段Prompt一次性输入模型并行计算Prompt里所有token对应的中间状态。这个阶段是纯矩阵乘法计算量在GPU眼里更像一个大一点的矩阵乘整体算力利用率可以打得很高。decode阶段逐token串行生成。每一步只产生一个新的Q要与历史上所有K、V做attention。这个阶段矩阵的某个维度很小GPU的并行度很难拉满而且显存访问次数和上下文长度强相关。所以做推理加速最忌讳的就是一刀切地看待性能。你会发现很多框架会把prefill和decode拆开调度因为前者是计算密集后者更接近访存密集。KV Cache的作用正是在decode阶段把历史上所有token的K、V提前算好、存到显存里每生成一个新token只需要对这个token计算一次K、V并追加到缓存尾部而不再把整个历史重新算一遍。2. KV Cache到底缓存了什么从QKV分工说起KV Cache说起来简单就是把Key和Value存下来。但为什么偏偏是Cache K和V而不是缓存整个中间结果这就得把Q、K、V在注意力机制里的真实分工弄明白。之前有同事跟我说他看了一堆资料知道要缓存KV但总感觉K、V和Q长得完全一样都是个向量凭啥只缓存两个。我打了个比方你可以把Q理解成我当前想查什么问题K是资料库里每份文档的目录条目V则是每份文档的正文内容。注意力机制做的事情是先拿Q去比对所有K找出哪些文档和自己想查的问题最相关然后取出对应文档的正文V来做加权汇总。在生成下一个token时当前token的Q是新产生的没法提前缓存但它要去比对的所有历史token的K以及要根据权重取出来的历史token的V都是固定不变的——至少在当前生成步内不会变化。既然如此把这两个东西重新算一遍纯属浪费它们正是KV Cache要存储的对象。2.1 Q是当前问题K是索引目录V是正文内容再往深一层说K和V不是单独一个数而是张量。模型有多层注意力层每层都有自己的K和V。每生成一个token每一层都会为这个token计算一份K向量和一份V向量并往缓存里追加。所谓Cache本质上就是把这些K、V按序列位置、按层、按注意力头整齐地排放在显存里。这里有一点值得注意缓存的不是模型的权重而是某个输入序列在若干层注意力中产生的中间激活值。模型权重是固定的只有几十GB训练好之后不动而KV Cache是跟每次请求动态绑定的请求长度越长、并发数越高缓存占的显存就越大。我在真实服务里见过一种典型的显存意外打满模型权重只占了6GB结果KV Cache悄悄吃掉了十几个GB最后OOM。2.2 缓存是如何被建立和更新的KV Cache的构建过程跟随上面前面说的prefill和decode分阶段展开prefill阶段模型处理完整Prompt时会把Prompt里所有token的K、V算出来一次性写入缓存。这一步相当于初始化了一份完整的KV Cache长度等于Prompt长度。decode阶段每生成一个新token模型会为这个token算出新的一份K、V追加到缓存末尾同时用更新后的缓存去计算attention。这个过程在每一层都会发生所以缓存是逐层、逐token不断增长的。用伪代码描述单步decode是这样的# 伪代码单步解码时 KV Cache 的工作方式 for layer in model.layers: # 当前token经过线性变换得到 q, k, v q layer.q_proj(hidden_state) k layer.k_proj(hidden_state) v layer.v_proj(hidden_state) # 把新的 k, v 追加到本层缓存尾部 layer.kv_cache_k[layer.cache_len] k layer.kv_cache_v[layer.cache_len] v # 使用完整缓存参与注意力计算 attn_output attention(q, layer.kv_cache_k[:layer.cache_len1], layer.kv_cache_v[:layer.cache_len1])也就是说每步decode只需要做追加和读取历史K、V不再重新计算。这也解释了为什么有KV Cache的推理服务单token生成速度能保持相对平稳不会因为对话变长而急剧恶化。2.3 一张公式看懂缓存到底占多少显存说完原理来看点实际的。KV Cache占显存这件事很多新手会低估。计算单个token的KV Cache占用可以用一个不算复杂的公式每个token的KV缓存字节数 2 × 层数 × KV头数 × 每个头维度 × 字节数乘以2是因为K和V各一份。以常见的7B开源模型配置来举例假设32层、8个KV头、每个头维度128、权重用FP162字节那么单个token的KV缓存为2 × 32 × 8 × 128 × 2 131072字节 128KB单个token就要128KB看起来好像不多。但当你把序列长度和并发batch乘上去数字就非常可观了。序列长度batch数KV Cache总占用估算10248128KB × 1024 × 8 1GB20488128KB × 2048 × 8 2GB409616128KB × 4096 × 16 8GB819216128KB × 8192 × 16 16GB上面这张表用的是估算值真实显存占用还会因为padding、对齐方式有些出入但量级是差不多的。你如果用的模型KV头更多、层数更深比如32层直接配32个KV头那么上表中的数字还要再翻几倍。显存规划时KV Cache往往比想象中更能吃显存。2.4 缓存思路简单为什么框架实现差异这么大听到这里你可能觉得KV Cache原理太简单了——不就是把算过的K、V存下来吗确实原理门槛很低真正的差异全在怎么存、怎么管理。同一个模型复用一份KV Cache思路不同推理框架做出来的实际吞吐差距可以拉到两三倍核心差别就在存储布局、内存复用和调度策略。这也是为什么有了朴素KV Cache之后还会出现vLLM、PagedAttention这类工程方案。就像你有一个书架理论上每一本书都摆在书架上但书的摆放方式、你找书的方式、书架满了之后如何腾挪空间直接决定了你实际看书的速度。3. vLLM做了什么改变从存下来到存得聪明vLLM是我接触过的推理框架里在工程层面把KV Cache管理做到很精细的一个。它最核心的贡献是PagedAttention它解决的不再是要不要缓存的问题而是缓存怎么分配才不浪费。要理解它为什么重要得先看看朴素KV Cache在真实服务里会踩什么坑。假设一个推理服务同时处理多个请求为了给每个请求预留缓存空间传统做法是在显存里按最大序列长度一次性分配一块连续的缓冲区。很多框架还会按一个最大并发数×最大长度的池子统一分配。问题在于用户的请求长度差异极大有的只问一句你好有的塞进一篇几千字的长文。按最大值预分配短请求占着的多余空间永远用不上显存利用率奇低。3.1 连续分配导致的碎片化与浪费我实际搭服务时观察过一个很典型的案例一台24GB显存的卡模型权重占了大约8GB剩下的16GB按理说全都可以给KV Cache用。但采用传统的预分配方案后同时并发20个请求每个请求预留4096长度的缓存最终实际用到的总长度可能连30%都不到。换句话说一大半显存被预留浪费掉了而浪费的结果就是不敢继续接更多请求服务吞吐上不去。另一个问题是碎片化。即使显存里还有很多空闲空间但这些空间不连续时传统方案依然没法继续为大请求分配一整块缓存。这和操作系统里老式的连续内存分配会遇到外部碎片是一个道理。很多开发者遇到这种问题第一反应是显存不够了加卡但其实是显存管理方式不对浪费了。3.2 PagedAttention像虚拟内存一样管理KV CachePagedAttention的思路很直接既然一次性分配一大块连续显存不现实那就把KV Cache切成固定大小的块按需分配同时用一张映射表记录逻辑位置和物理位置的对应关系。每个块内部是一小段连续的KV数据但块与块之间不需要物理连续。这个设计几乎就是操作系统中虚拟内存分页的翻版。操作系统允许一个进程的逻辑地址空间远大于物理内存并通过页表把逻辑页映射到分散的物理页PagedAttention也是类似的它把一个请求的KV Cache拆成多个block每个block默认可以承载若干个token常见设定是一块16个token然后通过block table维护从逻辑block到物理block的映射。解码时GPU侧只需要按block table的顺序遍历而不需要把散落的物理块拷贝成连续块。这个转换带来的直接收益是显存利用率的大幅提升。物理块按需申请请求结束了立刻释放别的请求马上复用没有大块连续内存也能正常运作。我在同一台24GB显卡上把传统方案换成PagedAttention方案之后并发请求数差不多翻了一倍靠的不是加显存而是把之前被浪费的预留给未来长序列的存储空间真正用了起来。3.3 连续批处理带来的额外收益PagedAttention还顺带解锁了一个重要能力连续批处理。传统批处理会等整个batch的所有请求全部完成之后才一起释放资源、载入下一批请求。如果batch里正好有一个超长请求在慢慢生成其他已经生成完的短请求也只能干等白白浪费算力。连续批处理continuous batching打破了这种同进同出的模式每解码一个token框架都会检查当前batch里哪些请求已经结束结束的立刻从batch中移除释放占用的block并把等待队列里的新请求补进来。这样一来GPU在任何时刻都在尽可能多地处理有效token而不是空转等待最慢的那个请求。这个机制我之前在传统框架里很难做到因为KV Cache的物理空间是预分配死的想动态插一个新请求进去往往要找一大块空闲的连续空间。有了PagedAttention之后新增请求只需要按block粒度分配物理块逻辑上还是连续的一块虚拟KV Cache实现成本低很多。3.4 为什么说它主要提升吞吐而不是单请求延迟要给原理做个总结的话vLLM这类优化方案的直接效果更多体现在系统吞吐而非单次请求延迟。PagedAttention没有减少某个请求生成单个token所需的flops它减少的是显存浪费、排队等资源碎片问题。单位时间内能完成的请求数量提升了用户排队等待的时间自然也就下降了但单请求的decode速度不会因为换框架就产生质的飞跃。这也是我经常跟人说的如果你的应用场景是每次只有一个用户、一个请求、跑一个很长很长的生成任务那PagedAttention带来的直观感受可能不明显但如果你的业务是API服务、多人并发、长短请求混合那它的价值大得多。选推理优化方案之前先想清楚自己的请求模式。4. KV Cache进阶量化、GQA、前缀复用这些搭配怎么用理解了核心原理后真正的工程问题来了KV Cache虽然省了重复计算但它占显存是真的多。我在前面表格里算过一个中等规模的模型配上16并发、8K上下文KV Cache能吃掉十几个GB。怎么在省显存、提速和保质量之间找平衡这里分享几个我实际用过的搭配方案。4.1 KV Cache量化显存立刻省一半以上KV Cache里面的数值要参与后续attention计算当然不能乱压精度。但工程实践发现KV Cache对低比特的容忍度比权重量化高一些。常见做法是给K、V单独做8bit或4bit量化并配合分组scale因子保证每一小段缓存都能被还原到大致正确的数值范围。我试过对某个7B模型做KV Cache 8bit量化显存占用直接砍半困惑度perplexity的劣化几乎可以忽略。4bit量化也能跑但长文本生成时偶尔会出现上下文漂移的问题模型开始说着说着忘记前面聊过什么或者回答变得有点飘。这说明4bit量化对某些模型和数据集并不完全无损需要有校准集配合。如果显存吃紧我一般建议先从8bit开始无损收益大如果还想继续压缩再考虑4bit但上线前一定要做长上下文的回归测试。把KV Cache量化比作压缩聊天记录是很贴切的——压缩了历史记录整体对话占的存储小了但细节可能有所丢失。至于丢多少取决于你的压缩算法和聊天内容的复杂程度。4.2 GQA会让缓存占用大幅变化还有一个从模型结构层面减少KV Cache的手段叫GQA分组查询注意力。传统多头注意力MHA里每个query头都有自己独立的K、V头KV头数与query头数相等缓存量自然大。而GQA让多个query头共享同一组K、V头比如32个query头共享8个KV头KV Cache的理论占用直接降到原来的四分之一。前面公式里的KV头数这项就是GQA起作用的变量。这也是为什么现在很多开源模型的配置里会专门标出KV heads的数量。如果同一个模型有MHA和GQA两个版本部署时优先选GQA的因为它的长上下文生成成本和显存占用都友好得多。MQAMulti-Query Attention是更极端的一种所有query头共用一个KV头缓存最小但质量上不如GQA平衡。4.3 前缀复用跑聊天系统容易被忽略的优化还有一个我初期忽略、后来真香的功能是前缀缓存复用prefix caching。在聊天或多轮对话场景里很多请求共享同一段前缀比如系统提示词、固定的开场白、人设设定。这些内容对应的KV Cache如果不复用每个新请求都要重新计算一遍浪费非常可观。具体做法是框架把KV Cache按块缓存下来并记录块之间的前缀关系。如果新请求的前缀与之前某个请求一致直接沿用已有的cache块只计算新增部分。比如某内部聊天机器人系统提示词固定开头又有一大段设定说明开启前缀复用之后请求的首token延迟大幅下降因为前面一大段KV计算结果都被直接继承了。不过前缀复用也有它自己的复杂度缓存什么时候淘汰、多长的前缀值得复用、跨请求共享会不会被显存打满都需要调。我在生产环境里的经验是当固定系统提示词超过几百token时这个优化性价比极高如果只是几十token的前缀收益就不明显。4.4 不同请求模式下的优化侧重点说了这么多最后还是要强调没有一套配置能吃遍所有场景。我自己会按请求模式粗略分三类来取舍短问题、高并发比如客服、简单问答。KV Cache总量不大瓶颈更多在并发调度和batch管理PagedAttention带来的收益最大。长文档、多轮对话比如RAG、长文总结。KV Cache随上下文快速膨胀优先做KV Cache量化、用带GQA的模型必要时配合前缀复用。单请求、超长生成比如离线写长文。此时连续批处理的价值有限反而是KV Cache的绝对容量和decode单步速度更关键。这张表不是标准答案但能帮你快速定位自己该优先试哪个方向。我经常看到有人拿着人家发的benchmark就照搬参数结果跟自己的场景根本不匹配然后跑出一堆优化了个寂寞的数据。5. 实践记录一次从卡顿到平稳的推理优化过程原理讲再多不如拿一次真实调优的过程收尾。这里不涉及具体公司和项目我按当时在内部测试环境里的做法把大致的排查和优化链路写出来给同样在折腾推理加速的朋友一个参考。5.1 评测前先把三个指标搞清楚很多人调推理速度嘴上说很慢但根本说不清慢在哪。我当时习惯先盯三个指标TTFTTime To First Token从请求发出到返回第一个token的时间。它主要受prefill阶段和排队情况影响也就是模型要读完你的整个Prompt才开始说话。TPOTTime Per Output Token生成每个token的平均耗时对应decode阶段的单步速度。吞吐量单位时间内整个系统能生成的token总数反映的是一个服务同时扛多个请求的能力。如果TTFT高先查prefill或排队长度如果TPOT高就要看decode的单步计算效率和KV Cache读取带宽如果整体吞吐上不去多半是显存管理或者batch策略的问题。那次测试里我最开始只盯总耗时结果怎么调都有点打地鼠后来拆开指标才定位到问题主要在TTFT而非decode速度。5.2 我踩过的两个坑第一个坑是KV Cache上量之后OOM。当时给一个服务开了较大的并发窗口想着batch越大吞吐越高结果请求一上来KV Cache涨得飞快显存撑爆直接OOM。后来才意识到模型权重只是显存占用的固定部分KV Cache是随并发和上下文长度动态增长的弹性部分。如果不给KV Cache设置量上限再大的卡也有被填满的一天。正确做法是给缓存池设一个可接受的上限超出的请求排队等待而不是无脑扩大batch。第二个坑是量化后长上下文质量劣化不明显但确实在掉。8bit量化跑短文本很稳几乎看不出差别但有一段时间做长文档摘要发现输出开始出现前言不搭后语的情况排查了很久最后才回归测试到KV Cache量化精度上。调整了量化分组大小、做了一个简单的校准集之后才恢复。这件事给我最大的教训是显存优化和生成质量的关系往往是短期无感、长期明显一定不能省掉长上下文的回归测试。5.3 个人体会如果让我给刚开始做推理优化的朋友一句建议我会说先别沉迷于各种花式算子优化先把KV Cache这条主线吃透。你把重复计算这个事想清楚了再回头看推理框架的设计会发现很多方案的出发点都很朴素只是工程实现里藏着大量细节。真正常见的瓶颈并不神秘——要么是重复计算要么是显存浪费要么是调度不灵活。KV Cache把第一个问题解决了一部分PagedAttention把第二个问题解决了一部分连续批处理把第三个问题解决了一部分。剩下的大多是在这些方案之间做取舍和调参。再分享一个小技巧做性能测试时别只测一个固定长度的请求。混合长度、混合并发才是真实世界的形态。跑一组长短混合的压测你会对KV Cache和框架调度策略的价值有更直观的感受。我自己那边把这套链路整理完之后同样的显存下并发能力提升了一倍多生成速度也从长对话后期明显卡顿变成了全程相对平稳这个结果基本符合这套原理的预期。
阅读完成 · 觉得有帮助?
咨询建站