ProofWright：用 AI Agent 给 LLM 生成的 CUDA 代码做形式化验证

一句话总结

LLM 能高效生成 CUDA 代码，但"能跑通测试"远不等于"正确"——ProofWright 用 AI Agent 驱动形式化验证工具，为 74% 的 LLM 生成内核提供了内存安全、线程安全乃至语义等价的数学级证明。

论文信息

机构：佐治亚理工学院、NVIDIA Research、斯坦福大学
作者：Bodhisatwa Chatterjee、Drew Zagieboylo、Sana Damani、Siva Hari、Christos Kozyrakis
发表刊物：J. ACM（ACM 期刊，正式论文）
关键词：CUDA 验证、LLM Agent、形式化方法、VerCors、Rocq

背景：LLM 生成 GPU 代码的"信任危机"

过去几年，以 GPT、Claude 为代表的大模型在 GPU 内核代码生成上取得了惊人进展。KernelBench 等基准测试表明，LLM 能从高层 PyTorch 规格自动生成高度优化的 CUDA 内核，极大提升了开发效率。

但这里有一个根本性的麻烦：LLM 生成的代码，我们怎么知道它真的是对的？

目前业界通行的做法是「跑几组单元测试，通过了就接受」。这看起来合理，却有两个致命弱点：

弱点一：测试覆盖率不足（False Negative）

以论文中给出的真实案例为例。GPT-4.1 生成了一个 sigmoid 内核，代码会对向量化处理（每个线程处理 4 个 FP16 元素），末尾元素用第一个线程（threadIdx.x == 0）处理。

这段逻辑在 N 是 4 的倍数时运行完全正常；但当 N 不是 4 的倍数时，多个线程块的 thread 0 会同时写入同一个输出位置，触发 race condition（竞争条件）。

这个 bug 会骗过几乎所有的单元测试——因为测试数据通常选 N=1024、N=4096 这样"整齐"的值。只有形式化验证才能穷举所有可能的输入规模，找到这个问题。

弱点二：Reward Hacking（奖励函数操控）

在以通过测试为目标的迭代优化中，LLM 有可能找到"让测试通过但实际不正确"的实现方式——例如硬编码特定输入的输出值，或者针对测试用到的具体数值进行特化处理。这类现象在 RL 训练中有充分记录，在代码生成中同样存在。

传统的动态测试工具（如 NVIDIA Compute Sanitizer）和符号执行工具各有局限：前者只检查给定输入下的执行，后者则在大型内核上无法规模化。

ProofWright 是什么

ProofWright 是一个端到端的自动化形式化验证框架，专门针对 LLM 生成的 CUDA 内核。它的核心思路是：用 AI Agent 来驱动形式化验证工具，从而弥合「LLM 高速生成」与「人工验证速度慢」之间的鸿沟。

整个系统由两个主要组件构成：

组件一：VerCors Agent（内存与线程安全验证）

VerCors 是一个程序验证工具，支持 GPU 程序（CUDA/OpenCL）的形式化安全性证明。它的工作方式是：程序员在代码中添加「权限注解」（annotations），说明每个线程对哪些内存位置有读/写权限，以及这些权限如何在同步点前后转移，然后 VerCors 用分离逻辑（Separation Logic）证明这些注解是一致的，从而保证：

内存安全：不存在越界访问、非法内存读写
线程安全：没有数据竞争（data race），线程间内存访问有序

问题在于，手写这些注解本身就需要专家知识——而 LLM 大量生成代码，靠人力完全跟不上。

ProofWright 的 VerCors Agent 解决了这个问题。它维护两个关键数据结构：

知识库（Knowledge Base）：包含 VerCors 语法文档和已验证成功的标注样例（包括来自 VerCors 官方文档的 8 个例子和 KernelBench 的 2 个例子），以及一个不断扩充的「已知错误及修复方案」数据集。

标注指南（Annotation Guide）：一个会随验证结果动态演化的指导文档。当 Agent 成功验证一个新内核时，它会把这次学到的标注模式（permission pattern）更新到指南中，供后续任务参考。

Agent 的工作流程大致如下：

接收 CUDA 内核代码
先对代码做预处理（简化 VerCors 不支持的 C++ 语法和最新 CUDA 特性）
参考知识库和标注指南，生成初始权限注解
调用 VerCors 验证；若失败，读取错误信息，查找知识库中的已知修复方案，修改注解后重试
若成功，将当前标注模式归纳，更新标注指南

组件二：语义等价框架（Semantic Equivalence Framework）

内存安全和线程安全解决的是"代码不会崩溃"的问题，但还没有回答"代码算的对不对"——即 CUDA 实现是否与原始 PyTorch 规格在数学上等价。

语义等价框架由前端和后端两部分构成：

前端：PyTorch → Rocq 翻译

PyTorch 静态分析器：利用 torch.fx 对 PyTorch 程序做符号追踪，提取带有张量元数据（形状、维度等）的数据流图（dataflow graph）
MLRocq 库：一个形式化验证库，用 Rocq 定理证明器实现了接近 100 个深度学习常用算子的精确数学定义，包括激活函数、归一化层、池化与归约操作、损失函数等。张量用递归类型 TensorND 表示（0维张量=标量整数，n维张量=n-1维张量的列表）
PyTorch to Rocq 翻译工具：将静态分析器提取的数据流图，借助 MLRocq 库翻译为对应的 Rocq 定理（theorem）

这个前端不使用任何 LLM，确保形式化规格的生成过程本身是可信的。

后端：语义等价 Agent（Rocq Agent）

后端接收 PyTorch 规格对应的 Rocq 定理，以及待验证的 CUDA 内核，执行以下步骤：

为 CUDA 内核生成对应的 Rocq 表示
构建「CUDA Rocq 表示 ≡ PyTorch Rocq 表示」的等价定理
自动构造定理证明（调用 Rocq 证明器）
将已证明的语义属性「下降」（lower）为 VerCors 的功能后置条件
用 VerCors 验证带注解的 CUDA 代码，建立 Rocq 规格与实际 CUDA 程序之间的桥接

以 HardTanh 函数为例，Rocq Agent 会自动生成如下 VerCors 注解：

/* @
// 函数正确性：HardTanh(x) = clamp(x, min_val, max_val)
ensures \gtid < size ==> {:output[\gtid]:} == (input[\gtid] < -1.0f ? -1.0f :
    (input[\gtid] > 1.0f ? 1.0f : input[\gtid]));
// 属性1：输出有界
ensures \gtid < size ==> {:output[\gtid]:} >= -1.0f && {:output[\gtid]:} <= 1.0f;
// 属性2：区间内保持原值
ensures \gtid < size && input[\gtid] >= -1.0f && input[\gtid] <= 1.0f ==>
    {:output[\gtid]:} == input[\gtid];
// ...
@ */

实验结果

数据集与基准

基准：KernelBench L1（100 个 PyTorch 程序，任务是生成等价的 CUDA 实现）
基线 CUDA 内核：由 Claude-4-Sonnet 和 GPT-5-Codex 生成
实验环境：MacBook Pro（Apple M4 Pro，macOS Redwood）

核心结论

验证层次	成功率
内存安全 + 线程安全	74%（74/100 个内核）
语义等价（全部正确）	14%（14/100 个内核）
语义等价（部分，两个内核之一）	3%
未能验证	26%

74% 的内存与线程安全：Agent 成功为 74 个内核生成了正确的 VerCors 权限注解，涵盖矩阵乘法、MaxPooling、归一化等多种算法类型。

成功案例中有一个值得关注的细节：对于 MaxPooling2D 内核，Agent 正确识别出「每个线程写入唯一的输出索引，但读取范围可能重叠」这一模式，并生成了相应的读/写权限配置——这需要对 CUDA 并发语义有相当深入的理解。

26% 未验证的失败原因：

约束无法满足（Unsatisfiable Constraints）：某些内核本身就存在 race condition 等真实 bug，VerCors 正确地拒绝了它们——这实际上证明了 ProofWright 能发现 LLM 生成代码中常规测试遗漏的真实问题
VerCors 工具局限性：VerCors 尚不支持原子操作、warp-level 原语等高级 CUDA 特性
Agent 生成注解失败：复杂内核（如共享内存+全局内存混合访问）偶尔超出 Agent 的推理能力

14% 的语义等价：全量验证（内存安全 + 线程安全 + 语义正确性）在元素级一对一映射的内核上效果最好（如 relu、matrix-scalar multiply、hardtanh）。对于需要归约、共享内存聚合等复杂计算模式的内核，Rocq Agent 目前还无法生成足够强的全局不变量（global invariant）。

时间开销

每个内核约 3 分钟，作者认为这是可接受的——相对于 LLM 生成代码的高速率，这只是一个适度的验证税。

技术亮点深析

知识库驱动的自适应学习

VerCors Agent 不是每次都从零开始推理。它维护一个不断更新的「已知错误-修复方案」数据集——当 Agent 遇到新的验证错误时，先检索知识库里是否有类似错误的修复先例，若有则直接套用，大幅减少重试次数。

这个设计思路很有意思：它把「如何写形式化验证注解」这项专家知识，转化为了一个可以机器学习的知识库，实现了人类专家经验的系统性积累和自动化迁移。

标注指南的渐进演化

Annotation Guide 是一个随成功案例增长而不断细化的指导文档。Agent 完成验证后，会主动提炼「这次用了什么 permission pattern，适用于什么类型的内核」，更新到指南中。

为了防止过拟合（训练集只有少量手工验证的内核），论文在这一更新步骤引入了人工监督——这是一个有效的工程折中。

前端无 LLM 的信任保证

PyTorch → Rocq 翻译的前端（静态分析器 + MLRocq 库 + 翻译工具）完全不使用 LLM。这意味着：「我们要验证的目标规格」本身是可信的，系统的可信基础（Trusted Computing Base）清晰且有限。LLM 只参与「生成注解」和「构造证明」这两步，即便出现幻觉（hallucination），VerCors 和 Rocq 的检验也会兜底。

这一设计原则（把 LLM 限制在「助手」角色，用确定性工具作为最终裁判）值得所有 AI 安全系统借鉴。

局限性与开放问题

论文作者对系统局限性相当坦诚：

VerCors 工具自身的覆盖边界：不支持 Tensor Core 操作、warp-level 同步、原子操作，现代 GPU 优化内核大量使用这些特性，意味着 26% 的验证失败中有相当部分是工具局限而非代码问题
语义等价验证的局限：当前的 Rocq Agent 只能处理「每个线程计算一个输出元素」的一对一映射模式；对于 reduction、shared memory tile、scan 等模式，建立全局不变量所需的数学推理超出了当前 LLM 的自动证明能力
Annotation Guide 的小样本风险：初始知识库只有 10 个手工验证的内核，Agent 的泛化能力依赖于这个有限的基础，在分布外的新型内核上可能退化
语义等价桥接的手工验证：从 Rocq 到 VerCors 的「下降」（lowering）步骤目前是由 Agent 完成的，但论文承认这些生成的注解是人工校验而非自动证明的——留下了一个小的可信缺口

我的思考与启示

「AI 守门人」模式的价值

ProofWright 体现了一个重要的设计范式：用 AI 来验证 AI 的输出。这不是说 AI 自己能证明自己是对的——而是说，AI 可以驱动传统的数学工具（VerCors、Rocq）去做验证，把「人工注解的成本」这个瓶颈突破掉。

这个范式在安全关键系统中有巨大潜力。GPU 内核只是开始——编译器优化、操作系统内核、密码学实现，都是类似的高风险、高专业性领域。

测试覆盖率的本质限制

论文对「测试不够可靠」的论证很有说服力。那个 sigmoid 内核的 race condition 案例是一个很好的教学样本：bug 只在 N 不是 4 的倍数时出现，而大多数测试用的都是"整齐"的数字。

更深一层：LLM 生成代码的过程本身就是在「训练分布上推理」，而测试数据往往也来自相似分布。测试分布与 bug 触发条件的重叠越少，测试漏掉 bug 的概率越高。形式化验证恰好不依赖分布——它对所有可能的输入都成立。

AI 能力边界的清晰化

14% 的完整语义等价率，听起来不高。但我觉得这是一个诚实的数字——它清晰地划出了当前 LLM 自动证明能力的边界：一对一线程-元素映射的简单内核可以做，有全局状态依赖的复杂内核目前还不行。

这比一个模糊的「我们的方法在大多数情况下有效」要有价值得多。清晰的边界才能指导工程实践：在 AI 能覆盖的范围内用 ProofWright，在覆盖不了的范围内用其他手段。

对 LLM 代码生成系统设计的启示

ProofWright 提示了一个关于 AI 代码生成系统架构的重要原则：生成与验证应该是两个独立的流水线，而不是把所有信任都压在「生成质量」上。

当前很多 AI 编程工具（包括 GitHub Copilot、Cursor 等）主要靠测试反馈来迭代，验证的严格程度取决于用户是否写了足够的测试。ProofWright 的实践表明：把形式化验证集成进生成流水线是可行的，开销约 3 分钟每内核，完全在工程可接受范围内。

下一步的有趣问题是：能不能把验证失败的信息作为反馈信号，驱动 LLM 重新生成更容易被验证的代码？ 目前论文没有涉及这个闭环，但这是一个很自然的扩展。

MLRocq 的独立价值

论文中有一个不太显眼但很有价值的贡献：MLRocq 库——接近 100 个深度学习算子的 Rocq 形式化定义。

这个库的价值超出 ProofWright 本身。它为整个深度学习系统栈提供了一个可用于形式化推理的数学基础。理论上，任何声称实现了某个算子的代码，都可以用 MLRocq 来验证其语义是否与数学定义一致——不限于 CUDA，也可以用于 Triton、MLIR 等其他 GPU 编程范式。

总结

ProofWright 是一篇时机恰好的论文：在 LLM 大量生成 GPU 代码的当下，它给出了「如何系统性地信任这些代码」的第一个可规模化答案。74% 的内存/线程安全验证率说明这套方法在工程上是可行的；14% 的语义等价率则清晰地标出了当前技术边界。

这项工作最深的意义或许在于：它建立了一个人类专家知识可以被机器积累和迁移的框架——通过知识库和标注指南的动态演化，每一次成功的验证都让下一次更容易。在「AI 生成代码」与「人类信任 AI 代码」之间，ProofWright 是目前最认真的一个桥接尝试。