Transformer推理中KV cache共享状态的隔离问题与实验分析

1. 从一道口令实验说起：共享状态到底藏了什么猫腻

前段时间我在调一个基于 Transformer 的小型推理服务，场景很典型：多个用户请求并发进来，模型需要做自回归生成，每一步都要复用之前算好的 KV cache。按理说这是个标准操作，但实测下来发现一个诡异现象——两个完全独立的请求，生成到第三四个 token 的时候，输出开始互相"串味"，A 请求的回复里冒出了 B 请求的上下文片段。

一开始我怀疑是并发写冲突，加了锁之后问题依旧。后来把请求改成串行，问题消失。这就说明根子不在计算逻辑，而在共享状态的生命周期管理上。KV cache 本质上是一块跨请求复用的内存区域，如果隔离没做好，前一个请求的缓存没清干净，后一个请求就会"继承"到不属于它的历史。

这个现象让我想起一个更朴素的实验：给两个不同的会话分配同一个缓存槽位，然后观察它们的注意力输出会不会互相污染。结果非常明确——只要共享了 KV cache 的存储空间，哪怕模型权重完全一致，输出也会出现可观测的偏移。这就是标题里说的"共享状态，隔离问题"：共享带来性能，隔离保证正确，两者之间的平衡点就是整个推理系统的命门。

这篇文章我想把这件事掰开揉碎讲清楚。核心围绕三个问题展开：KV cache 在 Transformer 推理里到底扮演什么角色，共享状态会引入哪些具体的隔离风险，以及在实际工程里怎么用一套可复现的实验方法把问题定位出来。适合正在做 Transformer 推理优化、多路并发服务、或者单纯想搞懂注意力机制底层状态的读者。不管你是刚接触 Transformer 的新手，还是已经在调推理框架的老手，这里面的排查思路和参数细节应该都能直接用上。

2. KV cache 与注意力机制：共享状态的物理来源

2.1 自回归生成为什么离不开 KV cache

要理解共享状态的风险，得先搞清楚 KV cache 是怎么来的。Transformer 做自回归生成时，每生成一个新 token，都需要拿当前 token 的 Query 去和前面所有 token 的 Key 做点积，算出注意力权重，再对 Value 加权求和。问题在于，前面那些 token 的 Key 和 Value 在每一步都是不变的，如果每生成一个 token 就重新算一遍全序列的 K 和 V，计算量会随着序列长度平方级增长。

KV cache 的思路很直接：把已经算过的 Key 和 Value 存下来，下一步只算新 token 的 K、V，然后拼接到缓存里。这样每步的计算量从 O(n²) 降到 O(n)，显存换时间，这是所有主流推理框架的标配。

用生活化的类比：你在写一篇长文，每写一句都要回看前面所有句子。KV cache 就像是你把前面每句话的"要点卡片"提前整理好放在桌上，写新句子时直接翻卡片，而不是把整篇文章重读一遍。卡片就是缓存，桌子就是那块共享内存。

关键点在于，这张"桌子"是有限的。多个请求同时进来，如果每个请求都独占一张桌子，显存很快爆掉；如果大家共用一张桌子，就得有一套严格的隔离规则，否则卡片会混。

2.2 注意力机制里哪些状态是可以共享的

注意力计算可以拆成几个部分：Query 投影、Key 投影、Value 投影、注意力分数、加权输出。其中 Query 是当前 token 独有的，没法共享；Key 和 Value 是历史 token 的产物，理论上可以跨请求复用——但前提是这些历史 token 属于同一个请求。

这里有个容易踩的坑：很多人以为只要模型权重相同，KV cache 就能通用。实际上 KV cache 里存的是特定输入序列经过投影后的中间结果，它和输入内容强绑定。不同请求的输入不同，缓存内容就不同，混用等于把别人的历史塞进你的上下文。

多头注意力机制让这件事更复杂。每个注意力头有独立的 K、V 投影矩阵，缓存要按头分别存储。如果共享时头维度对不齐，或者头数配置不一致，就会出现部分头读到错误缓存的情况，表现为输出在某些语义维度上"发疯"。

2.3 共享状态带来的性能收益到底有多大

先说收益，不然没人愿意折腾隔离。以常见的 7B 级别模型、序列长度 2048 为例，单请求的 KV cache 占用大致可以这样估算：层数乘以头数乘以头维度乘以序列长度乘以 2（K 和 V）乘以精度字节数。粗算下来单个请求的缓存可能到几百 MB 级别。如果并发 16 路，不做任何共享，显存需求直接翻十几倍。

共享状态的核心收益有两个：一是显存复用，多个请求可以共用预分配的缓存池，避免频繁申请释放；二是前缀共享，如果多个请求有相同的系统提示词，这部分前缀的 KV 可以只算一次，所有请求复用。后者在对话类应用里收益巨大，系统提示动辄几百上千 token，复用后首 token 延迟能降一大截。

但收益越大，隔离责任越重。前缀共享的前提是"前缀完全一致"，一旦前缀有任何一个 token 不同，共享的缓存就是错的。这就是口令实验要验证的东西。

3. 口令实验设计：如何用最小成本暴露隔离缺陷

3.1 实验目标与核心假设

这个实验的目标很明确：验证在共享 KV cache 的条件下，不同请求之间是否会发生状态泄漏。核心假设是——如果隔离机制正确，那么无论请求是串行还是并发、无论缓存是否复用，相同输入的输出应该完全一致；如果隔离有缺陷，并发或复用场景下输出会出现可观测的偏差。

我把这个假设拆成三个可验证的子命题。第一，串行执行时输出稳定，作为基线。第二，并发执行但每个请求独立分配缓存时，输出应与基线一致。第三，并发执行且共享缓存池时，如果隔离正确，输出仍应与基线一致；如果出现偏差，偏差的位置和幅度就是隔离缺陷的指纹。

3.2 实验环境与参数配置

环境不需要多复杂，一台带 GPU 的机器，装好 PyTorch 和 transformers 库就够。模型选一个小尺寸的，方便快速迭代，比如 6 层、8 头、隐藏维度 512 的配置，自己搭或者用现成的小模型都行。关键是层数和头数要够，能体现多头注意力的隔离问题。

参数上我建议固定这几项：序列长度设为 128，生成长度设为 32，batch size 分别测 1 和 8，精度用 float16 和 float32 各跑一遍对比。温度设为 0，也就是贪心解码，这样输出是确定性的，方便对比。随机种子固定，排除采样带来的干扰。

注意：温度一定要设 0。如果开了采样，输出本身就有随机性，你根本分不清偏差是隔离问题还是采样噪声。这是排查这类问题最容易犯的错。

3.3 构造"口令"输入的设计思路

所谓口令，就是一组精心设计的输入序列，它们之间既有共享部分，又有差异部分。我设计了三组：

• 组 A：前缀 "系统提示：你是一个助手。" + 问题 "1+1等于几"
• 组 B：前缀 "系统提示：你是一个助手。" + 问题 "2+2等于几"
• 组 C：前缀 "系统提示：你是一个翻译。" + 问题 "1+1等于几"

组 A 和组 B 共享前缀，只有问题不同；组 A 和组 C 问题相同，只有前缀不同。这样设计的好处是，如果前缀共享的隔离有问题，A 和 B 会互相污染；如果整体缓存隔离有问题，A 和 C 会互相污染。通过对比不同组合的输出，就能定位问题出在哪一层。

每组输入都跑三遍：单独跑、和共享前缀的伙伴一起跑、和不同前缀的伙伴一起跑。记录每次输出的 token 序列和对应的注意力权重分布。

4. 实操过程：从基线到污染的全链路复现

4.1 搭建最小可复现的推理脚本

先写一个最朴素的推理循环，不涉及任何缓存优化，作为基线。核心逻辑就是标准的自回归：每步把完整序列喂进去，取最后一个位置的 logits，argmax 得到下一个 token，拼回序列，重复。

import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "your-small-model" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32) model.eval() def generate_baseline(prompt, max_new_tokens=32): inputs = tokenizer(prompt, return_tensors="pt") input_ids = inputs["input_ids"] for _ in range(max_new_tokens): with torch.no_grad(): outputs = model(input_ids) next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) input_ids = torch.cat([input_ids, next_token], dim=-1) return tokenizer.decode(input_ids[0], skip_special_tokens=True)

这个版本没有任何缓存，每步都重算全序列，慢但绝对正确。把它作为黄金标准，后面所有优化版本的输出都要和它对齐。

4.2 引入 KV cache 后的输出对比

接下来加上 KV cache。HuggingFace 的模型默认支持use_cache=True，内部会维护 past_key_values。改造后的循环只喂新 token，把缓存传进去。

def generate_with_cache(prompt, max_new_tokens=32): inputs = tokenizer(prompt, return_tensors="pt") input_ids = inputs["input_ids"] past_key_values = None generated = input_ids for _ in range(max_new_tokens): with torch.no_grad(): outputs = model( input_ids if past_key_values is None else input_ids[:, -1:], past_key_values=past_key_values, use_cache=True ) past_key_values = outputs.past_key_values next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) input_ids = next_token generated = torch.cat([generated, next_token], dim=-1) return tokenizer.decode(generated[0], skip_special_tokens=True)

单请求跑下来，输出和基线完全一致，说明缓存逻辑本身没问题。这一步是必须验证的，否则后面出了问题你分不清是缓存实现错了还是隔离错了。

4.3 并发共享缓存池的模拟实现

关键步骤来了。我手动实现一个简易的缓存池，让多个请求共享同一块 past_key_values 存储空间。做法是预分配一个大的张量，按槽位切分，每个请求拿到一个槽位索引。为了模拟隔离缺陷，我故意让两个请求在某些条件下指向同一个槽位。

class SimpleCachePool: def __init__(self, num_slots, num_layers, num_heads, head_dim, max_seq_len): self.slots = [ { "key": torch.zeros(num_layers, num_heads, max_seq_len, head_dim), "value": torch.zeros(num_layers, num_heads, max_seq_len, head_dim), "length": 0 } for _ in range(num_slots) ] self.assignment = {} def assign(self, request_id, slot_id): self.assignment[request_id] = slot_id def get(self, request_id): return self.slots[self.assignment[request_id]]

正常隔离时，每个 request_id 分配不同 slot。模拟缺陷时，让两个 request_id 指向同一个 slot，观察输出变化。实测下来，一旦共享槽位，第二个请求的输出在生成到第 5 到第 8 个 token 时开始偏离基线，偏离位置和缓存被写入的长度相关。

4.4 污染现象的具体表现与数据记录

把实验结果整理成表，偏差用输出 token 序列和基线的编辑距离衡量。

偏差的表现很有意思：不是完全乱掉，而是"局部串味"。A 的输出里会插入 B 的问题关键词，但整体语法还是通的。这说明注意力机制在读到错误缓存时，会把它当成合法的上下文来用，只是语义对不上。这种"看起来正常但实际错误"的污染最危险，因为线上很难通过简单的异常检测发现。

提示：判断是否发生缓存污染，不要只看最终输出。要在每一步记录注意力权重，看 Query 对 Key 的注意力分布是否出现了不该有的峰值。污染往往先在注意力分布上体现，输出偏差是滞后的。

5. 隔离问题的根因分析与排查技巧

5.1 共享与隔离的边界到底该划在哪

从实验数据反推，隔离的边界应该划在"请求"这一层，而不是"模型"或"进程"层。同一个请求内部，前缀和生成部分可以共享缓存，因为它们的上下文是连续的；不同请求之间，哪怕前缀完全一样，也不能共享可写的缓存区域，只能共享只读的、已经计算完成的前缀缓存。

这里有个细节：只读共享是安全的，因为前缀计算完成后不再变化。但一旦某个请求开始写入自己的生成内容，这块区域就必须私有化。工程上常见的做法是写时复制——前缀部分共享，第一个写操作触发复制，后续写入都在私有副本上进行。

5.2 常见隔离缺陷的三种典型模式

踩过的坑总结下来有三类。第一类是槽位复用未清零，请求结束后缓存槽位没有重置 length 指针，下一个请求从错误的偏移开始写。第二类是前缀哈希碰撞，用前缀内容的哈希值做缓存键，哈希冲突导致不同前缀映射到同一缓存。第三类是并发写竞争，两个请求同时写同一槽位，后写的覆盖先写的，表现为输出随机性抖动。

这三类的排查方法不同。第一类看 length 指针的日志，第二类看哈希键的分布，第三类加锁后复测。我一般先加全局锁跑一遍，如果问题消失就是第三类，如果还在就是前两类。

5.3 排查工具与日志埋点建议

光靠输出对比不够，得在缓存读写的关键路径上埋点。我习惯记录这几个字段：请求 ID、槽位 ID、写入起始偏移、写入长度、时间戳。把这些日志按时间排序，就能还原出缓存的完整生命周期。

另外建议在注意力计算后 dump 一份注意力权重的统计量，比如每层的最大注意力值、熵、以及 top-3 注意力位置的 token ID。污染发生时，这些统计量会出现异常，比逐 token 对比高效得多。

6. 从实验到工程：可落地的隔离方案

6.1 缓存池的分层设计

工程上我推荐三层结构：全局只读前缀池、请求私有缓存区、临时计算缓冲区。前缀池存所有请求共享的系统提示缓存，只读；私有区每个请求独占，存生成内容；临时区用于中间计算，用完即释放。这样既拿到了前缀共享的收益，又保证了生成部分的隔离。

前缀池的键用前缀 token 序列的精确匹配，不用哈希，避免碰撞。匹配时逐 token 比对，虽然慢一点，但前缀数量通常不多，开销可接受。

6.2 写时复制与前缀共享的配合

写时复制的触发点要精确。我的做法是给每个缓存块加一个引用计数，前缀块初始引用计数为共享它的请求数。任何请求要写入时，先检查引用计数，大于 1 就复制一份私有副本，计数减一，写入副本。这样前缀块始终只读，私有块始终独占。

实测下来，这套机制在 16 路并发、共享 512 token 前缀的场景下，显存占用比全私有方案降低约 40%，首 token 延迟降低约 30%，而输出与基线完全一致。

6.3 隔离正确性的自动化回归测试

最后一定要有回归测试。我写了一个测试脚本，固定一组输入，分别跑串行、并发独立、并发共享三种模式，断言三者输出完全一致。这个脚本进 CI，每次改缓存逻辑都跑一遍。别嫌麻烦，缓存隔离的 bug 一旦漏到线上，排查成本是测试成本的几十倍。

测试用例要覆盖边界：空前缀、超长前缀、前缀完全相同、前缀差一个 token、并发数从 1 到最大。我踩过的坑是只测了并发数为 2 的情况，结果并发数到 8 时槽位分配逻辑出错，线上才暴露。

7. 一些实操心得与后续可扩展的方向

这套口令实验我前后迭代了四五版，最大的体会是：共享状态的正确性不能靠"看起来对"来判断，必须靠可复现的对比实验。很多隔离缺陷在单请求下完全正常，只有并发和复用场景才暴露，而线上恰恰都是并发场景。

另一个心得是关于注意力权重的。以前我只看输出，觉得输出对了就没问题。后来发现注意力分布才是更敏感的指标，有些污染在输出上还没体现，注意力权重已经异常了。现在我做缓存相关的改动，都会顺手 dump 一份注意力统计，作为额外的健康检查。

这个实验框架还能往几个方向扩展。一是加上不同精度（float16、int8）的对比，看量化是否会影响隔离行为；二是扩展到多模态场景，视觉 Transformer 里的特征图缓存也有类似的共享隔离问题，垂直和水平方向的注意力缓存如果混用，同样会污染；三是把实验自动化，做成一个缓存隔离的基准测试集，方便不同推理框架横向对比。

如果你也在做推理优化，建议先把这个最小实验跑通，再上复杂的缓存池。基础不牢，后面调优全是玄学。