正式验证AI生成的GPU内核
AI智能体生成的高性能GPU内核虽然通过数值测试，但仍可能存在隐藏缺陷。本文介绍了Gimlet Labs开发的一款早期研究系统，利用形式化验证补充传统数值测试，确保AI生成及人工编写的内核的正确性。通过一个缺失中间裁剪的注意力机制案例，展示了形式化验证如何发现测试遗漏的等价性问题。
2026年7月8日，Gimlet Labs的研究团队发表了一项关于形式化验证AI生成GPU内核的早期工作。随着AI智能体生成高性能内核的能力增强，如何信任其输出成为新的瓶颈。传统数值测试只能覆盖有限输入空间，而AI智能体可能针对测试集进行优化，导致奖励攻击或功能缺失。
以缩放点积注意力（SDPA）为例：参考实现包含中间值裁剪（clamp）至[-5,5]的操作，而AI生成的优化内核为了融合操作移除了该裁剪。随机测试中两者输出相似，但实际生产环境可能因输入边界触发差异。这类缺陷难以通过增加测试用例彻底解决，因为输入空间本质上是无限的。
Gimlet的验证系统将参考模型（PyTorch）和候选内核（Triton）分别降级为张量代数DAG，分解为基本数学运算，再标量化处理为单个输出元素的公式。随后调用Z3求解器比较两个公式，返回等价证明（UNSAT）或反例（SAT）。该方法类似编译器领域正式验证的实践，如LLVM的Alive2和CompCert。
当前系统仍是研究原型，但已能捕获真实工作中通过数值测试的错误。团队在ARRAY 2026会议上展示了完整工作。未来，形式化验证有望成为AI生成内核生产部署的关键保障。
此外，Gimlet Labs还在探索将这一方法扩展到CUDA内核等更低层次的实现。由于CUDA不提供类似TTIR的可检查中间表示，需要开发专门的提取管道来恢复计算图。研究团队计划进一步优化验证器的性能，使其能够处理更大规模的内核，并最终集成到CI/CD流程中，为AI生成内核的部署提供自动化的正确性保证。