用Lean4形式化验证滑动窗口KV Cache的正确性

做推理引擎的同学肯定遇到过这种纠结：显存快被打满，Batch size提不上去，这时候看到滑动窗口缓存能把KV Cache的上限锁死在窗口大小上，内存立刻可控，很难不心动。但真把它放进生产环境之前，你一定会问自己一个问题：把窗口外的K/V悄悄丢掉，输出到底还对不对？

我这次选了一个比较硬核的回答方式——用Lean4把「滑动窗口下KV Cache的正确性」做成了形式化验证。不是再写一轮单元测试，也不是拿随机样例轰边界，而是把缓存状态、位置对应关系、注意力语义全部建模成可推理的数学对象，最后得到一个机器检查过的定理：无论生成多少token，滑动窗口缓存每一步选择的K/V集合，都与全量历史做窗口截取后的结果完全一致。

这篇文章是这次验证过程的完整复盘。适合两类人看：一类是在LLM推理侧做KV Cache优化、想搞清楚形式化验证能不能落地到工程问题的人；另一类是正在学Lean4、想知道一个真实问题在Lean4里长什么样的人。我会把建模思路、定理陈述、证明策略和踩过的坑都讲一遍。

1. 为什么要为一个看似简单的缓存策略做形式化验证

1.1 KV Cache与滑动窗口：省显存但引入状态

先对齐一下背景。Transformer自回归解码时，每个token的注意力计算都需要之前的Key和Value，直接重新算一遍代价太高，于是大家都把它们缓存起来，这就是KV Cache。问题在于KV Cache的大小随序列长度线性增长，公式大概是这样：

2 × num_layers × num_heads × head_dim × seq_len × batch_size × dtype_size

当seq_len到几万、batch再一放大，显存消耗非常可观。滑动窗口注意力就是为了对抗这个增长：每个query只和最近的W个token做注意力，于是缓存只需要保留最近W个token的K/V，内存占用变为常数。

这里有个关键认知需要先说清楚：如果模型本身训练时就使用了滑动窗口注意力（比如Mistral、Longformer这类），那么滑动窗口缓存不是近似，而是这种注意力模式下的精确实现。窗口外的K/V在标准注意力中会被掩码掉，softmax权重为0，因此丢弃它们不会改变任何数学结果。我们验证的正是这个精确性，而不是去证明「全局注意力近似成滑窗的误差界」——后者是另一个完全不同的问题，难度也大得多。

1.2 缓存Bug的隐蔽性：不崩溃，但悄悄变差

这类缓存逻辑的Bug和普通程序Bug有个非常大的区别：它通常不会让程序崩溃，也不会产生明显的数值NaN或Inf，它只是让生成结果在长文本场景下悄悄变差。

我实际见到过的错误类型包括：索引偏移一个位置导致注意力窗口整体错位；窗口未满时基准位置算错；滚动位置编码的绝对位置记录错误；环形缓冲区在覆盖旧数据时head指针和tail指针的更新顺序反转。这些错误在生成短文本时完全不会暴露——因为窗口还没满，代码路径根本不会走到逐出分支；在长文本生成中，错误可能表现为某层某头的一两个位置注意力分布异常，但整体困惑度只会轻微劣化。

最麻烦的是这类问题难以回归验证。你修完一个偏移Bug，跑一遍长文本生成，看几个案例感觉「好像差不多」，很难量化之前到底对不对。我见过团队在线上模型里长期带着一个滑动窗口的索引Bug，输出质量一直比理论预期差一点，但没人定位到缓存层。

1.3 为什么测试不能代替这里的验证

有人会问：写单元测试不就行了吗？比如生成几百条随机序列，对比滑动窗口缓存实现和朴素全量实现的输出。这确实能抓到一部分问题，但覆盖范围有限。

首要问题是状态空间太大。窗口位置随着生成步数不断右移，未满、刚好满、已满后持续替换，这三个阶段发生在每个不同长度上。随机测试撞上「满窗后的第一次替换」这个精确边界的概率并不高；而这类跨阶段的边界恰恰是最多Bug的地方。

另一个问题是LLM输出对细微错误的鲁棒性。由于softmax的分布特性，一个位置的logit稍微出错，采样结果可能不变，也可能变，但评估指标对单步微小偏差不敏感。单元测试适合抓「结果完全错误」的问题，而缓存类Bug经常是「结果基本对、偶尔偏一点」，这说明测试预言本身就没定清楚。

形式化验证提供的是对全部合法输入序列的保证，不是抽样保证。它把「窗口边界处理正确」从概率事件变成定理。这也是我这次选择Lean4的根本原因。

1.4 为什么是Lean4而不是其他证明助手

做形式化验证，可选工具其实不少。我简单对比过Coq、Isabelle/HOL、Agda和Lean4，做了个表格：

选Lean4的核心原因有三个。第一，Mathlib对代数、列表、有限类型这些基础设施覆盖得非常全，很多关于List的引理可以直接用，不用自己从零开始。第二，Lean4既可以当证明助手也可以当函数式编程语言，模型里的函数能被真正执行，方便先写可运行的规范再补证明。第三，VS Code插件成熟，交互式证明的体验好，目标状态随时能看，这对调试证明过程帮助巨大。

2. 先把「正确性」说清楚：滑窗缓存的两种语义模型

2.1 参考实现：全量KV Cache加掩码

做形式化验证的第一步，是先把非形式的「输出应该正确」翻译成一个精确的数学命题。我采用的方法是定义两种实现模型，证明它们语义等价。

第一种模型是「参考实现」：假设缓存保存了从开始到当前的所有K/V，注意力计算时通过掩码把窗口外的位置排除掉。用公式表达就是：

Attn(q_t, K_1..t, V_1..t) = softmax( (q_t K_1..t^T / √d) + mask_window + mask_causal ) V_1..t

其中mask_window把位置小于t-W+1的logit设为负无穷，mask_causal保证只看到过去。

这个模型的正确性是「显然」的，因为它就是滑动窗口注意力的定义本身。但它也是不可扩展的——缓存无限增长。它存在的意义是作为规范（specification），给所有其他实现提供一个对照基准。

2.2 目标实现：定长滑动窗口缓存

第二种模型是「目标实现」：缓存固定为长度为W的列表，每一步生成新token时，把新的K/V追加到末尾；如果缓存长度超过W，就把最开头的KV逐出。这个模型对应实际工程里真正的缓存数据结构。

这个模型必须满足的核心性质非常简洁：在任意时刻，目标实现缓存中的K/V序列，恰好等于参考实现中全量历史的后W个K/V。如果这个性质成立，那么两个模型在任意query上的注意力输出必然一致——因为参考实现里窗口外的权重本来就是0。

2.3 正确性的正式定义：逐步模拟关系

形式化上，我定义了一个谓词Valid来表示缓存与规范之间的对应关系：

def lastWindow (w : Nat) (hs : List (KV Cell)) : List (KV Cell) := hs.drop (hs.length - min w hs.length) def Valid (w : Nat) (hs : List (KV Cell)) (c : List (KV Cell)) : Prop := c = lastWindow w hs

意思是：给定全量历史KV列表hs，合法缓存状态c必须等于hs去掉开头的hs.length - min w hs.length个元素之后剩下的部分。这个定义把「窗口语义」用一段精确代码钉死了。

这类模拟关系是形式化验证里最核心的手法：证明程序正确，其实就是在证明程序状态和规范状态之间存在某种逐步保持的对应关系。这样定义的自然之处在于，它完全刻画出「滑动窗口」的语义——所谓窗口，就是全量历史上从某个位置到末尾的子序列。

2.4 证明责任拆解：三个子目标，一个核心等式

把正确性命题正式展开，可以拆成三个子目标：

1. 长度不变式：缓存长度永远不超过W。这是最基础的，保证内存上界成立。
2. 内容对应：缓存列表始终等于全量历史的最后W个K/V。这是核心，也是证明工作量最大的部分。
3. 位置对应：如果K/V携带位置信息，缓存中第i个元素的位置必须等于当前绝对位置 - 窗口大小 + 1 + i。这个性质在RoPE场景下尤其重要。

三个子目标最后会汇总成一个核心等式。假设当前历史是xs，新生成token的KV是kv，那么需要证明：

lastWindow w (xs ++ [kv]) = step w (lastWindow w xs) kv

其中step就是目标实现的更新函数。这个等式说明了「规范先追加再截断」和「实现先截断再追加再处理溢出」两条路径殊途同归。整个验证的骨架，就是围绕这个等式展开的。

3. Lean4建模：用类型和不变式把缓存写进逻辑

3.1 定长缓存的类型表达：List长度约束与Vector

建模时第一个选择是用List还是Vector。Vector α n是长度n的定长列表，长度信息在类型层面就固定了，安全但笨重——一旦需要动态长度检查，类型层面就非常啰嗦。List α则长度在运行时，需要手动证明长度条件。

我最终采用List加长度约束的组合。这样step函数可以写得很直接：

structure KVCache (w : Nat) where items : List (KV Cell) length_le : items.length ≤ w

length_le这个字段本身就是不变式的声明：任何KVCache类型的合法值，其items长度必然不超过w。这是Lean4比普通语言强的地方——不变量可以被编码进类型，编译器强制所有构造路径满足它。

但这里要提醒一点：把不变量编码进structure后，每次构造和更新都要额外提供证明，会给函数定义增加噪音。如果目标是快速验证核心语义，可以先不把长度约束放进类型，只单独证明一个length_le定理。这是我的建议——先让函数好写，再补证明。

3.2 缓存不变式：槽位与token位置的对应关系

长度约束只是第一层，实际工程中更关键的对应关系是：缓存里第i个槽位到底存的是哪个token的K/V。如果模型使用绝对位置编码，这个对应关系直接决定了注意力的正确性。

在Lean4里，我把位置信息显式建模成KV Cell的一部分：

structure KV Cell where pos : Nat key : Vector Float d val : Vector Float d

然后定义位置不变式：对任意合法缓存，如果缓存的第一个元素是窗口内最早的token，那么第i个槽位的pos必须等于base + i，其中base是窗口内最早token的绝对位置。

def PositionInvariant (c : List (KV Cell)) : Prop := ∃ base : Nat, ∀ i : Nat, i < c.length → (c.get i).pos = base + i

这个不变式的价值在于：它把实现中隐式依赖的「位置对应关系」变成显式的、可检查的命题。实际代码里可能通过数组下标或偏移量间接表达这个关系，一旦写错，普通测试很难发现，但定理证明会在编译期抓住。

3.3 核心函数骨架：step、lookup、attention

目标实现的更新函数在Lean4里可以写成这样：

def step (w : Nat) (c : List (KV Cell)) (kv : KV Cell) : List (KV Cell) := let c' := c ++ [kv] if h : w < c'.length then c'.drop 1 else c'

这里用了依赖类型的分支：当c'.length > w时丢弃头部，否则原样保留。实际工程中的环形缓冲区实现会和这个列表模型不同，但语义等价——列表模型方便证明，环形缓冲区贴近硬件。

注意力查询函数可以抽象成参数，不必在这里展开浮点运算：

def attentionWithCache (attn : List (KV Cell) → Output) (c : List (KV Cell)) : Output := attn c

把attn作为参数传入，好处是验证注意力计算的实现时可以单独进行，缓存正确性验证不依赖具体数值运算。这是很关键的分层设计：缓存逻辑和数值计算解耦，证明工作量大幅下降。

3.4 从模型到真实实现的映射：环形缓冲区与GPU kernel

你可能已经发现，上面这个列表模型和真实C++/CUDA里常用的环形缓冲区还有距离。真实实现通常是一个定长数组加head/tail指针，写入时覆盖最旧元素。

这里需要补一层「物理实现到逻辑模型」的模拟关系。定义一个函数把环形缓冲区的数组和head指针解释成列表：

def ringToList (w : Nat) (buf : Array (KV Cell)) (head : Nat) : List (KV Cell) := (List.range w).map (fun i => buf.get! ((head + i) % w))

然后证明：环形缓冲区的push操作在ringToList解释下，等价于列表模型的step。这层证明需要处理取模运算，是工作量比较大的部分。如果只是验证逻辑正确性，可以先证明列表模型，物理层用人工review确认对应关系；但如果要做到完全可信，这层证明值得补上。

实际项目里我建议采用两层验证模式：第一层验证列表模型的语义正确性，第二层验证环形缓冲区对列表模型的模拟。这样任何一个环节出错，都能定位到具体层。

4. 定理陈述与证明策略：核心引理到归纳完成

4.1 单步正确性引理：逐出不改变注意力输出

最核心的单步引理可以这样陈述：如果当前缓存合法，那么执行一步生成后，新缓存依然合法。

lemma step_preserves_valid (w : Nat) (xs : List (KV Cell)) (kv : KV Cell) (h : Valid w xs c) : Valid w (xs ++ [kv]) (step w c kv)

证明时展开Valid和step的定义，分成xs.length < w和xs.length ≥ w两种情况。前者相当于窗口未满，直接由h推得；后者需要证明一个关于List.drop和append交互的等式：

(xs ++ [kv]).drop (xs.length + 1 - w) = xs.drop (xs.length - w) ++ [kv]

这个等式是单步证明的心脏。它说明：在全量历史上先追加新元素再截断窗口，等价于先截断旧窗口再追加新元素。直观上就是「后W个元素」这个操作在追加操作下的结合性质。

4.2 全序列正确性定理：从初始缓存归纳

单步引理搭好之后，全序列正确性定理就可以用归纳法收尾。定义整个解码过程：

def run (w : Nat) (tokens : List (KV Cell)) : List (KV Cell) := tokens.foldl (fun c kv => step w c kv) []

然后证明：

theorem run_correct (w : Nat) (tokens : List (KV Cell)) : Valid w tokens (run w tokens)

证明通过对tokens做结构归纳：空列表时平凡；cons时用归纳假设加上单步引理，两步就完成。最终再补一个推论：由于参考实现的注意力在窗口外的权重为0，所以滑动窗口缓存的注意力输出与全量缓存的注意力输出完全一致。

4.3 证明过程拆解：simp、omega、induction的分工

Lean4的证明自动化相当强，但也不能全指望simp一把梭。我的经验是两个自动化策略分工明确：simp负责等式化简和定义展开，omega负责线性整数算术。两者配合可以处理掉大量机械证明。

上面那个核心等式，在Lean4里可以这样拆：

lemma drop_append_of_ge (xs : List α) (a : α) (w : Nat) (h : w ≤ xs.length) : (xs ++ [a]).drop (xs.length + 1 - w) = xs.drop (xs.length - w) ++ [a]

这个引理本身通过对xs做归纳证明。归纳步里关键是omega处理Nat的减法关系，然后用simp做列表运算化简。整个过程没有特别高深的技巧，但需要对List.drop的行为非常熟悉。

还有一个典型陷阱：Nat的减法在Lean4里是截断的。3 - 5 = 0，这在普通编程里可能无所谓，但在证明中会破坏直觉。好消息是，在做drop (length - w)这类表达式时，截断减法的行为恰好符合需要——当length < w时，length - w = 0，drop 0返回整个列表，窗口未满时不发生逐出。但这也意味着你不能直接断言xs.length - w在length < w时是「负的」——类型系统里没有负数，你必须通过条件分支来处理。

4.4 证明工作量与噪音：哪些值得写，哪些只是机械劳动

实际验证下来，如果只做列表模型的核心证明，代码量并不夸张，大概两三百行Lean4就够。但这个数字有迷惑性——我花在调试证明脚本上的时间，远比写定义的时间多。

我把证明分成两类：一类是「有价值」的证明，比如单步引理和位置不变式，它们揭示了实现设计中的关键契约；另一类是「噪音」，比如反复证明List.drop的分配性质、Nat加减法的边界条件。后者本质上是标准库覆盖不足导致的重复劳动，不算有智力含量，但绕不开。

减少噪音的方法有两个。第一，把常用引理抽出来写成simp定理，让后续证明自动使用；第二，尽量避免展开大定义，尽量用rw [引理]而不是simp [大定义]精确控制证明步骤。后者的重要性容易被低估——无脑simp在大定义上会显著拖慢编译速度，甚至导致内存占用暴涨。

5. 实测中的教训：索引算术、未满状态与值得吐槽的坑

5.1 窗口未满：大多数Bug都出在优雅退化路径

我在验证过程中发现，真正容易出Bug的不是窗口满之后的替换逻辑，而是从未满到满的过渡阶段。

很多实现在窗口未满时，缓存数组只有前一部分被填了有效数据，head指针可能是0，也可能指向某个未定义位置。如果代码统一用(head + i) % w计算槽位，那么未满阶段读出来的可能是一堆垃圾KV；如果代码对满/未满分别处理，那么「满了之后第一次替换」时基准位置必须切换，这个切换点最容易出错。

Lean4的好处是，当你在证明step_preserves_valid时，必须分情况处理xs.length < w和xs.length ≥ w，这会强制你意识到过渡路径的存在。但如果你只是用C++写循环，这段过渡路径很容易被忽略，因为短文本测试全部走的是未满分支。

5.2 Nat截断减法：位置编码错乱的隐藏来源

Lean4里的Nat减法是截断的，这让我在证明位置不变式时吃了不少苦头。比如要证明pos = base + i，其中base = totalLen - w + 1，当totalLen < w时，这个base会加一个巨大的偏差，而截断语义不会报错，只是算出来的结果不对。

普通语言里，totalLen - w如果是负的，通常会溢出或变成无符号大数，至少能引起注意；Lean4的截断会把负值变成0，反而让「错误的base」看起来合理。我遇到过的情况是：base的计算在不同路径上用了不同的边界条件，导致相同逻辑在不同长度下的位置记录不一致，而simp和omega对这种不变量之间的冲突会给出非常清晰的不可证明目标——这其实是在帮你定位实现中的不一致。

5.3 标准库的缺口：给Vector写引理的日常

如果说哪里最影响验证效率，那就是Lean4标准库对Vector操作的支持不如List完善。List有现成的drop、take、append、getLast?等大量引理；Vector的get、set、modify、追加和截断操作，很多引理都需要自己补。

我的建议是：核心语义验证用List，物理层用Vector或数组。如果因为类型安全想用Vector直接建模，你会陷入大量关于Fin索引和Vector操作性质的证明中，分心且收益有限。等以后Mathlib对Vector支持更成熟时，这个建议可以改。

5.4 证明开发效率：小步快跑与增量编译

最后分享几个提升Lean4验证效率的小技巧。

第一，多用#guard做快速数值验证。#guard (step 4 [] kv1) = [kv1]这类命令可以直接执行证明脚本中的函数，先验证定义的行为符合预期，再开始证明。这能提前拦截掉很多规格错误。第二，用native_decide验证有限域内的具体命题，它把命题编译成可执行代码再判定，对布尔性质的检查非常快。第三，避免给大定义加simp标签，否则每次simp都会展开它，编译速度会指数级下降；改成按需rw [defn]更可控。

还有一个开发流程上的建议：保持每个证明目标尽量小，引理拆细。宁可多写几个中间引理，也不要写一个巨型引理一口气完成。这在Lean4里尤其重要，因为当你面对一个复杂目标时，策略的选择和调试都会变得困难。

总结一下踩坑经验：

6. 这套验证思路能带走什么：从滑窗到更复杂的缓存策略

6.1 PagedAttention：块粒度逐出的形式化验证空间

滑动窗口KV Cache只是缓存管理的一种。vLLM的PagedAttention把KV Cache划分成固定大小的块，通过block table做逻辑块到物理块的映射，这种机制的正确性问题本质上也是「逻辑视图是否等于物理数据的正确解释」。

用同样的方法，可以定义block table的合法状态：物理块的内容加上block table的映射，解释出来的KV序列必须等于「全量历史的滑动窗口片段」。然后证明插入、逐出、复制等操作保持这个合法状态。这个验证比滑动窗口复杂得多，因为涉及逻辑块索引到物理块索引的映射，还有块内部的偏移，但思路是一致的：显式定义模拟关系，逐步证明操作保持它。

6.2 Prefix Caching：前缀复用的等价条件

另一个值得做形式化验证的场景是Prefix Caching——多个请求共享同一个prompt前缀的KV Cache。这里最微妙的正确性条件与位置编码有关：如果使用RoPE这类相对位置编码，前缀中同一个token在不同序列里的有效位置可能不同，复用时必须确保位置信息不会被错误继承。

形式化验证可以精确刻画「什么时候复用前缀是安全的」这个条件。比如可以证明：当且仅当两个序列的共享前缀后续部分在模型视角下位置完全一致时，复用的输出才与不复用的输出完全相同。这个命题如果不写清楚，很容易在长上下文下产生边界Bug。

6.3 不变量驱动的缓存设计：对普通工程的启发

做完这次验证，我最大的感想是：形式化验证对生产代码的最大影响，不是让你写出一个Lean4版本的缓存，而是逼你在写第一行C++之前定义清楚合法状态是什么。这种「不变量先行」的开发方式，即使后面完全不用证明工具，也能显著减少边界Bug。

一个可落地的做法是：在代码注释中明确写出核心不变量，然后包装成assert或运行时检查。比如滑动窗口缓存的核心不变量「缓存内容永远等于全量历史的最后W个K/V」，它可以直接翻译成一个debug_assert：把缓存导出成列表，和参考计算比对。这个检查在单测和模糊测试中非常有效。

6.4 形式化规范当测试预言：与模糊测试结合

最后分享一个成本很低但收益很高的实践：把Lean4里的规范函数作为模糊测试的预言（oracle）。让实现跑随机生成的token序列，然后用Lean4模型算出期望缓存，和真实实现比对。

这个做法的好处是：即使你不打算把证明写完，规范模型本身也是极好的一致性工具。由于规范函数和实现是分开写的，两个代码里同时出现同一个偏移错误的概率很低。我实际用这个方法抓出过一个C++实现里head指针更新顺序的问题，而那正是之前单元测试一直没有覆盖的边界路径。

如果你正在做KV Cache相关优化，并且对正确性要求很高，我强烈建议先从「规范函数+模糊测试」开始，再逐步把核心不变量补成形式化证明。这个路线投入可控，收益立竿见影，也为后续更深的验证留好了地基。