1. 从“口令实验”说起一个被忽视的共享状态陷阱第一次看到“共享状态隔离问题”这个说法是在一次内部技术复盘会上。当时有个同学分享了一个“口令实验”——他写了一段极简的 Transformer 推理代码用同一组权重、同一份输入跑两次结果居然不一样。第一次输出是[0.231, 0.874, ...]第二次变成了[0.229, 0.881, ...]。差异在小数点后第三位肉眼几乎看不出来但用torch.allclose一比对直接返回False。这个现象在视觉任务里可能只是精度抖动但在自回归生成任务里就是灾难。因为生成是逐 token 递推的前一步的微小偏差会被后续步骤放大最终导致完全不同的输出序列。他花了整整两天排查最后发现问题出在KV cache 的共享状态上——他在实现多头注意力时把 K 和 V 的缓存张量在多个注意力头之间共享了同一块内存而不同头在计算时又对这块内存做了原地更新in-place update。第一次推理时缓存被写脏了第二次推理读到的就是脏数据。这个“口令实验”揭开的是 Transformer 推理优化中一个非常隐蔽的角落共享状态带来的隔离问题。KV cache 本身是为了加速自回归解码而设计的缓存机制它把已经计算过的 Key 和 Value 存下来避免每一步都重新计算整个前缀。但一旦缓存的生命周期管理不当或者多个请求、多个头、多个层之间共享了不该共享的状态就会引发难以复现的 bug。这篇文章适合三类人看一是正在手写 Transformer 推理代码、被 KV cache 折磨过的工程师二是做模型部署、遇到“同样输入输出不一致”问题的同学三是对注意力机制底层实现感兴趣、想搞清楚“缓存到底缓了什么”的读者。我会从 KV cache 的基本原理讲起拆解共享状态在哪些环节容易出问题给出可复现的排查步骤最后分享几个我在实际项目中踩过的坑和对应的隔离方案。2. KV cache 与注意力机制缓存到底缓了什么2.1 自回归解码为什么需要 KV cacheTransformer 的解码器是自回归的给定前缀[x1, x2, ..., xt]预测下一个 tokenxt1然后把xt1拼到前缀后面继续预测xt2如此循环。如果每一步都重新计算整个前缀的注意力计算量会随着序列长度平方增长。具体来说第t步需要计算t个 token 的 Q、K、V第t1步又要重新计算t1个大量重复计算浪费在已经算过的前缀上。KV cache 的思路很直接既然每一步只有最后一个 token 是新的那前面t-1个 token 的 K 和 V 就没必要重算。把它们缓存下来每一步只需要计算新 token 的 Q、K、V然后把新的 K、V 追加到缓存里用新 Q 和完整的 K、V 做注意力计算。这样每一步的计算量从O(t^2)降到O(t)显存换时间非常划算。用生活化的类比你在做一道需要反复查阅前文的阅读理解题。没有 KV cache 时你每回答一个新问题都要把整篇文章重新读一遍有了 KV cache你只读新段落前面读过的内容记在草稿纸上直接查草稿纸就行。草稿纸就是 KV cache。2.2 多头注意力中的 K、V 形状与内存布局要理解共享状态的问题必须先搞清楚 K、V 在多头注意力里的形状。假设 batch size 为B序列长度为S隐藏维度为D注意力头数为H每个头的维度为d D / H。输入张量形状是[B, S, D]经过线性投影后Q、K、V 的形状都是[B, S, D]。然后 reshape 成[B, S, H, d]再 transpose 成[B, H, S, d]这样每个头独立计算注意力。KV cache 通常存储的是投影并 reshape 之后的 K 和 V形状为[B, H, S_cache, d]。在解码第t步时新 token 的 K、V 形状是[B, H, 1, d]追加到缓存后面变成[B, H, t, d]。注意力分数计算为Q K^T形状[B, H, 1, t]再经过 softmax 和 V得到[B, H, 1, d]最后合并所有头得到[B, 1, D]。这里的关键点是每个头有自己独立的 K、V 缓存。如果实现时为了省内存让所有头共享同一块 K、V 缓存或者让不同层共享就会出问题。因为不同头的投影矩阵不同算出来的 K、V 语义完全不同共享等于把不同头的特征混在一起注意力分数就乱了。2.3 共享状态在哪些环节容易出问题根据我的经验共享状态引发的隔离问题主要集中在四个环节第一多头之间的缓存共享。有些实现为了减少显存占用把[B, H, S, d]的缓存压缩成[B, S, D]然后在计算时再 split。如果 split 操作返回的是 view 而不是 copy后续的原地更新会污染原始缓存。第二多层之间的缓存共享。Transformer 有L层每层都有自己的 KV cache。如果用一个全局列表存储所有层的缓存但索引写错了比如第l层读到了第l1层的缓存输出就会完全错乱。第三多请求之间的缓存共享。在 batch 推理时不同请求的序列长度不同需要 padding 对齐。如果 padding 部分的缓存没有正确 mask或者不同请求共享了同一块缓存区域就会互相干扰。第四原地更新与视图共享。PyTorch 的expand、view、transpose等操作返回的是视图不是新张量。如果对视图做原地操作会影响到原始张量。KV cache 的追加操作如果用cat还好如果用切片赋值cache[:, :, t, :] new_kv而cache又是某个更大张量的视图就会出问题。3. 口令实验的完整复现与排查过程3.1 实验环境与最小复现代码为了让大家能亲手复现这个问题我写了一个最小化的例子。环境是 Python 3.9 PyTorch 2.0CPU 和 GPU 都能跑。代码只保留最核心的多头注意力加 KV cache 逻辑去掉了 LayerNorm、FFN 等无关部分。import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttentionWithCache(nn.Module): def __init__(self, d_model64, n_heads4): super().__init__() self.d_model d_model self.n_heads n_heads self.d_head d_model // n_heads self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) # 故意用一个共享的缓存张量来模拟错误实现 self.shared_cache None def forward(self, x, use_cacheFalse): B, S, D x.shape q self.q_proj(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) k self.k_proj(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) v self.v_proj(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) if use_cache: if self.shared_cache is None: self.shared_cache (k, v) else: cached_k, cached_v self.shared_cache # 错误直接对缓存做原地拼接且所有头共享 k torch.cat([cached_k, k], dim2) v torch.cat([cached_v, v], dim2) self.shared_cache (k, v) attn torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5) attn F.softmax(attn, dim-1) out torch.matmul(attn, v) out out.transpose(1, 2).contiguous().view(B, S, D) return self.out_proj(out)这段代码里shared_cache是一个实例变量所有头共用同一个缓存元组。第一次调用use_cacheTrue时缓存被初始化为当前输入的 K、V。第二次调用时新的 K、V 被拼接到缓存后面。问题在于第二次调用的输入如果和第一次不同缓存里混入了两次输入的 K、V注意力计算就会把两次输入关联起来。3.2 两次推理结果不一致的定位方法复现步骤很简单构造两个不同的输入x1和x2先跑x1开启缓存再跑x2开启缓存观察x2的输出。然后重置模型只跑x2开启缓存对比两次x2的输出。如果不一样说明缓存被污染了。torch.manual_seed(42) model MultiHeadAttentionWithCache() x1 torch.randn(1, 3, 64) x2 torch.randn(1, 3, 64) # 第一次先跑 x1再跑 x2 out1_x2 model(x2, use_cacheTrue) # 重置模型 model_reset MultiHeadAttentionWithCache() model_reset.load_state_dict(model.state_dict()) model_reset.shared_cache None out2_x2 model_reset(x2, use_cacheTrue) print(torch.allclose(out1_x2, out2_x2)) # 输出 False定位方法就是二分法先确认问题是否由缓存引起关掉缓存看是否一致再确认是哪个维度的共享头、层、请求最后确认是视图还是拷贝的问题。我当时的排查路径是先打印缓存的data_ptr()发现两次推理的缓存地址相同说明是同一块内存再打印缓存的形状发现第二次比第一次长说明发生了拼接最后检查拼接逻辑发现没有按请求隔离。3.3 从缓存生命周期看隔离的必要性缓存的正确生命周期应该和请求绑定一个请求开始解码时创建缓存解码结束时释放缓存。不同请求的缓存必须物理隔离不能共享同一块内存。同一个请求内部不同层、不同头的缓存也必须逻辑隔离虽然可以放在一个大张量里但索引必须严格对应。我后来改写的正确实现是这样的class CorrectMHAWithCache(nn.Module): def __init__(self, d_model64, n_heads4): super().__init__() self.d_model d_model self.n_heads n_heads self.d_head d_model // n_heads self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, x, cacheNone): B, S, D x.shape q self.q_proj(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) k self.k_proj(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) v self.v_proj(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2) if cache is not None: cached_k, cached_v cache k torch.cat([cached_k, k], dim2) v torch.cat([cached_v, v], dim2) new_cache (k, v) attn torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5) attn F.softmax(attn, dim-1) out torch.matmul(attn, v) out out.transpose(1, 2).contiguous().view(B, S, D) return self.out_proj(out), new_cache关键改动是缓存不再作为实例变量而是作为输入参数传入、作为返回值传出。这样每个请求自己管理自己的缓存天然隔离。torch.cat返回的是新张量不会污染旧缓存。如果你追求极致性能可以用预分配的大张量加切片赋值但一定要确保切片范围正确且不要对视图做跨请求的原地操作。4. 共享状态引发的典型问题与排查技巧4.1 问题速查表从现象到根因下面这张表是我在实际项目中整理出来的覆盖了 KV cache 共享状态最常见的几类问题。你可以把它当成排查清单遇到“同样输入输出不一致”时逐条对照。现象可能根因排查方法修复方案两次推理输出不一致缓存未重置或跨请求共享打印缓存data_ptr和形状缓存与请求绑定用后即弃生成结果重复或乱码多头缓存共享特征混淆检查 K、V 的 head 维度是否独立每个头独立缓存不压缩长序列生成质量下降缓存索引错位读到 padding检查 attention mask 是否覆盖缓存mask 要包含缓存部分显存占用异常增长缓存未释放或视图持有引用用torch.cuda.memory_summary查看显式置None避免循环引用多 batch 推理互相干扰batch 内不同请求共享缓存检查 batch 维度的缓存是否独立每个请求独立缓存padding 对齐这张表里的每一条我都至少踩过一次。最坑的是“长序列生成质量下降”当时模型在短序列上表现正常一到 512 token 以上就开始胡言乱语。排查了很久才发现是 attention mask 只覆盖了当前步的 token没有覆盖缓存里的历史 token导致模型“看到”了 padding 位置。4.2 隔离方案对比物理隔离 vs 逻辑隔离隔离共享状态有两种思路物理隔离和逻辑隔离。物理隔离是每个请求、每个头、每层都分配独立的内存块互不干扰。逻辑隔离是把所有缓存放在一个大张量里但通过索引和 mask 保证访问范围正确。物理隔离的优点是简单、安全不容易出 bug缺点是显存碎片化分配和释放开销大。逻辑隔离的优点是显存利用率高适合大规模部署缺点是实现复杂索引和 mask 容易写错。我的建议是原型阶段用物理隔离生产环境用逻辑隔离。原型阶段最重要的是正确性显存多点无所谓。生产环境再优化但优化后一定要做充分的回归测试特别是边界情况空输入、单 token、超长序列、batch 内长度差异大等。4.3 实操心得三个避免缓存污染的习惯第一个习惯缓存永远不作为模块的实例变量。实例变量意味着生命周期和模块绑定而模块通常是全局单例多个请求会共享。把缓存作为 forward 的输入和输出让调用方管理生命周期是最安全的做法。第二个习惯追加缓存时用torch.cat不用切片赋值。torch.cat返回新张量天然避免视图共享问题。切片赋值虽然省内存但要求预分配的张量足够大且切片范围不能重叠。如果一定要用切片赋值确保缓存张量是contiguous的且不要对切片结果做原地操作。第三个习惯每次推理前打印缓存的形状和data_ptr。这个习惯帮我省了无数排查时间。形状能告诉你缓存是否被意外拼接data_ptr能告诉你是否共享了内存。在 CI 里加一个断言检查不同请求的缓存data_ptr不同能提前拦截大部分隔离问题。5. 从注意力机制看共享与隔离的设计哲学5.1 注意力本质是一种加权查询把注意力机制拆开看它做的事情就是用 Query 去查询 Key得到权重再用权重去加权 Value。这个过程天然要求 Query 和 Key 是“配对”的——哪个 Query 对应哪些 Key必须明确。如果多个 Query 共享了同一组 Key或者多个请求的 Key 混在一起查询结果就没有意义了。这就像你在图书馆查书每本书有书名Key和内容Value你的检索词Query只能匹配到你想要的那本书。如果所有书的书名都混在一起你检索“深度学习”可能返回“红烧肉做法”因为书名被共享了。KV cache 的隔离本质上就是保证每个 Query 只能看到自己应该看到的 Key。5.2 共享状态在视觉 Transformer 中的特殊表现视觉 Transformer如 ViT、Swin Transformer和 NLP Transformer 在缓存共享上有不同的表现。NLP 是自回归生成有明确的 KV cache 需求视觉任务通常是整图输入、一次性前向不需要缓存。但视觉领域有另一种共享状态问题窗口注意力中的窗口间共享。Swin Transformer 把特征图划分成多个窗口每个窗口内部做自注意力。如果实现时窗口之间的相对位置编码共享了同一组参数或者窗口边界的 padding 处理不当就会导致窗口间信息泄漏。这和 KV cache 的共享问题本质相同都是“本应隔离的状态被共享了”。视觉领域还有一个概念叫“高度方向、宽度方向注意力”专门针对特征图的垂直和水平两个维度分别做注意力。这种分解注意力的实现里如果 H 方向和 W 方向的注意力共享了同一组投影权重或者缓存了对方的中间结果也会出问题。我在做一个目标检测项目时就遇到过因为 H、W 方向注意力共享了 LayerNorm 参数导致检测框在垂直方向偏移的情况。5.3 因果自注意力中的 mask 与缓存隔离因果自注意力Causal Self-Attention要求每个位置只能看到自己和之前的位置不能看到未来。在训练时这通过一个上三角为负无穷的 mask 实现。在推理时KV cache 里存的是历史 token 的 K、V当前 token 的 Q 只能和缓存里的 K 做注意力不能和未来 token 做。如果 mask 写错了比如缓存部分的 mask 全是 0允许注意力当前 token 就会“看到”缓存里所有位置包括那些在因果意义上不应该看到的位置。更隐蔽的是如果多个请求的缓存拼在同一个 batch 里mask 必须区分不同请求的缓存边界否则请求 A 的 token 会看到请求 B 的缓存。我处理这个问题的经验是mask 的形状要和注意力分数的形状严格一致。注意力分数是[B, H, S_q, S_k]mask 也必须是[B, 1, S_q, S_k]或能广播到这个形状。S_k包含缓存长度和当前长度mask 要精确覆盖每个位置。不要图省事用torch.triu生成一个固定大小的 mask因为缓存长度是动态变化的。6. 工程落地如何设计一个隔离良好的推理缓存6.1 缓存数据结构的设计取舍设计缓存数据结构时核心取舍是用一个统一的大张量还是每个请求独立的小张量。统一大张量的好处是显存连续、便于批量操作坏处是不同请求长度不同需要 padding且索引复杂。独立小张量的好处是隔离彻底、逻辑简单坏处是显存碎片化batch 操作需要手动拼接。我目前的方案是折中按请求分组组内用统一张量组间物理隔离。具体来说维护一个cache_pool每个请求分配一个cache_id缓存张量按cache_id索引。同一个请求内部所有层的缓存放在一个列表里每层的 K、V 形状为[1, H, S, d]batch 维为 1因为每个请求独立。推理时把同一 batch 的请求缓存临时拼成[B, H, S_max, d]用 mask 处理长度差异算完再拆回去。这种设计的好处是请求之间物理隔离不会互相污染请求内部批量计算效率不低。缺点是拼接和拆分有开销但相比注意力计算本身这个开销可以接受。6.2 缓存分配与释放的时机控制缓存的分配时机应该是请求进入解码阶段时而不是请求刚到达时。因为编码阶段如果有 encoder不需要 KV cache过早分配浪费显存。释放时机应该是请求生成结束或超时而不是等整个 batch 结束。因为不同请求的生成长度不同早结束的请求应该尽早释放缓存给后续请求腾空间。我在实际项目里用了一个简单的引用计数每个缓存张量有一个ref_count请求持有缓存时加一释放时减一减到零时真正释放显存。这样可以避免缓存被提前释放导致悬空引用也可以避免缓存泄漏。注意PyTorch 的 CUDA 缓存分配器有自己的缓存池del tensor不一定立即归还显存给系统但会归还给 PyTorch 的缓存池后续分配可以复用。所以监控显存时要看torch.cuda.memory_allocated而不是nvidia-smi。6.3 多请求并发下的隔离验证方法验证隔离是否彻底最直接的方法是并发压力测试同时发起多个请求每个请求的输入不同记录每个请求的输出。然后单独跑每个请求对比输出是否一致。如果一致说明隔离良好如果不一致说明有共享状态污染。更严格的测试是交错测试请求 A 跑一步请求 B 跑一步再请求 A 跑一步模拟真实并发场景。这种测试能发现那些“只在特定交错顺序下才出现”的 bug。我写过一个测试脚本用随机种子控制交错顺序跑一千次每次检查输出一致性。这个脚本帮我抓到了至少三个隐蔽的缓存共享 bug。还有一个技巧是在缓存张量上打标记给每个请求的缓存张量设置一个唯一的request_id属性在每次注意力计算前断言cache.request_id current_request_id。这个断言在开发阶段开启生产环境关闭能提前拦截跨请求污染。7. 一些踩坑记录与个人体会7.1 那些年我们误用的 view 和 expandPyTorch 的view和expand是缓存共享问题的重灾区。view返回的是共享内存的视图修改视图会影响原始张量。expand返回的是广播后的视图虽然不能直接原地修改但如果你对它做contiguous()再修改有时会意外影响原始张量。我踩过的一个坑是用expand把[1, H, S, d]的缓存扩展到[B, H, S, d]以为每个 batch 有独立副本实际上所有 batch 共享同一块内存。后来改成repeat才解决。repeat会复制数据expand不会。在缓存场景下如果你需要每个请求独立修改缓存必须用repeat或clone。另一个坑是transpose之后忘记contiguous。transpose返回的是视图内存不连续。后续如果对这个视图做view操作PyTorch 会报错或隐式拷贝。隐式拷贝有时会破坏缓存的共享关系导致你以为在修改缓存实际上修改的是拷贝。7.2 缓存长度与位置编码的配合位置编码和缓存长度必须严格对应。如果缓存里存了S个 token 的 K、V当前 token 的位置编码应该是S从 0 开始计数而不是 0。如果位置编码写错模型会以为当前 token 在序列开头注意力分布完全错乱。我在一个项目里用过旋转位置编码RoPERoPE 的特点是位置信息直接乘在 Q、K 上不额外加位置嵌入。使用缓存时缓存的 K 已经乘过旧位置的旋转因子新 token 的 K 要乘新位置的旋转因子。如果缓存里的 K 被重新计算或覆盖旋转因子就丢了。所以 RoPE 场景下缓存必须存乘过旋转因子的 K不能存原始 K。7.3 一个反直觉的发现隔离反而更快最后分享一个反直觉的发现在某些场景下物理隔离的缓存比共享缓存更快。原因是共享缓存需要复杂的索引和 mask 操作这些操作本身有开销而物理隔离的缓存逻辑简单虽然显存占用多但计算路径短CPU 和 GPU 之间的同步少。在 batch size 较小、序列较短的场景下隔离方案的总延迟反而更低。这个发现让我重新思考了“优化”的定义优化不是一味省显存而是在显存、计算、延迟之间找平衡。KV cache 的共享状态问题表面上是显存优化引发的 bug深层原因是我们在优化时忽略了正确性这个前提。先保证隔离再谈共享这是我用两次线上事故换来的教训。后续如果要做进一步扩展可以考虑把缓存管理抽象成一个独立的CacheManager类负责分配、索引、释放和隔离验证。这样注意力模块只关心计算缓存管理只关心生命周期职责清晰bug 也更容易定位。