从KV cache共享到批量调度：注意力机制下的隔离问题与口令实验

1. 从“共享状态”说起一个被忽视的工程陷阱“共享状态隔离问题”这个说法第一次听到是在一次内部技术复盘会上。当时团队在排查一个推理服务的性能抖动问题现象很诡异单条请求跑得好好的一旦并发上来延迟就飙升而且不是线性上升是那种突然翻倍的跳变。日志里看不出任何报错CPU和内存也都在合理范围。最后定位到的根因是KV cache的共享策略出了问题——多个请求共用了同一块缓存区域导致注意力计算时读到了不属于自己的历史状态。这件事让我意识到一个很朴素但容易被忽略的道理共享状态能省资源但也会把本来隔离的问题搅在一起。KV cache、批量调度、注意力机制这几个词看起来是三个独立的技术点实际上它们在推理系统里是咬合在一起的。你动了其中一个另外两个的行为就会跟着变。这篇内容就想把这条链路拆开讲清楚从注意力机制里KV是怎么来的到KV cache为什么要共享再到批量调度怎么在共享和隔离之间做权衡最后落到一个可以自己动手复现的口令实验上。适合谁来读如果你正在做推理服务的性能优化或者你在看vLLM、TensorRT-LLM这类框架的源码又或者你只是好奇“为什么大模型推理这么吃显存”那这篇内容应该能给你一些可以直接上手的东西。不需要你精通CUDA但需要你对Transformer的基本结构有个概念。我会尽量用生活化的类比把原理讲透再给可复现的操作步骤。先给一个全局的图景。注意力机制的核心是Q、K、V三个矩阵Q是“我在找什么”K是“我有什么可以被找到”V是“我实际提供的内容”。自注意力机制里每个token都会生成自己的Q、K、V然后用自己的Q去和所有token的K做点积得到注意力权重再对V加权求和。这个过程里K和V是会被反复用到的——当前token要和前面所有token的K做计算也要用前面所有token的V做加权。所以K和V天然就有“缓存”的价值这就是KV cache的由来。问题在于当多个请求同时进来的时候它们的KV cache怎么管理如果每个请求都独立分配一块显存很快就爆了如果共享又会出现我开头说的那种“读到别人状态”的问题。批量调度就是在这个矛盾里找平衡的机制。而“口令实验”是我设计的一个最小化复现方案用一段固定的口令文本构造多请求并发的场景观察KV cache共享策略对输出结果的影响。这个实验不需要大模型用一个小型的Transformer就能跑出可观察的差异。2. 注意力机制里的KV到底是怎么算出来的2.1 从自注意力到多头注意力QKV的生成逻辑要理解KV cache得先把QKV的生成过程捋一遍。假设输入序列长度为L隐藏维度为d_model。自注意力机制的第一步是线性变换用三个不同的权重矩阵W_Q、W_V、W_K把输入X映射成Q、K、V三个矩阵。形状上X是(L, d_model)W_Q是(d_model, d_k)乘完得到Q是(L, d_k)。K和V同理。这里有个细节很多人会忽略Q、K、V的维度不一定相同。在标准的多头注意力机制里d_k和d_v通常相等但和d_model的关系是d_k d_model / num_heads。比如d_model512num_heads8那每个头的d_k就是64。多头自注意力机制原理的核心就是把d_model切成num_heads份每一份独立做注意力计算最后拼接再过一个线性层。这样做的目的是让不同的头关注不同的子空间有的头可能关注语法结构有的头可能关注语义相似度。自注意力机制QKV的计算公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。这个sqrt(d_k)的缩放很关键因为当d_k很大的时候QK^T的点积结果会很大softmax之后会变得非常尖锐梯度就消失了。除以sqrt(d_k)相当于把方差拉回到1附近让softmax的输出分布更平滑。因果自注意力是在这个基础上加了一个mask把未来位置的注意力权重置为负无穷softmax之后变成0。这样每个token只能看到自己和前面的token保证了自回归生成的因果性。这个mask是一个上三角矩阵对角线及以下为0以上为负无穷。2.2 KV cache的本质用空间换时间的经典操作理解了QKV的生成KV cache就很好懂了。在自回归生成的时候每生成一个新token都需要用它的Q去和前面所有token的K做点积。如果没有缓存每生成一个token就要重新计算前面所有token的K和V计算量是O(L^2)的累积。有了KV cache前面token的K和V算一次就存下来新token只需要算自己的Q、K、V然后用Q去和缓存里的K做点积再把新的K、V追加到缓存里。这个操作的本质是用显存换计算。缓存的大小是2 * batch_size * num_layers * num_heads * seq_len * d_k * sizeof(dtype)。以一个7B模型为例num_layers32num_heads32d_k128seq_len2048batch_size1fp16精度算下来大概是2 * 1 * 32 * 32 * 2048 * 128 * 2字节约1GB。如果batch_size开到32就是32GB直接吃满一张卡。这就是为什么KV cache是推理显存的大头。这里有个容易踩的坑KV cache的布局对性能影响极大。常见的布局有两种一种是(num_layers, batch, num_heads, seq_len, d_k)另一种是(num_layers, batch, seq_len, num_heads, d_k)。前者在追加新token的时候seq_len维度是连续的写入效率高后者在注意力计算的时候num_heads和d_k连续读取效率高。实际框架里会根据硬件和kernel实现做取舍比如FlashAttention就更偏好后者。你在自己实现的时候如果发现性能不对先检查一下cache的memory layout。2.3 多头注意力机制下的KV共享MQA和GQA标准的多头自注意力机制里每个头都有自己的K和V缓存量是num_heads份。为了省显存后来出现了MQAMulti-Query Attention所有头共享一份K和V只有Q是每个头独立的。这样KV cache直接缩小到1/num_heads。再后来是GQAGrouped-Query Attention把num_heads分成若干组每组共享一份K和V是MQA和MHA的折中。这三种方案的对比可以用一个表格说清楚方案Q的头数K/V的头数KV cache大小表达能力MHAnum_headsnum_heads最大最强GQAnum_headsnum_groups中等中等MQAnum_heads1最小最弱实测下来GQA在大多数场景下是性价比最高的选择。Llama 2 70B用的就是GQAnum_heads64num_groups8KV cache缩小到1/8而效果损失很小。这个取舍的逻辑是Q负责“提问”需要多样性K和V负责“被检索”共享一些不会太影响检索质量。3. 批量调度共享与隔离的博弈场3.1 连续批处理与PagedAttention批量调度的核心问题是多个请求的KV cache怎么组织。最朴素的做法是每个请求预分配一块最大长度的连续显存但这样浪费严重因为大多数请求的实际长度远小于最大长度。vLLM提出的PagedAttention借鉴了操作系统的虚拟内存分页思想把KV cache切成固定大小的block每个请求的cache由若干block组成block之间不需要连续。这样就消除了外部碎片显存利用率能到90%以上。连续批处理是另一个关键机制。传统批处理要等一个batch里所有请求都完成才能开始下一批而连续批处理是动态的一个请求生成完了立刻把新的请求插进来不用等整个batch结束。这样GPU的利用率大幅提升因为不会出现“长请求拖死短请求”的情况。但连续批处理引入了一个新问题不同请求的KV cache在物理上可能是交错的。PagedAttention的block table记录了每个请求的逻辑block到物理block的映射注意力计算的时候要通过block table去查。这个查表操作如果实现不好会成为性能瓶颈。vLLM的kernel是把block table传到GPU上在kernel内部做地址计算避免了CPU和GPU之间的频繁同步。3.2 前缀共享共享状态带来的收益与风险前缀共享是共享状态最典型的应用。如果多个请求有相同的前缀比如相同的system prompt那这段前缀的KV cache可以只存一份所有请求共享。这样显存节省是线性的N个请求共享长度为P的前缀就省了(N-1)*P的cache。但共享状态的风险也在这里。如果共享的cache被某个请求修改了比如追加了新token其他请求就会读到脏数据。所以共享的cache必须是只读的或者用copy-on-write机制。vLLM的实现里共享的block有一个引用计数当某个请求要写入时如果引用计数大于1就先复制一份再写。这个机制保证了隔离性但复制的开销在写入频繁的时候会累积。我踩过的一个坑是前缀共享的hash碰撞。有些实现用前缀的hash值作为key来查找可共享的block如果hash函数不好不同前缀可能映射到同一个key导致错误的共享。这个bug很隐蔽因为大多数时候输出看起来是对的只有在特定输入组合下才会出错。排查的时候可以用一个固定的口令文本构造两个只有细微差异的前缀观察输出是否互相污染。3.3 批量调度中的注意力计算变长序列的处理批量调度里另一个难点是变长序列的注意力计算。一个batch里不同请求的序列长度不同如果padding到最大长度浪费计算如果不padding就需要用变长的kernel。FlashAttention的varlen版本就是干这个的它用一个cu_seqlens数组记录每个序列的起始位置在kernel内部按实际长度计算。这里的关键参数是cu_seqlens它是一个长度为batch_size1的数组cu_seqlens[i]表示第i个序列的起始偏移cu_seqlens[i1]-cu_seqlens[i]就是第i个序列的长度。所有序列的token拼成一个长序列注意力计算的时候通过cu_seqlens来划分边界保证每个序列只attend到自己内部的token。这个设计很巧妙既避免了padding又保证了隔离性。但cu_seqlens的构造容易出错。如果某个序列长度为0cu_seqlens会出现连续相同的值kernel里做除法或者索引的时候可能出问题。我在实现的时候就遇到过这个一个空请求导致整个batch的输出全乱。后来加了一个过滤把长度为0的序列直接跳过才解决。4. 口令实验用最小化方案复现共享状态问题4.1 实验设计思路与口令构造这个实验的目标是用最小的代价复现“共享状态导致隔离失效”的现象。不需要大模型用一个小型的Transformer比如4层d_model128num_heads4就够了。关键是构造一个口令文本让不同请求的前缀有重叠但不完全相同这样才能触发共享逻辑。口令的设计原则是前缀高度相似后缀有区分度。比如请求A的口令是“请把以下内容翻译成英文今天天气很好”请求B的口令是“请把以下内容翻译成英文今天天气很差”。两个请求的前缀“请把以下内容翻译成英文今天天气”是完全相同的只有最后一个词不同。如果前缀共享逻辑正确两个请求应该各自输出正确的翻译如果共享逻辑有bug可能会出现A的输出被B污染或者两个请求的输出相同。我实际用的口令更长一些前缀大概50个token后缀10个token。这样前缀共享的收益明显同时后缀的差异足够触发隔离问题。实验的观察指标有三个输出的正确性、KV cache的block引用计数、以及注意力权重的分布。4.2 实验环境搭建与关键配置环境用Python 3.10PyTorch 2.1CUDA 12.1。模型用一个自己实现的MiniTransformer不依赖任何推理框架这样能完全控制KV cache的管理逻辑。代码结构大概是一个Model类一个KVCache类一个Scheduler类。KVCache类里实现block的分配、引用计数、copy-on-write。关键配置参数如下block_size 16每个block存16个token的KVnum_blocks 256总共有256个blockmax_seq_len 128最大序列长度share_prefix True开启前缀共享copy_on_write True开启写时复制这些参数的选择逻辑是block_size太小会导致block数量过多管理开销大太大则内部碎片严重。16是一个经验值和vLLM的默认值一致。num_blocks要足够大能容纳所有并发请求的cache但也不能太大否则显存浪费。max_seq_len根据口令长度设定留一些余量。4.3 实验步骤与观察到的现象实验分三组第一组关闭前缀共享每个请求独立分配cache第二组开启前缀共享但关闭copy-on-write第三组开启前缀共享和copy-on-write。每组跑100次每次并发2个请求记录输出和cache状态。第一组的结果是预期的两个请求输出都正确cache各自独立引用计数都是1。第二组出现了问题大约30%的情况下两个请求的输出相同而且都是请求A的翻译结果。检查cache发现请求B的block table指向了请求A的block因为前缀共享的时候没有做引用计数请求B直接复用了请求A的block写入的时候覆盖了请求A的数据。第三组的结果是正确的两个请求输出都正确共享的block引用计数为2写入时触发了copy-on-write各自有了独立的block。这个实验最直观的收获是共享状态本身不是问题问题是共享状态的生命周期管理。引用计数和copy-on-write是两个必须的机制缺一不可。引用计数保证知道有多少人在用copy-on-write保证写入时不污染别人。4.4 从实验结果反推注意力计算的影响实验里还观察了注意力权重的变化。在第二组出问题的时候请求B的注意力权重分布和请求A几乎一样因为它的K和V就是请求A的。这导致请求B的Q虽然在问“今天天气很差怎么翻译”但K和V提供的是“今天天气很好”的上下文输出自然就错了。这个现象说明KV cache的隔离性直接决定了注意力计算的正确性。注意力机制的本质是Q和K的匹配如果K被污染了匹配结果就错了后面的加权求和再精确也没用。所以做推理优化的时候省显存的前提是不能破坏隔离性。任何共享策略都要回答一个问题这段共享的状态会不会被某个请求的写入操作影响5. 常见问题与排查技巧实录5.1 KV cache相关的典型故障速查表现象可能原因排查方法解决方案并发时输出乱码cache共享未做隔离检查block引用计数开启copy-on-write显存占用随并发线性增长未开启前缀共享打印block分配日志开启前缀共享长序列推理变慢cache布局不优profile注意力kernel调整memory layout输出重复或截断cu_seqlens构造错误检查变长序列边界过滤空序列显存碎片严重预分配连续显存查看显存分配图改用PagedAttention这个表是我在实际排查中总结的每一条都对应至少一次真实的debug经历。其中“输出重复或截断”那条最隐蔽因为cu_seqlens的错误不会报异常只是计算结果不对。我的做法是在构造cu_seqlens之后加一个assert检查所有值是否单调递增且最后一个值等于总token数。5.2 共享状态下的调试技巧调试共享状态的问题最有效的方法是给每个请求打上唯一标识然后在cache的读写路径上打日志。日志里记录请求ID、block ID、操作类型读/写、引用计数。这样一旦出现污染能直接定位到是哪个请求在什么时候写了哪个block。另一个技巧是用固定口令做回归测试。每次修改cache管理逻辑都跑一遍口令实验对比输出是否和基线一致。这个测试很快几秒钟就能跑完但能抓住大多数隔离性问题。我把这个测试集成到了CI里每次提交都跑省了很多事后排查的时间。还有一个容易忽略的点copy-on-write的触发时机。有些实现是在写入前检查引用计数有些是在分配时检查。前者更精确但每次写入都要检查开销大后者开销小但可能多复制一些不必要的block。实测下来写入前检查在大多数场景下更划算因为写入操作远少于读取操作。5.3 批量调度的参数调优经验批量调度里最影响性能的参数是max_num_seqs和max_num_batched_tokens。前者控制一个batch里最多有多少个请求后者控制一个batch里最多有多少个token。这两个参数决定了GPU的利用率。我的调优经验是先用小batch跑通再逐步加大。max_num_seqs从8开始每次翻倍观察吞吐和延迟的变化。通常会在某个点之后吞吐不再增长而延迟开始飙升那个点就是拐点。max_num_batched_tokens的设置要参考模型的max_position_embeddings不要超过它否则会触发位置编码的越界。还有一个参数是block_size前面提到过。16是默认值但如果你的序列普遍很短比如平均长度小于32可以调到8减少内部碎片如果序列很长比如大于512可以调到32减少block数量和管理开销。这个需要根据实际负载来定没有万能值。6. 从口令实验延伸出的工程思考口令实验虽然简单但它揭示的问题在真实系统里会被放大。真实场景下并发请求成百上千前缀共享的粒度更细cache的block数量更多引用计数的管理更复杂。任何一个环节的疏漏都可能导致输出污染而且这种污染往往是概率性的难以复现。我在实际项目里总结了几条原则。第一共享的状态必须是不可变的或者有明确的写时复制机制。任何“先共享再修改”的设计都要慎重因为修改的副作用会传播到所有共享者。第二隔离性要用测试来保证不能靠代码审查。人的注意力有限复杂的引用计数逻辑很容易看漏但一个自动化的口令测试能每次都跑一遍。第三性能优化要在正确性之后。先保证每个请求的输出是对的再去优化显存和吞吐。反过来做很容易为了省显存而引入难以排查的bug。这个口令实验的代码我后来扩展成了一个小的测试框架支持自定义口令、自定义并发数、自定义共享策略。每次改cache管理逻辑就跑一遍这个框架看输出是否一致。这个习惯帮我避免了好几次线上事故。如果你也在做推理优化建议你也搭一个类似的测试成本不高但收益很大。最后分享一个小技巧在观察注意力权重的时候可以把每个头的权重分布画出来。正常情况下不同的头应该有不同的关注模式如果发现多个头的权重分布几乎一样那可能是KV cache共享过度导致头的多样性丢失了。这个现象在MQA里比较常见GQA会好一些。如果你在用MQA可以关注一下这个指标。