研究团队在 Lean 中形式化 Nvidia PTX 以证明 GPU 内核正确性

LLM 在优化 GPU kernel 方面已经相当擅长了。但在大规模训练任务中，我们仍然希望确信这些 kernel 确实做对了事情。测试 kernel 只能覆盖有限的情形，会漏掉竞态条件这类微妙的 bug。为了解决这个问题，过去几个月我们构建了一些系统，用来形式化证明 GPU kernel 是正确的。

我们此前在这一方向上的工作已经能证明许多 kernel 是正确的，但在可验证的 kernel 类型上仍有根本性限制。此后，我们探索了一种新的证明 kernel 正确性的方法：在 Lean 中形式化其语义，然后写一个经过计算机检查的形式化证明，证明它们是正确的。对于任意给定的 kernel，我们让 LLM 用 Lean 写一个证明，说明该 kernel 计算出了正确的结果，然后由 Lean 验证这个证明。这样一来，我们不需要信任证明本身；Lean 会替我们检查它是否正确。Lean 的 kernel 是一个小巧、经过充分审计的程序。如果它接受了这个证明，我们就可以非常有信心地认为该 kernel 是正确的——前提是在我们对 PTX 的形式化范围之内。我们选择在 PTX 层面工作，是因为所有 CUDA 程序都编译到 PTX，而且 PTX 有一个精确规定的内存一致性模型，我们可以将其形式化。我们使用 Lean，是因为它的证明检查器小巧且可信，也因为 LLM 已经相当擅长写 Lean 证明。这一新方法支持许多我们旧工作不支持的东西，包括数据相关的控制流、TMA、异步操作、使用 acquire-release 模式的细粒度同步、proxy 与 proxy-fence、生产者/消费者模式、megakernel 以及多 kernel 流水线。我们用它证明了 Hopper tensor-core GEMM 的简单实现、带 online softmax 的 FlashAttention、一个三 kernel 的 Mixture-of-Experts 流水线，以及一个带数据相关控制流的 megakernel 的正确性。

GPU 线程并不共享一个简单的、全局一致的内存视图。当线程 A 写入一个值、线程 B 读取同一地址时，B 能否看到 A 的写入，取决于它们操作之间的顺序关系。PTX 内存模型规定了这些规则。

该模型建立在内存事件之间的关系之上。一个事件是由某个特定线程在其执行中的某个特定点执行的一次内存操作（一次 load、一次 store、一次 barrier 到达）。我们在下面描述这些关系。

Program order（po）刻画单个线程内的顺序。如果线程 0 先写入 x，然后到达一个 barrier，那么这两个事件就由 program order 排序。

Reads-from（rf）记录一次读取从哪里获得它的值。如果读 R 返回了写 W 所存储的值，那么 W 就是 R 的 reads-from 来源。

Coherence order（co）对所有对同一地址的写入做部分排序。它表示在每个内存位置上，哪个写入“发生在”哪个之后。

Synchronizes-with（sw）刻画来自同步原语的跨线程顺序。当线程 0 到达一个 bar.sync barrier（PTX 中相当于 __syncthreads()），而线程 1 从同一个 barrier 返回时，线程 0 的到达 synchronizes-with 线程 1 的返回。

这些关系组合成 causality：program order 与 synchronizes-with 的传递闭包（再加上一些额外的边）。在推理内存可见性时，真正重要的是 causality。如果事件 A 在因果上先于事件 B，那么 A 的效果对 B 可见。

内存模型的核心约束是 causality 公理：如果写 W 在因果上先于对相同字节的读 R，那么 R 不能从任何在 coherence 上先于 W 的写读取。简单说，如果一个写在保证对某次读可见，那么这次读必须看到那个写或更新的值。

一个具体例子。两个线程共享一个全局变量 x，初始值为 0：

线程 0 线程 1 st.global.u32 [x], 1 bar.sync 0 bar.sync 0 ld.global.u32 r, [x]

线程 0 写入 x=1，然后到达 barrier。线程 1 从 barrier 返回，然后读取 x。我们把线程 0 执行的写记作“W”，把线程 1 的读记作“R”。类似地，把线程 0 中的 barrier 到达记作“A0”，把线程 1 中的 barrier 返回记作“A1”。Program order 将 W 与 A0 关联起来，也将 A1 与 R 关联起来。此外，synchronizes-with 将 A0 与 A1 关联起来。合在一起，这些关系意味着 causality 将 W 与 R 关联起来。换句话说，W 已经“传播到 R”，意思是 R 不能读取任何比 W 更旧的值。因此，它不能读到初始值 0。由于 W 是唯一的另一个写，R 必须读到 1。

这就是 barrier 有用的机制：它们在不同线程的内存操作之间建立了一条因果链。没有 barrier 的话，线程 1 合法地可以读到 0，因为写与读之间没有因果顺序。

完整的 PTX 模型还包含其他关系和公理（coherence、single-copy atomicity、SC-per-location 等）。我们形式化了除两个之外的所有这些公理：no-thin-air 公理和 Fence-SC 公理。我们省略它们是因为它们难以建模（更多细节见我们的论文）。省略这些公理可能会使我们无法证明某些 kernel 是正确的，但不会使我们把实际上错误的 kernel 证明成“正确”的。

PTX 内存模型是一种事后规范。给定来自所有线程的完整事件集合，它说明哪些 reads-from 和 coherence-order 配置是可能的（这被称为“consistent”）。它并不定义逐步执行过程。这给验证带来了循环依赖。一个线程产生哪些事件，取决于它的 load 返回了什么，因为存在分支和数据相关的控制流。但 load 返回什么，又取决于所有线程的完整事件集合，因为内存模型推理的正是这个集合。

我们用一个所谓的 oracle 来打破这个循环依赖。我们给每个线程一个“oracle 函数”。每当线程需要外部信息（一个 load 结果、一个 atomic 返回值、一个 warp shuffle 结果），

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）