💻 computer science

Equivalence Checking of ML GPU Kernels

本文介绍了 Volta，这是首个针对 GPU 核函数的完备且可靠的等价性检查器，用于形式化验证由人工、编译器或大语言模型（LLM）优化的机器学习计算的正确性。

原作者： Benjamin Driscoll, Kshitij Dubey, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken

发布于 2026-08-18

📖 1 分钟阅读☕ 轻松阅读

原作者： Benjamin Driscoll, Kshitij Dubey, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

在现代人工智能那庞大且隐形的机器运作中，最关键的工作并非发生在云端，而是在被称为 GPU 的专用计算机芯片上。这些芯片旨在同时执行数百万次微小的计算，这是训练那些如今能够编写代码、翻译语言和生成艺术的大型语言模型的必要条件。为了让这些模型运行得足够快以发挥作用，工程师必须编写高度专业化的指令，即所谓的“内核”（kernels），这些指令精确地告诉 GPU 如何移动数据以及如何进行数学运算。在过去几年中，公司已开始使用人工智能本身来编写这些内核，希望能找到比人类工程师更高效的实现方式。然而，这种速度提升伴随着风险：当人工智能或编译器重写一段代码以追求更高速度时，可能会无意中引入细微的错误。这些错误可能导致计算机产生错误的答案，或者更糟的是，以一种通过标准测试几乎无法发现的方式发生静默崩溃。核心挑战在于，这些芯片同时执行着数千个线程的工作，如果它们不能完美协调，就会互相干扰，从而产生“竞态条件”（race condition），使得最终结果取决于事件发生的不可预测顺序。

一组研究人员开发了一种名为 Volta 的新工具来解决这个问题。Volta 并非通过猜测一个更快的内核版本是否正确，而是作为一个形式化验证器（formal verifier），从数学上证明两个版本会产生完全相同的结果。研究人员构建了一个系统，该系统获取参考内核（即原始、受信任的版本）和优化内核（即新的、更快的版本）的底层指令，并将其通过一个符号引擎（symbolic engine）进行处理。该引擎并非向代码输入具体的数字并观察输出，而是将输入视为抽象符号。它追踪代码可能采取的每一条路径，记录数据如何在数千个并行线程中移动，以及它们如何相互同步。如果代码尝试以可能导致冲突的方式访问内存，或者如果线程陷入永久等待，该工具会立即标记错误。如果代码运行顺畅，该工具会将两个内核的最终输出转化为复杂的数学表达式，并检查无论输入什么具体的数字，这些表达式在本质上是否相同。

研究人员在各种现实世界的机器学习任务上测试了 Volta，包括矩阵乘法、卷积以及驱动大型语言模型的注意力机制（attention mechanisms）。他们发现，该工具能够成功验证由人工优化、编译器优化甚至由大型语言模型生成的内核。在一个案例中，他们检查了一个经过十三轮自动化改进优化的 AI 生成内核。Volta 证实了该 AI 生成的代码在数学上与原始的人类编写的参考内核是等价的，证明了激进的优化并未破坏逻辑。该工具也通过捕捉其他方法遗漏的错误证明了其价值。例如，它检测到了一个在 GPU 编程领域被广泛引用、且已被数千名开发者使用了数年的流行教程中的数据竞争（data races）问题。这些错误之所以隐藏，是因为它们只在非常特定的定时条件下才会出现，而标准测试很少能捕捉到这种情况。该工具还识别出了一个 AI 生成内核中的漏洞：代码试图从一个不存在的内存位置读取数据；虽然当前的硬件恰好忽略了这个错误，但研究人员指出，该代码在本质上是不安全的，并且可能会在未来的机器上失效。

这种方法的优势在于其处理 GPU 编程独特复杂性的能力，即数千个线程必须协调其动作。以往的工具可以检查单线程程序或高层数学运算，但在将 GPU 的大规模并行性分解为易于管理的部分方面却表现不佳。Volta 通过假设其分析的内核遵循机器学习中常见的特定结构化模式（即线程数量和数据大小是预先已知的）来克服这一难题。在这一框架内，该工具可以确切地证明如果线程没有得到适当同步，则存在竞态条件；并且如果两个不同版本的程序产生相同的符号结果，它可以证明它们是等价的。研究人员展示了该工具即使对于包含数十万条指令的内核，也能在几秒钟或几分钟内完成验证。他们还证明了该工具背后的数学逻辑是完备的，这意味着如果工具判定两个程序相等，那么对于所有可能的输入，它们确实是相等的。

这项工作代表了使人工智能开发变得更加安全和可靠的重要一步。随着越来越多的公司依赖自动化系统来生成驱动其模型的代码，建立一种严谨的检查机制变得至关重要。研究人员表明，通过超越只能检查有限场景的简单测试，转向提供正确性形式化保证的方法是可行的。通过验证优化内核的等价性，Volta 让开发者能够放心地使用更快、更激进的优化手段，而不必担心引入静默漏洞。该工具目前已可投入使用，研究人员已将其代码及背后的证明过程向公众开放，以便他人在此基础上进行开发。虽然该工具目前尚未涵盖所有类型的 GPU 代码，但它已能成功处理驱动现代机器学习的大部分内核，为高性能计算代码的自动化生成树立了新的信任标准。

技术摘要：ML GPU 核函数的等价性检查

问题陈述

深度学习和大型语言模型（LLM）的快速发展导致执行 GPU 核函数（kernels）的开销巨大，这使得它们成为了激进优化的首要目标。这些优化由人工、编译器以及日益增多的大型语言模型（LLM）来完成。虽然经验测试和程序分析被用于验证这些核函数，但它们缺乏形式化保证。鉴于 GPU 编程的独特挑战，这一点尤为关键：

并发错误： 数千个线程之间的细粒度同步引入了难以通过测试检测到的微妙数据竞争（data races）和死锁（deadlocks）。这些错误可能仅在罕见的执行调度中显现，或者依赖于在不同硬件版本上无法保证的隐式同步。
浮点数语义： 优化过程经常重排算术操作，导致参考实现与优化实现之间产生微小的数值差异。虽然测试通常会容忍这些差异，但它们可能会掩盖真正的语义错误（例如，遗漏了裁剪操作）。
缺乏形式化工具： 现有的等价性检查技术支持单线程整数程序或高层张量操作，但无法处理 GPU 核函数中固有的并行性、同步和非确定性调度。目前尚不存在针对 GPU 程序的等价性检查器。

方法论

作者提出了 Volta，这是第一个专门为一类实用的 ML GPU 核函数设计的等价性检查器。Volota 作为一个黑盒验证器运行，通过分析 PTX 汇编（NVIDIA GPU 栈中最底层的文档化层面），而无需了解优化过程。

1. 目标类别：Structured-CTAs

Volta 针对 Structured-CTAs（协作线程阵列，Cooperative Thread Arrays）。该类别假设：

线程数量、张量大小和指针目标是静态已知的。
分支目标和内存访问地址可以根据线程 ID 和循环索引进行静态解析。
控制流不依赖于运行时数据值（例如，高效的 ML 核函数不会做出依赖数据的运行时决策）。
这些假设适用于许多高性能 ML 核函数（如 JAX 和 XLA 中的核函数），并且由 Volta 进行检查；如果违反这些假设，则会抛出异常。

2. 符号执行与汇合性（Confluence）

Volta 的核心是一个符号执行引擎，它将张量值视为实数（模拟优化器如 --ffast-math 使用的约定）。

符号状态： 引擎对参考核函数和优化后的核函数进行符号执行，推导出输出张量元素作为输入张量函数的表达式。
调度： 它采用轮询（round-robin）调度策略。线程执行直到遇到屏障（例如 syncthreads、syncwarp）。
汇合性： 一个关键的技术结果是证明了汇合性：符号评估器保证所有可能的执行调度都会导致相同的符号表达式（或报告竞争/死锁）。这一属性在 Agda 证明助手中得到了形式化验证，确保了检查器的结果独立于分析过程中选择的具体调度。

3. 竞争与死锁检测

Volta 通过扩展符号执行来检测利用屏障同步的数据竞争和死锁：

上下文追踪： 它维护一个上下文 X，用于追踪每个线程的内存事件（读/写）和同步状态。
竞争检测： 如果一个线程在未与上一个写入者/之前的读取者进行同步的情况下读写内存位置，则检测到竞争。
死锁检测： 如果线程在等待一个无法满足的屏障时进入阻塞状态，则报告死锁。
保证： 该系统被证明是可靠的（如果报告了竞争，则确实存在具体的竞争），并满足“真阳性”属性。

4. 判定程序

一旦生成符号表达式，Volta 必须判定参考输出和优化输出在数学上是否等价。

表达式形式： 这些表达式涉及多元多项式、加法、乘法和指数运算（例如 $p(\bar{x})e^{h(\bar{x})}$ ）。
可判定性： 作者证明了形式为 $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ 的等价性的可判定性。该证明依赖于这样一个事实：如果指数 $h_i$ 是不同的多项式，那么为了使总和恒等于零，系数 $p_i$ 必须为零。这避免了沉重的数论机制（如 Lindemann–Weierstrass），并能处理任意实系数。
规范化： 实现使用一个规范化器来归一化有理函数和指数项，处理诸如 softmax 重缩放的情况。

核心贡献

线性时间竞争/死锁算法： 一种利用屏更同步在 Structured-CTAs 中检测竞争和死锁的线性时间算法。
可判定性证明： 一个关于实数域上多元多项式和指数的等价性（特别是对于 $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ 形式）的可判定性证明。
Volta 实现： 首个针对实用 GPU 核函数的等价性检查器。它验证了卷积、矩阵乘法、归约（reduction）和注意力机制（attention mechanisms）的正确性。
形式化验证： 关键的汇合性引理在 Agda 中得到了形式化验证，且系统为定义的程序类提供了完备性和可靠性保证。

结果

Volta 在包括以下内容在内的多样化基准测试集上进行了评估：

人工优化核函数： 验证了来自标准教程的矩阵乘法和归约核函数的等价性。它正确地拒绝了包含由于弃用的 warp 同步执行模式导致的数据竞争的归约核函数版本（Red-5, Red-6, Red-7）。
LLM 生成的核函数： 验证了由 LLM 生成的核函数（例如用于 2D 卷积和矩阵乘法的核函数）。在其中一个案例中，Volta 检测到了 LLM 生成核函数中的越界共享内存读取，而 NVIDIA 的 Compute Sanitizer 或 LLM 自身的测试并未发现该问题。
编译器生成的核函数： 验证了由 TileLang 编译器生成的核函数的等价性，包括使用 Tensor Cores 的核函数。
注意力机制： 成功验证了标准 softmax 实现与优化的在线 softmax 公式（用于 FlashAttention）之间的等价性，尽管这些公式涉及复杂的指数和最大值运算的符号表达式。
性能： Volta 在不到 3 分钟内验证了复杂的注意力核函数（例如 FlashAttention-2）。在处理包含指数项的基准测试时，通过在多个输出元素间分摊规范化成本，其判定程序显著优于通用 SMT 求解器（如 Z3）。

意义与主张

论文声称提供了第一个针对 GPU 核函数的正式等价性检查器。其意义在于：

形式化保证： 提供比经验测试更强的信心，这对于在错误代价极高的生产部署中至关重要。
处理并发： 解决了以往等价性检查器无法处理的 GPU 同步（屏障、warp 级同步）的独特挑战。
对优化的鲁棒性： 通过将值建模为实数并检查数学恒等性而非位级相等，验证即使在优化重排浮点运算时的正确性。
实际适用性： 证明了“Structured-CTA”假设涵盖了广泛的现实世界 ML 工作负载（卷积、矩阵乘法、注意力机制），并且该方法可以扩展到现实规模的问题。

作者指出，虽然这项工作专注于 NVIDIA GPU（通过 PTX），但其技术也适用于其他运行时。他们同时也承认了局限性：该方法不支持具有动态控制流（如 top-k/argsort）或异步原语（如 pipeline, tma）的核函数，这些属于未来的研究方向。