资讯详情

资讯详情

KV Cache与PagedAttention:大模型推理显存优化实战解析

我最早对“显存爆炸”有切肤之痛是在本地显卡上跑7B模型的时候单条请求响应十几秒看着还能忍同事一上来并发测试服务不是卡死就是直接报CUDA out of memory。后来我把账算明白了——模型权重只是显存账本的一半另一半大头是KV Cache而KV Cache的分配方式又决定了并发一上来会崩得多快。这篇文章就从推理加速和显存治理两个角度把KV Cache和PagedAttention彻底拆开讲适合正在部署大模型推理服务、或自学大模型技术栈时对着这两个名词头疼的朋友。1. 推理慢、显存爆根子不在算力而在于“重复劳动”1.1 一个7B模型的自白单请求正常、并发即崩先说一个被很多人忽略的事实Transformer解码是逐token生成的。你看到模型“唰唰”输出一大段话实际底层是“预测下一个词”这个动作重复了几百上千次。每次只生成一个token然后把新token拼到序列后面再预测下一个。这个循环叫自回归解码。在这个循环里每一步都要做注意力计算。注意力公式长这样Attention(Q, K, V) softmax(Q * K^T / sqrt(d)) * V其中Q是当前正在处理的token的查询向量K和V是序列中所有token的键和值向量。问题就在这做第100个token的注意力时你需要第1到第99个token的K和V做第200个token时你需要第1到第199个token的K和V。如果每一步都重新算历史token的K和V那计算量会随着序列长度二次方增长慢是必然的。我在刚接触推理优化时犯过一个典型错误以为解码慢纯粹是GPU算力不够。后来用性能分析工具一看大量计算时间花在重复计算历史token的Key矩阵和Value矩阵上。这属于典型的“重复劳动”——结果一模一样但因为没存下来每一步都在白算。1.2 显存都去哪了模型权重只是冰山一角推理时显存占用主要有三块模型权重7B模型FP16权重约14GB如果量化成INT4/GGUF Q4约4-5GB。KV Cache每一层、每一个序列token都要存K和V随上下文长度和并发数线性增长。激活值和中间张量前向传播过程中的临时数据一般占比不大但预填充长prompt时会出现峰值。很多人部署时只盯着模型权重大小。8GB显卡跑7B模型量化后权重4GB看起来还剩4GB很宽裕。实际上并发一上来KV Cache直接把剩余显存吃穿。我试过在8GB卡上跑Qwen2.5-7B-Instruct的GGUF Q4版本上下文拉到4096并发8个请求OOM前毫无征兆日志只留下一句“No available memory for the cache distribution”。1.3 没有缓存的朴素方案为什么会慢到崩溃假设我们没有KV Cache解码第t步时模型需要重新计算第1到t-1个token的K、V。这等于每一步都在重复做前面所有步骤的Key投影和Value投影。算一笔简单账生成1024个token没有缓存的话KV投影总计算量大约是有缓存的512倍因为第t步要算t-1套K/V累计约L^2/2次投影而有缓存只需L次。我知道有人会杠现代GPU矩阵乘法很快多算点也没事。但请记住两个约束一是推理服务有延迟要求用户等不了1024个token慢慢算二是多算的每一步都在烧电、占算力吞吐量直接被打骨折。所以KV Cache不是为了“锦上添花”而是自回归解码里必不可少的工程设施。2. KV Cache用显存换重复计算的加速器2.1 一句话说清KV Cache的逻辑KV Cache的思路极其朴素把历史token算好的K矩阵和V矩阵存下来下一步直接从显存里读而不是重新算。第t步时只有新token需要算自己的K、V然后拼接到缓存末尾。用伪代码表示就是# 简化示例带KV Cache的逐token解码 def decode_one_token(model, new_token_id, past_kvNone): # past_kv: 历史所有层的 K、V 缓存结构为 [(K_layer0, V_layer0), ...] outputs model( new_token_id, past_key_valuespast_kv, use_cacheTrue ) # 模型内部逻辑 # 1. 只对当前token计算新的K_new、V_new # 2. 将K_new、V_new拼接到past_kv对应层 # 3. attention时从拼接后的完整K、V中读取 next_logits outputs.logits[:, -1, :] next_token torch.argmax(next_logits, dim-1) return next_token, outputs.past_key_values这段代码如果你用HuggingFace Transformers跑过生成应该很眼熟。past_key_values、use_cache这些参数就是KV Cache在框架层的体现。2.2 KV Cache的数学账本那个shape是怎么来的很多人看源码时对KV Cache的维度一头雾水。其实推导特别直接。每一层的K、V缓存shape是[batch_size, num_kv_heads, seq_len, head_dim]为什么是四维batch_size是当前同时处理的序列数num_kv_heads是KV头的数量seq_len是已经生成的token数head_dim是每个头的维度。注意这里用的是num_kv_heads不是num_attention_heads——在多查询注意力GQA/MQA架构里KV头数小于Q头数后面会细说。以Llama-2-7B为例32层32个注意力头每个head_dim128FP16存储每元素2字节。单个序列上下文长度4096时单层KV Cache大小2K和V两个矩阵× 1batch× 32KV头× 4096seq_len× 128head_dim× 2字节 2 × 1 × 32 × 4096 × 128 × 2 67,108,864 字节 ≈ 64MB乘以32层单序列4096上下文的KV Cache约2GB。你没看错7B模型FP16权重14GB可能用不起但KV Cache单序列就能干到2GB。2.3 8G显卡跑7B模型的真实显存账单我把不同上下文和并发下的KV Cache需求列了一张表按7B模型、FP16、MHA架构32 KV头估算并发序列数上下文长度KV Cache总占用1512256 MB120481 GB140962 GB440968 GB8409616 GB1128K8 GB×32最后一行我故意没填数因为你一看规律就能自己算KV Cache随上下文长度线性增长、随并发数线性增长。8GB显卡量化后权重占4GB剩4GB可用跑单序列4096上下文勉强够2GB KV Cache 激活值一旦并发4个序列KV Cache直接要8GBOOM就是必然的。这也是为什么很多私有化部署方案“看起来模型不大一上线就崩”核心原因就是没把KV Cache算进显存预算。2.4 MHA、GQA、MQA头部共享如何改变显存需求如果你一直用MHA多头注意力估算KV Cache会发现自己部署Llama-3-8B时账对不上。因为新一代模型大量采用GQA分组查询注意力KV头数大幅缩小。注意力架构Q头数KV头数KV Cache相对大小代表模型MHA32321xLlama-2-7B大多数早期模型GQA3281/4Llama-3-8BQwen2.5-7BDeepSeek系列MQA3211/32PaLM部分早期流式模型Qwen2.5-7B-Instruct就是GQA架构32个Q头但只有8个KV头。同样7B规模、4096上下文、单序列KV Cache从2GB降到512MB差距非常可观。所以选模型时不能只看参数量。同样7B一个MHA一个GQA长上下文场景显存需求差4倍。我现在的经验是凡是面向长上下文的部署优先选GQA架构KV Cache小省下来的显存能多接几路并发。2.5 一个关键承诺KV Cache不会改变生成结果有朋友问过我缓存K和V会不会让模型输出变差答案是不会。KV Cache存的是历史token做完投影后的K、V矩阵不是采样结果数学上跟重新计算完全等价。唯一的风险来自量化KV Cache比如INT8、FP8那是精度近似另说。这一点很重要它意味着你可以放心大胆用缓存不用在“加速”和“质量”之间做取舍。3. PagedAttention把操作系统的分页哲学搬进GPU3.1 连续式缓存为什么吞掉了近60%显存KV Cache解决了重复计算但KV Cache本身怎么分配也是个问题。传统推理框架比如HuggingFace Transformers原生的past_key_values采用连续式显存分配一个请求进来直接按最大序列长度一次性分一整块连续显存。比如配置里max_seq_len设4096那每个请求就预留4096长度的KV空间哪怕它实际只生成128个token剩下的空间也占用着。这在并发量大、序列长度差异大的场景下浪费极其严重。我第一次意识到这个问题是在长文本生成任务里同时跑4个请求有的请求要生成2000 token有的只生成200 token结果200 token那个也占着4096长度的预留空间整体显存浪费接近一半。vLLM论文里给过一个数字连续式缓存在真实生产场景中浪费的显存高达60%-80%。这个数字我实测下来并不夸张。3.2 分页的思路逻辑块与物理块的映射PagedAttention的核心做法是把KV Cache切成固定大小的“块”block每个块装固定数量的token比如16个。一个长序列的KV Cache在逻辑上是连续的但在物理显存中可以散落在任何位置通过一张“块表”block table记录逻辑块到物理块的映射。逻辑块0 - 物理块7 逻辑块1 - 物理块1 逻辑块2 - 物理块4 逻辑块3 - 物理块9这就是操作系统虚拟内存的分页思想进程看到的虚拟地址连续实际物理页可以离散。KV Cache不需要一整块连续显存碎片被天然消除。生成新token时只需要按需分配新的空块填满一块再申请下一块分配粒度从“整个最大序列长度”降到了“16个token”。你还记得频繁分配、释放显存会导致碎片化吗连续式缓存分配时序列结束后释放一整块大显存留下空洞新请求又进来如果找不到足够大的连续空洞就会OOM。PagedAttention的块很小释放和分配都以块为单位碎片问题基本不存在。3.3 Copy-on-Writebeam search和多序列场景的意外红利PagedAttention还有个特别巧妙的设计物理块可以共享。beam search时多个分支序列共享同一个前缀这个前缀的KV Cache就放在同一批物理块里多个beam引用同一份数据完全不用重复分配。只有当一个beam需要在这个共享块上继续写新token时才复制该块不写的beam继续引用原块。这就是写时复制Copy-on-Write机制。这个特性对并行采样同样有用。比如做few-shot评测、多个候选答案生成或者Agent多路探索时多个序列共享同一段system prompt和对话历史KV Cache共享能省出可观的显存。我实际测试过同样8个并发序列共享前缀越长显存占用比朴素实现低得越多。3.4 除了省显存PagedAttention还提升了什么碎片被消除后发现一个连带好处可以用更小的显存跑更大的batch。batch大了GPU计算利用率就上去了。做推理的人都知道小batch是GPU利用率杀手——矩阵太小算力喂不饱大部时间是显存带宽在跑计算单元闲着。vLLM论文公开的数据是在A100上vLLMPagedAttention相比FasterTransformer/TGI吞吐量提升2-4倍长序列、长度分布不均匀的场景收益更明显。我自己的复测结论类似优势不在单请求加速而在“同样的显存能塞下更多并发”吞吐量翻倍不是夸张。另外PagedAttention让按需分配成为可能。传统连续式缓存里一个长序列如果中途用掉所有预留空间就必须立刻终止或重算。PagedAttention没有这种硬性的“预留上限”只要有空闲块就能继续生长长序列更稳。3.5 和FlashAttention是什么关系很多人会把PagedAttention和FlashAttention弄混。简单区分FlashAttention解决的是“attention计算时怎么减少显存读写、利用SRAM”的问题属于算得快PagedAttention解决的是“KV Cache怎么分配、怎么减少碎片”的问题属于装得下。两者并不冲突vLLM里就是PagedAttention管理KV Cache分配而FlashAttention内核做实际注意力计算配合使用。4. 落地实操vLLM部署、参数调优与真实效果4.1 用vLLM替换原生推理代码改动有多小我在项目中把原生HuggingFace推理迁移到vLLM最大的感受是API层几乎零改动收益却是立竿见影的。一个最小可用的部署示例# 启动vLLM服务OpenAI兼容API vllm serve Qwen/Qwen2.5-7B-Instruct \ --gpu-memory-utilization 0.88 \ --max-model-len 8192 \ --max-num-seqs 16 \ --block-size 16 \ --kv-cache-dtype fp8 \ --enable-prefix-caching然后业务侧用OpenAI SDK直接接from openai import OpenAI client OpenAI( base_urlhttp://127.0.0.1:8000/v1, api_keyEMPTY ) resp client.chat.completions.create( modelQwen/Qwen2.5-7B-Instruct, messages[{role: user, content: 讲个冷笑话}], max_tokens512 ) print(resp.choices[0].message.content)对比基准同机同卡原生Transformers推理batch size跑不大就OOMvLLM在同样显存限制下能跑4-8倍的并发序列首token延迟还更稳定。这是PagedAttention直接带来的收益——显存碎片少了KV Cache按需分配同样的显存“塞”下更多请求。4.2 一眼看懂vLLM关键参数gpu_memory_utilization、max-num-seqs、block-size用vLLM不是把参数抄上去就完事这些参数背后的逻辑要明白gpu_memory_utilization控制KV Cache最多占显存的比例。默认0.9但如果你还在这张卡上跑别的任务建议调到0.8-0.85。我踩过坑设0.9跑单模型没问题但同一进程里还有嵌入模型结果预热阶段直接OOM。预留一点余量永远没错。max-num-seqs限制并发序列数。很多人理解成“用户并发上限”其实是“KV Cache可容纳同时计算的序列数”。设越大KV Cache预留越激进显存分配越紧张。如果你的场景是几百路用户并发靠的是服务层的排队机制不要靠无限调大这个值硬扛。block-size每个物理块容纳的token数。默认16一般不用动。块越大分配粒度越粗、管理开销越小但碎片可能多一点块越小越灵活调度开销会上去。vLLM默认值是在吞吐和显存之间调过的手痒的话可以试试8或32但别指望质变。max-model-len限制上下文长度。这是最容易被忽略的显存治理手段。如果你业务实际只需要8K上下文就别把max-model-len设成32K。KV Cache按“最大长度”预留长度设置越保守能跑的并发越多。4.3 KV Cache量化INT8/FP8到底省了多少显存KV Cache默认用FP16存储每元素2字节。如果量化到FP8直接减半INT8同样。但量化KV Cache会引入精度损失生成质量在部分任务上可能下降。用vLLM开启FP8 KV Cache就是启动命令里加一行--kv-cache-dtype fp8。这个选项不是所有GPU都支持。我实测在支持FP8的卡上7B模型、8K上下文、16并发的场景KV Cache占用从约8GB降到4GB显存富余度提升非常明显生成结果的BLEU/ROUGE得分基本持平看不出明显退化。如果GPU不支持FP8还可以考虑另一种思路选择原生就带量化KV支持的工具链。更保守的做法是换GQA架构模型——先把“头数”降下来再谈量化。顺序很重要结构省显存最稳量化省显存要验收精度两者不要反过来。4.4 长上下文和高并发场景的进阶开关Prefix Caching与Chunked Prefill两个容易忽视但实战特别有用的特性Prefix Caching前缀缓存同一段前缀的KV Cache可以被不同请求复用。多轮对话每次请求都带着历史消息Agent的System Prompt经常是几千token的长指令这部分的KV计算完全重复。vLLM开启--enable-prefix-caching后命中缓存的请求首token延迟能降很多。Chunked Prefill预填充分块把长prompt的预填充阶段切成小块跟解码请求交错执行。传统框架里一个长prompt预填充时GPU要独占很长时间后面的解码请求全部排队首token延迟暴涨。Chunked Prefill就是为了避免“一个长请求把整艘船堵死”。不过它本质是让长prompt预填充变慢以换取整体吞吐稳定要不要开取决于你的请求长度分布。4.5 一个容易被忽略的显存观察工具vLLM启动时会打印KV Cache分配的详细报告--print-usage参数也能在运行时看。我养成了每次调参先看三行数字的习惯模型权重占用、KV Cache上限、实际已分配KV Cache。三个数加起来大约就是整卡显存。很多“莫名其妙OOM”的排查从这里开始比看日志更直观。5. 踩坑自查清单KV Cache和PagedAttention实战中常见的几个问题5.1 上下文长度不是越长越好把max-model-len从8K调到32KKV Cache预留立刻翻4倍。不做长文本业务别为了“以防万一”调长上下文。这是显存治理里成本最低、收益最直接的一刀。我见过不少部署事故追根溯源就是“反正支持32K我就设32K”模型没多大KV Cache先把显存放满了。5.2 并发数和KV Cache的预留关系max-num-seqs不是业务并发数而是同时计算的最大序列数。外部来100个请求vLLM会排队真正同时在GPU里的只有max-num-seqs个。你想提高吞吐要综合看显存剩余和gpu_memory_utilization想提高并发体验业务层要做的可能是更聪明的排队和优先级调度而不是硬调这个值。5.3 释放KV Cache长对话服务的隐形杀手用HuggingFace原生态跑服务时past_key_values是显存大户。如果你在循环里不断调model.generate但每次都不重新初始化prompt旧的KV Cache可能还占着显存。我出错的那次就是多轮对话循环里把past_key_values传给下一轮但某些分支路径忘记释放几个小时后服务悄悄OOM。在vLLM这类框架里KV Cache由调度器管理用完自动回收到空闲块池不会漏但如果你还在用原生Transformers写服务务必仔细管理这个对象。5.4 Prefix Caching不是银弹Prefix Caching命中率取决于前缀是否完全一致。多轮对话里只要有一处token不同比如时间戳、随机ID公共前缀就会变短缓存收益下降。想提高命中率设计Prompt时把变化的字段放在消息末尾把固定指令放在开头都是实战里验证有效的做法。5.5 量化KV Cache后一定要跑评测FP8 KV Cache在多数任务上精度损失可接受但代码生成、数学推理这类对数值敏感的任务退化更明显。不要直接上生产先拿你的评测集对比FP16和FP8的输出用实际数据判断。任何量化方案都遵循这个原则——省显存的目的是为了更多并发或更长上下文如果精度退化让业务不可用省下来的显存没有意义。5.6 显存不足时的降级路径显存捉襟见肘时我按这个顺序排查先砍上下文长度再换GQA架构模型再看要不要量化KV Cache最后才是换更小的模型。顺序逻辑是前两步不影响生成质量第三步需要验收第四步直接影响模型能力。不要一上来就换小模型——明明7B能干的活换成3B效果降一个档次还不一定省多少显存。最后聊点我自己的体会。KV Cache和PagedAttention这两个概念初看像两个孤立的“优化技巧”实际是同一件事的两面一个解决“算两遍”的浪费一个解决“分不到”的浪费。理解了这两个浪费再看vLLM的配置参数、看模型的GQA头数、看各种KV Cache量化方案都会通透很多。如果你准备部署自己的推理服务我建议第一步不是抄启动命令而是先把第2节那张显存账单算清楚——你的场景、你的并发、你的上下文长度KV Cache要吃多少心里有数了后面所有参数都不会调得太离谱。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →