某个自研推理项目的调试过程让我在一种极其隐蔽的 KV Cache 问题上耗了整整一天。事情背景并不复杂为了提升吞吐我把批处理推理里的缓存从每请求独立申请改成共享显存池复用顺带为适配某个优化 kernel 改了一下 attention mask 的传参方式。改完之后单条 prompt 怎么跑都正常但只要 batch 里出现两条长度差异较大的输入生成结果就开始出现大段重复、乱码甚至一头扎进某个无意义 token 的循环里出不来。这篇文章把这次排查的完整链路还原出来。它不只想讲一个 bug 的修复更想聊清楚一个我此前几乎没认真想过的问题KV Cache 在推理优化里承担的角色远不止减少重复计算这么简单。它实质上定义了大模型每一步能看到的上下文范围和边界是一份跨阶段的上下文契约。这份契约一旦乱了表面上看不到任何 crash但生成质量会以极其阴险的方式慢慢崩掉。这个盲区才是这次折腾最大的收获。1. 先搞清楚KV Cache 缓存的是什么优化为什么总在它身上翻车1.1 没有缓存的时代每次 decode 都在重复计算历史大模型生成是自回归的预测下一个 token 时输入是当前整个 token 序列。假设现在序列长度是 100你要生成第 101 个 token在没有任何缓存的情况下模型要把这 100 个 token 全部过一遍 Transformer。每个 token 都要计算自己的 Q、K、V再彼此做注意力。生成第 102 个 token 时前面 101 个 token 的 QKV 又要全部重新算一遍。把复杂度拉出来更直观。如果 prefill 阶段处理 N 个 token之后还要生成 M 个 token无缓存时每一步都要对当前全量序列做注意力计算总计算量大致是 O(N² (N1)² … (NM)²)差不多是 O(M·N² M²·N M³/3) 的量级。有 KV Cache 之后prefill 只算一次 O(N²) 的注意力后续每一步只计算新 token 的 Q、K、V然后跟历史 K/V 做一次长度不断增长的注意力。省掉的不是常数倍的冗余而是重复计算整段历史这个本质性浪费。这也是为什么任何大模型推理服务都不会关掉 KV Cache 的原因。它是吞吐和时延优化的地基几乎所有你看得到的推理加速手段——显存池、block 分块管理、连续批处理、prefix 复用——都在围绕它做文章。1.2 缓存不只是存张量它定义了每一步能看到什么KV Cache 的基本形态并不复杂把每一层的 K 张量、V 张量按序列长度维度拼起来单看一条样本大致是层数、注意力头数、当前序列长度、每个头维度这么两组数据。难的地方在于它的生命周期横跨两个行为完全不同的阶段。prefill 阶段是一次性喂入一整段 prompt计算量大但本质上是批量写入把整段历史 K/V 一次性算好塞进缓存。decode 阶段每次只输入一个新 token计算量小但必须先把新 token 的 K/V 追加到缓存末尾再去读一遍完整的历史 K/V 做注意力。一个是写入整段历史一个是追加一行并读取全部历史两者对缓存的操作模式完全不一样。各种推理优化恰恰喜欢在这两个阶段的衔接处引入新约定比如 mask 的约定、位置编码的约定、有效长度的约定。任何一个约定没对齐缓存里的 K/V 张量还是完整地躺在那里但它们的语义已经悄悄变质了。我这次踩的坑就是约定在 prefill 阶段出了问题。1.3 我在某次优化里踩进去的入口我踩坑的入口是给批量推理做显存池复用。为了减少反复申请和释放缓存带来的开销我把 cache 改成按 block 分配、池化复用的方式。同时为了适配一个优化 kernel我在 prefill 阶段决定不传完整的二维注意力掩码只传 padding 掩码指望那个 kernel 在内部自动补上因果部分。这个决定听起来很合理很多优化 kernel 的文档里也确实写着内部自动处理因果掩码。问题就出在我接的这个 kernel 恰好不会。它只是把我传入的 mask 当作完全体来使用。于是批量路径上的 prefill 阶段因果掩码整体丢失。当时我完全没意识到这一步埋了多深的坑。直到批量推理生成质量开始飘我才回头找它的麻烦。回头看如果一开始就把 mask 的构造逻辑收敛成一个公共函数让所有路径都共用这一天的时间完全可以省下来。2. 问题现场单条生成正常batch 一加 padding 就输出乱码2.1 触发条件与最诡异的特征最让人难受的是症状不稳定。同样一个 prompt单条跑出来质量很好语义通顺两条长度差不多的输入放进同一个 batch问题也不算明显一旦把一条较长文本和一条很短文本放在一起长文本后半段就开始崩。具体崩溃模式主要有三种大段重复同一句话突然蹦出几个无关 token 又绕回原句极端情况下在一个无意义 token 上反复打转。我一开始怀疑是不是显存池复用导致数据被覆盖于是关掉 pool 复用问题照旧。又怀疑是随机种子和采样排序的问题把温度调到 0 改成确定性生成乱码依然存在。这种不崩溃、不报错、只是结果变差的行为是所有排错场景里最磨人的类型。运行场景现象倾向性猜测单条 prompt生成完全正常模型本身没坏batch 内长度接近偶发轻微重复采样噪声batch 内长度悬殊 padding长文本后半段明显乱码疑似 mask / padding 边界问题2.2 第一步排查变量二分锁定 attention排错思路是先做二分。第一步把优化开关全部关掉恢复正常 attention 路径问题消失。第二步只开显存池问题依然存在。第三步只开自定义 mask 传参路径问题稳定复现。这样几轮下来嫌疑范围直接缩小到 attention 计算和 mask 处理这一段。这一步没有太多算法含量核心就是控制变量。但我强烈建议所有做推理优化的人都养成一个习惯任何优化上线前准备一个开关矩阵让每个优化点都能独立开启和关闭。否则等到问题出现你会在显存池、缓存复用、mask 压缩、kernel 替换、调度器逻辑里毫无头绪地打转。这次我运气好保留了完善的开关配置否则排查时间至少翻倍。2.3 最小复现两个 prompt一大一小锁定 attention 之后我把复现脚本压到最小两个 prompt一个短、一个长短的那个后面补 padding token 拼到相同长度然后用同一个模型、同一份日志在固定位置打印输出序列。这个最小复现非常稳定几乎 100% 能在长文本后半段复现质量崩塌。我还专门试过把短 prompt 放在 batch 的第一个位置和第二个位置结果发现无论怎么排出问题的主要都是长文本的后半段。这说明问题跟样本在 batch 里的排位无强相关而跟长度差、padding 区域强相关。这个观察给了后面的排查一个非常明确的指向去查 mask 在 padding 边界上到底发生了什么。3. 把 Cache 张量拉出来做尸检锁定 Attention Mask3.1 数值对比发现的污染痕迹猜测再多不如直接看数据。我把正常路径和异常路径在同一个权重、同一批输入上各跑一遍然后把某一层的 cache 张量按序列长度维度分成几段有效 token 区域、padding 区域、靠近末尾的最后几个 token。对比出来的结论非常明显异常路径的 padding 区域里K/V 不再是初始值而是出现了一堆非零数字有效 token 的尾部区域K/V 数值分布也跟正常路径差了一大截。这说明两件事。padding 位置被写入了本不该存在于注意力计算中的信息有效 token 的尾部也被多余的上下文污染了。正常 attention 里位置 i 只能看到 0 到 i-1 的位置如果能看到未来位置尾部 token 的 K/V 就会吸收未来信息。padding 位置同理它们本不该被任何 query 看到实际却被当成普通 token 参与了注意力。我还专门抽了一个 padding 位置把它对应的 attention 分数打印出来。正常路径里这一行应该是被屏蔽的无效值异常路径里分数是有限的甚至分布相当均匀。这一步基本实锤attention mask 没有正确生效。3.2 顺着 mask 函数追到根因再往上追代码很快就找到我写的那段 mask 构造函数。当时为了适配优化 kernel 的调用约定我把 mask 压缩成只带 padding 维度的一条写出来的逻辑和下面这段差不多# 批量优化路径只传 padding mask def build_mask_fast(input_ids, pad_token_id): padding (input_ids ! pad_token_id) # [batch, seq_len] mask padding[:, None, None, :] # [batch, 1, 1, seq_len] return mask这个 mask 能挡住 padding 位置但完全没有约束因果部分。如果底层 kernel 内部自动补上因果 mask那确实没问题但我实际接的 kernel 没有补。最终效果就是 prefill 阶段所有有效 token 互相可见包括自己之后的 token。正确版本其实只多了一步合并因果 maskdef build_mask_correct(input_ids, pad_token_id, seq_len): padding (input_ids ! pad_token_id) # [batch, seq_len] causal build_causal_mask(seq_len) # [1, 1, seq_len, seq_len] return padding[:, None, None, :] causal # [batch, 1, seq_len, seq_len]现在看这就是少合并了一个 mask的低级错误。但为什么 debug 了那么久因为它藏在路径分支里单条 prompt 走框架默认 attention内部自带因果处理批量优化路径走自定义 kernel我误以为它也会自带。两条路径的 mask 语义从一开始就不统一所以单条正常、批量翻车。3.3 为什么这类 Bug 不崩不报错只会让质量阴险地变差大多数 mask 错误如果发生在 decode 阶段会立刻导致 shape 不匹配或者产生 NaN表现非常显眼。但这次发生在 prefill 阶段而且不是形状错误是语义错误。模型不会崩溃loss 也不会当场爆炸只是把未来 token 的信息混进了历史 K/V 里。decode 阶段的新 token 做注意力时它 attend 到的是已经被污染过的历史 K/V。因此无论 decode 阶段的 mask 多么正确新 token 读到的上下文里都带着未来的影子。这种污染会随序列长度累积越往后被污染的 token 越多生成轨迹就偏得越厉害。于是就有了长文本后半段崩溃、短文本相对不明显的观感。这也解释了为什么我一开始会怀疑显存池和 beam search。这些地方出错通常会导致内存异常或明显错位而 mask 语义错误是一种没有异常信号的逻辑污染必须走到数值层面才能看见。排查这类问题最有效的工具就是把张量分段拉出来跟正常路径做数值对比。4. 修复与验证让 prefill 和 decode 共用同一份掩码语义4.1 修复方案统一 mask 构造逻辑修复手段本身并不复杂。我把 mask 构造收敛到一个公共函数prefill 和 decode 都从同一个函数拿 mask。核心原则是不信任某个 kernel 会自动补因果 mask这种说法除非亲眼在文档和源码里确认过它的调用约定否则一律把 mask 当作必须显式传全的量。改完后的代码结构大概是这样一个公共函数负责生成完整的注意力掩码包括 padding 部分和 causal 部分prefill 用完整掩码decode 阶段根据需维护的 active 长度生成对应的一次性掩码。单条路径、批量路径、不同 kernel 分支之间语义完全一致。更本质的原则是不要让两条代码路径对同一个序列该看到哪些上下文产生不同理解。如果单条路径和批量路径的 mask 语义天然不同那几乎等于埋了一个定时炸弹。后续再往里加量化、切块 prefill、prefix 复用这些优化都要保证这条原则不被破坏。4.2 给 cache 写入加一道自检闸门修 mask 是治本但我也顺手加了一道自检闸门防止以后再出现同类问题同时不影响极致的性能追求。具体做法是在 prefill 阶段算完 attention、准备把 K/V 写入缓存之前检查几条不变量缓存序列长度必须等于当前样本的实际有效长度有效区域不允许出现 NaN 或 Inf抽一两个 token 验证其 attention 分数在 padding 区域和未来区域都被正确屏蔽。这些检查在开发阶段非常值得保留。它们本质上就是在校验上下文契约的关键字段每次写入缓存前确认一遍能挡掉一整类悄悄污染缓存的 bug。等系统稳定后可以把断言编进 debug 模式release 模式跳过性能不损失。我还把最小回归用例保留了下来一长一短两个 prompt固定的 padding 配置每次改完代码先跑一遍检查后半段输出是否稳定、是否出现循环。这条用例到现在仍在使用后来帮我挡住了至少两三个潜在的回归。4.3 修复后的效果与性能回归修复之后之前 100% 复现的长文本后半段乱码直接消失。同一个 batch 配置下长文本输出恢复到单条时的质量水平我换了好几组采样种子都没有再出现循环重复或乱码。性能上attention 的计算量没有本质变化生成吞吐基本持平。但显存表现更好了padding 位置不再被错误地当作有效 token 参与注意力部分优化路径可以提前识别 padding 区域并复用这部分显存实测稳定运行时的峰值显存下降了 10% 左右。对比项修复前修复后长文本后半段输出重复、乱码、循环恢复稳定峰值显存偏高下降约 10%生成速度正常基本持平掩码语义两条路径不一致全局统一这个结果也印证了一个观点语义修正是性能正确性的地基。看起来只是补了一个 mask实际上避免了无效信息参与注意力计算反而让显存和后续的显存管理都变得更干净。5. 这个 Bug 暴露的优化盲区KV Cache 是一份跨阶段上下文契约5.1 核心洞察性能优化的前提是序列语义一致这次倒腾的最大收获不是别忘了加因果 mask这个具体知识点而是对 KV Cache 有了一个更整体的认知缓存里存的 K/V 张量只是载体真正被模型依赖的是这套张量隐含的上下文语义——哪些位置是可见的哪些位置是 padding当前位置在哪整个序列实际有多长。我把这套语义称作上下文契约。prefill 建立契约decode 遵守契约任何优化中间件都必须维护契约。显存池化、block 分块、量化压缩、prefix 缓存复用本质上都只是在换一种方式存放同一份契约。如果存放方式把契约搞乱了优化速度再快输出的也是垃圾。另一个相关的盲区是性能模型。很多人聊大模型推理优化时眼睛总盯着 FLOPs觉得减少计算量就是一切。但 decode 阶段其实是 memory-bound 的一个 batch 在 decode 时的主要瓶颈往往是 KV Cache 的读写带宽和显存容量而不是矩阵乘的多余 FLOPs。省 20% 计算有时候不如把 cache 的访存布局整理好、把无效 padding 的读写省掉。5.2 顺着契约视角看还有哪几个高频坑用上下文契约这个视角重新审视我发现很多经典问题都能统一起来。第一个是位置索引没有前进。使用 RoPE 这类位置编码时decode 阶段每前进一步position id 也应当加一。如果不加新 token 的 K/V 会被放到错误的位置上模型看到的位置信息错乱表现跟这次 bug 类似——一开始还行越往后越偏。排查方向就是检查 position id 是否随 active 长度递增。第二个是 beam search 回退 cache 时只回退了一部分状态。beam search 要维护多个候选序列剪枝时会把某个分支的 cache 回退到某个长度。如果只回退 K/V 张量却不同步回退 position id、active 长度、对应的 mask 状态下一步生成时缓存里就会出现看似完整、实际错位的数据进而污染所有候选分支。第三个是 prefix 缓存复用时的 mask 边界。复用一段公共前缀是省 prefill 时间的重要手段但复用之后第一个新 token 的因果边界在哪里、padding 怎么处理必须严格重新计算。边界差一位新 token 就会看到本不该看到的信息输出同样会越跑越偏。类似坑典型症状排查方向RoPE 位置索引未前进越往后越偏、位置错乱检查 position id 是否随 active_len 递增beam search 回退不完整候选序列状态奇异检查所有配套状态是否同步回退prefix 缓存复用边界错位复用后第一个新 token 异常重新计算 mask 边界与 active_len5.3 把契约落进工程实践两个长期有效的好习惯最后分享两个我认为长期有价值的工程习惯。第一个是把序列的语义状态显式建模而不是散落成一堆 K/V 张量。比如用一个结构体统一保存 cache、position id、active length、mask 元信息所有操作——追加、回退、扩充 batch、复用前缀——都通过这个结构体提供的方法来完成。这样做能大幅降低张量看起来正常但语义已经错乱的概率等于把契约变成了代码里看得见、摸得着的对象。第二个是保持一条廉价的语义回归测试。固定两个样本一个较长 prompt一个较短 prompt 加 padding组成最小 batch跑一次完整生成检查后半段输出是否稳定、是否出现循环。这条测试成本极低但几乎所有 cache 语义类问题都会在这条测试上露出马脚。这次 bug 之后我把这两个习惯带进了后续每一个优化项目。说实话它们的价值不亚于任何显存复用技巧——因为再快的方法如果跑在错误的上下文上也只是更快地生产垃圾。