从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_model=512，num_heads=8，那每个头的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_layers=32，num_heads=32，d_k=128，seq_len=2048，batch_size=1，fp16精度，算下来大概是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份。为了省显存，后来出现了MQA（Multi-Query Attention），所有头共享一份K和V，只有Q是每个头独立的。这样KV cache直接缩小到1/num_heads。再后来是GQA（Grouped-Query Attention），把num_heads分成若干组，每组共享一份K和V，是MQA和MHA的折中。

这三种方案的对比可以用一个表格说清楚：

实测下来，GQA在大多数场景下是性价比最高的选择。Llama 2 70B用的就是GQA，num_heads=64，num_groups=8，KV 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_size+1的数组，cu_seqlens[i]表示第i个序列的起始偏移，cu_seqlens[i+1]-cu_seqlens[i]就是第i个序列的长度。所有序列的token拼成一个长序列，注意力计算的时候通过cu_seqlens来划分边界，保证每个序列只attend到自己内部的token。这个设计很巧妙，既避免了padding，又保证了隔离性。

但cu_seqlens的构造容易出错。如果某个序列长度为0，cu_seqlens会出现连续相同的值，kernel里做除法或者索引的时候可能出问题。我在实现的时候就遇到过这个，一个空请求导致整个batch的输出全乱。后来加了一个过滤，把长度为0的序列直接跳过，才解决。

4. 口令实验：用最小化方案复现共享状态问题

4.1 实验设计思路与口令构造

这个实验的目标是用最小的代价复现“共享状态导致隔离失效”的现象。不需要大模型，用一个小型的Transformer（比如4层，d_model=128，num_heads=4）就够了。关键是构造一个口令文本，让不同请求的前缀有重叠但不完全相同，这样才能触发共享逻辑。

口令的设计原则是：前缀高度相似，后缀有区分度。比如请求A的口令是“请把以下内容翻译成英文：今天天气很好”，请求B的口令是“请把以下内容翻译成英文：今天天气很差”。两个请求的前缀“请把以下内容翻译成英文：今天天气”是完全相同的，只有最后一个词不同。如果前缀共享逻辑正确，两个请求应该各自输出正确的翻译；如果共享逻辑有bug，可能会出现A的输出被B污染，或者两个请求的输出相同。

我实际用的口令更长一些，前缀大概50个token，后缀10个token。这样前缀共享的收益明显，同时后缀的差异足够触发隔离问题。实验的观察指标有三个：输出的正确性、KV cache的block引用计数、以及注意力权重的分布。

4.2 实验环境搭建与关键配置

环境用Python 3.10，PyTorch 2.1，CUDA 12.1。模型用一个自己实现的MiniTransformer，不依赖任何推理框架，这样能完全控制KV cache的管理逻辑。代码结构大概是：一个Model类，一个KVCache类，一个Scheduler类。KVCache类里实现block的分配、引用计数、copy-on-write。

关键配置参数如下：

• block_size = 16，每个block存16个token的KV
• num_blocks = 256，总共有256个block
• max_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相关的典型故障速查表

这个表是我在实际排查中总结的，每一条都对应至少一次真实的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，可以关注一下这个指标。