利用NVIDIA Warp构建加速、可微分的AI物理仿真代码

2026年3月12日

先说几个趋势判断：计算机辅助工程正从人工驱动的工作流，加速迈向AI驱动的工作流，其中就包括那些能够在不同几何形状和运行条件下实现泛化的物理基础模型。不过，与大语言模型不同——这些物理模型高度依赖大量高保真、符合物理规律的数据，而非互联网上的文本。

近期关于计算流体动力学基础模型的规模定律研究揭示了一个关键瓶颈：仿真生成的训练数据，在实践中往往是成本的主要限制因素。问题来了——这对模拟器提出了怎样的要求？答案很明确：模拟器必须原生支持GPU、运算速度足够快，并且能直接接入机器学习工作流。

这时就要提到NVIDIA Warp了。它是一个用于加速仿真、数据生成和空间计算的框架，在CUDA与Python之间搭建了一座桥梁。开发者可以在Warp中编写高性能内核，它们看起来只是普通的Python函数，但这些函数会被实时编译成在GPU上高效执行的代码。与基于张量的框架不同——后者要求将所有计算表达为对整个N维数组的操作——在Warp里，你可以编写灵活的内核，这些内核会在计算网格的所有元素上同时执行，更符合物理直觉。

仿真内核通常在计算网格上表达，并且高度依赖数据相关的控制流，比如条件分支、提前退出、按元素变化的更新。这些模式在张量框架里处理起来较为棘手——你需要组合布尔掩码，不仅很快变得难以维护，还会在无关元素上浪费计算资源。但在Warp内核中，每个线程可以独立分支、跳过或退出，自然表达这些逻辑，完全不需要掩码这种变通方案。

更值得一提的是，Warp原生支持自动微分。这意味着，用Warp编写的求解器可以轻松实现可微分，直接集成到优化或训练工作流中。同时它还能与PyTorch、JAX、NumPy等框架互操作，适用于仿真、机器人、感知和几何处理等多种场景。

本篇内容将带你从头构建一个二维纳维-斯托克斯求解器，解释Warp的编程模型如何映射到偏微分方程求解器，然后通过对仿真进行微分来解决一个端到端的最优扰动问题。最后，我们会通过工业案例研究展示Warp在生产工作流中的实际能力。

如何使用Warp编写二维纳维-斯托克斯求解器

为了把重点放在Warp本身而非数值方法上，这里用一个教科书式的二维衰减湍流示例来说明，由不可压缩纳维-斯托克斯方程的涡量-流函数形式来描述。涡量ω根据输运方程演化：

∂ω/∂t + (∂ω/∂x)·(∂ψ/∂y) – (∂ω/∂y)·(∂ψ/∂x) = (1/Re) ∇²ω

流函数ψ通过泊松方程从涡量恢复：

∇²ψ = –ω

在周期性边界条件下，上述方程在傅里叶空间中将简化为代数方程，无需迭代求解器：

ˆψ{m, n} = ˆω{m, n} / (k²m + k²n)

这里 (km, kn) 是傅里叶空间中的波数对。求解器利用快速傅里叶变换（FFT），高效地在空间域和频率域之间来回切换。

每个时间步包含两个子步骤（图1）：首先，在L×L方形区域内的N×N网格上对涡量输运方程进行离散化，使用三阶强稳定性保持龙格-库塔（RK3）方案将解向前推进Δt，得到ω(t+Δt)；其次，在傅里叶空间中求解泊松方程，得到更新后的ψ(t+Δt)。

图1. 求解器单个时间步循环示意图

构建块1：有限差分离散化与时间推进

涡量输运方程中的对流项和扩散项，采用二阶中心有限差分来近似。高阶离散化当然也可行，但为简单起见，这里选择中心二阶格式。

图2. ω和ψ的有限差分模板

下面这个 rk3_update() 内核，负责计算扩散项和对流项，并执行单个RK3子步更新。而 step() 函数在每个时间步会调用这个内核三次——每个RK3阶段一次，每个阶段使用不同的系数。

@wp.kernel def rk3_update(n: int, h: float, re: float, dt: float, coeff0: float, coeff1: float, coeff2: float, omega_0: wp.array2d(dtype=float), omega_1: wp.array2d(dtype=float), psi: wp.array2d(dtype=float), omega_out: wp.array2d(dtype=float)): i, j = wp.tid() left = cyclic_index(i - 1, n) right = cyclic_index(i + 1, n) top = cyclic_index(j + 1, n) down = cyclic_index(j - 1, n) inv_h2 = 1.0 / (h * h) laplacian = (omega_1[right, j] + omega_1[left, j] + omega_1[i, top] + omega_1[i, down] - 4.0 * omega_1[i, j]) * inv_h2 inv_2h = 1.0 / (2.0 * h) j1 = ((omega_1[right, j] - omega_1[left, j]) * inv_2h) * ((psi[i, top] - psi[i, down]) * inv_2h) j2 = ((omega_1[i, top] - omega_1[i, down]) * inv_2h) * ((psi[right, j] - psi[left, j]) * inv_2h) rhs = (1.0 / re) * laplacian + j2 - j1 omega_out[i, j] = coeff0 * omega_0[i, j] + coeff1 * omega_1[i, j] + coeff2 * dt * rhs

rk3_update() 内核遵循单指令多线程（SIMT）范式——每个线程映射到计算域上的一个网格点，所有N×N个点通过一次 wp.launch() 调用同时更新。

wp.launch(rk3_update, dim=(self.n, self.n), inputs=[self.n, self.h, self.re, self.dt, stage_coeff[0], stage_coeff[1], stage_coeff[2], self.omega_0, self.omega_1, self.psi], outputs=[self.omega_tmp])

图3. 二维网格上ω的SIMT更新。线程(i,j)使用当前时间步模板中相邻单元格的值，将单元格(i,j)更新到下一个时间步

构建块2：FFT泊松求解器

Warp基于图块的原语，支持在傅里叶空间中求解泊松方程。关键操作是 wp.tile_fft() 和 wp.tile_ifft()，它们分别对加载到图块中的单行执行正向和反向FFT。一个N×N数组上的完整二维FFT可以分解为三个步骤：逐行FFT → 转置 → 逐行FFT。图4展示了 fft_tiled() 和 ifft_tiled() 如何计算正向和反向FFT。

@wp.kernel def fft_tiled(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)): i, _, _ = wp.tid() a = wp.tile_load(x, shape=(1, N_GRID), offset=(i, 0)) wp.tile_fft(a) wp.tile_store(y, a, offset=(i, 0)) @wp.kernel def ifft_tiled(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)): i, _, _ = wp.tid() a = wp.tile_load(x, shape=(1, N_GRID), offset=(i, 0)) wp.tile_ifft(a) wp.tile_store(y, a, offset=(i, 0))

图4. 在NxN网格上逐行tile_fft。每个块将一行加载到寄存器图块中，协作计算FFT，然后将结果存回全局内存

二维FFT还需要在逐行变换之间进行一次转置。可以用SIMT范式或者图块范式来实现。为简单起见，下面展示SIMT版本：

@wp.kernel def transpose(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)): i, j = wp.tid() y[i, j] = x[j, i]

把这三个内核组合起来：fft_tiled -> transpose -> fft_tiled，就得到了完整的二维正向FFT。逆向变换使用 ifft_tiled，遵循相同的模式即可。

整合构建块

示例中的 step() 函数还会依赖一些其他辅助内核。有了所有构建块，一次 step() 调用就能把仿真推进一个时间步。示例代码中的 self._solve_poisson() 方法，把 ω(t+Δt) → FFT → ˆω → (方程3) → ˆψ → IFFT → ψ(t+Δt) 这个流程封装成了模块化调用。

def step(self) -> None: for stage_coeff in self.rk3_coeffs: wp.launch(rk3_update, dim=(self.n, self.n), inputs=[self.n, self.h, self.re, self.dt, stage_coeff[0], stage_coeff[1], stage_coeff[2], self.omega_0, self.omega_1, self.psi], outputs=[self.omega_tmp]) self.omega_1, self.omega_tmp = self.omega_tmp, self.omega_1 self._solve_poisson() wp.copy(self.omega_0, self.omega_1)

运行求解器会产生图5所示的衰减湍流场。在GPU上，step() 函数通过 wp.ScopedCapture 被捕获到CUDA Graph中，后续所有帧使用 wp.capture_launch() 重放，消除了每次启动的开销。

图5. Re = 1,000时的二维衰减湍流

对求解器进行微分

现在已经构建了一个可工作的求解器，下一个问题很自然——如何让它变得可微分？

自动微分（AD）通过对计算图中每个基本运算应用链式法则，来计算程序的精确导数。与有限差分不同，AD不需要调整步长，产生的梯度可达机器精度。AD对PDE求解器的主要优势在于可扩展性：对于大规模网格上的复杂仿真，单次正向求解已非常昂贵，而有限差分等方法需要对N个输入执行O(N)次完整求解才能获得梯度。反向模式AD则只需大约一次正向传递加一次反向传递，即可计算出所有∂L/∂xi，这使得基于梯度的优化在生产分辨率下变得可行。这本质上与神经网络中的反向传播是同一思想——也是深度学习和大规模物理优化能够处理数百万自由度的根本原因。

Warp的自动微分系统在编译时，会为一个可微分仿真生成两个版本的程序：

• 正向版本：接收物理输入（初始条件、离散化控制律等），计算仿真输出（场量、导出量），以及伴随版本所需的中间数组。
• 伴随版本：正向仿真的自动生成对应部分，接收所选目标量相对于仿真输出的敏感性，并将其一直反向传播到输入。

这种反向传播会复用正向执行中的中间数组，在整个求解器上应用微分链式法则，无需构建大的符号表达式即可产生仿真伴随。

开发者只需编写正向物理过程，Warp负责梯度计算。任何需要可微分的 wp.array 在分配时设置 requires_grad=True，这会告诉Warp分配一个伴随存储数组。生成的伴随可以独立使用，或者与PyTorch、JAX互操作来实现端到端优化，包括训练机器学习模型。需要注意，目前Warp仅支持反向模式AD。

为了说明这个问题，我们来解决文献中提出的一个最优扰动问题。在湍流中，初始条件的微小扰动会随时间放大，显著改变流动的轨迹。识别哪些扰动增长最快，是迈向流动控制以及理解哪些流动结构具有动态重要性的基石。具体来说，目标是这样的：寻找初始涡量扰动 Δω，使得在提前时间T内，扰动轨迹与未扰动轨迹之间的差异最大化。

设FT 表示应用于T时间单位的前向求解器。未扰动轨迹为 ω∗ = FT(ω₀)，扰动轨迹为 ω̃ = FT(ω₀ + Δω)。目标是最小化均方误差：

MSE = –(1/N²)·‖ω∗ – ω̃‖₂²

负号的存在，将最大化轨迹差异转化成了最小化问题。为了约束优化，要求Δω的均方根不超过初始涡量场ω₀均方根的20%。

无就地修改

wp.Ta pe() 会记录前向传递中的内核启动，并在反向传递中重放它们以计算梯度。这只有在反向传递所需的中间值仍然可用时才有效——因此数组不能随意就地覆盖。这是与不可微分求解器的一个关键区别。在仅正向版本中，可以在每个时间步结束时交换两个数组 omega_0 和 omega_1。但对于可微分求解器，需要把RHS计算和RK3更新拆分成写入不同数组的独立内核。

在Warp中，所有中间数组都需要由用户显式定义。这意味着需要为每个时间步的每个RK子步预先分配单独的数组——这是任何可微分求解器的主要GPU内存成本。

self.omega_timestep = [wp.zeros((n, n), dtype=wp.float32, requires_grad=True) for _ in range(T + 1)] self.omega_stage = [] self.psi_stage = [] self.rhs_stage = [] self.fft_arrays = [] for _ in range(T): s_omega, s_psi, s_rhs, s_fft = [], [], [], [] for _ in range(3): s_omega.append(wp.zeros((n, n), dtype=wp.float32, requires_grad=True)) s_psi.append(wp.zeros((n, n), dtype=wp.float32, requires_grad=True)) s_rhs.append(wp.zeros((n, n), dtype=wp.float32, requires_grad=True)) s_fft.append({"omega_complex": wp.zeros((n, n), dtype=wp.vec2f, requires_grad=True)}) self.omega_stage.append(s_omega) self.psi_stage.append(s_psi) self.rhs_stage.append(s_rhs) self.fft_arrays.append(s_fft)

为每个中间状态存储Warp数组，会随时间步数线性增长，在长时间运行中会变得不太现实。一种常见的应对策略是梯度检查点：只保存选定的状态，然后在反向传递期间用正向求解器重新计算缺失的片段。这种方法以额外的正向计算换取更小的内存占用。

使用wp.Ta pe()记录梯度

有了预分配的数组，记录和微分正向传递就变得很简洁了：

with wp.Ta pe() as tape: forward() tape.backward(loss)

wp.Ta pe() 上下文会把每个 wp.launch() 调用记录到一个计算图中。tape.backward(loss) 反向遍历这个图，计算损失相对于Warp数组的导数。

优化循环

下面的代码展示了一个优化步骤。forward() 函数在扰动的初始涡量上运行，生成最终场和损失。优化器更新扰动，目标是减少损失。

with wp.Ta pe() as tape: forward() tape.backward(loss) optimizer.step([delta_omega.grad.flatten()]) tape.zero()

经过1000次迭代后，优化器发现了一个结构化的扰动Δω，让轨迹差异显著放大，MSE从接近零上升到了大约250。

图6. 1000次迭代中的优化过程及发现的扰动（右上）

Warp实践：AI驱动工业工作流的案例研究

XLB：可微分格子玻尔兹曼求解器

有机构构建了XLB——一个具有Warp和JAX两种后端的可微分格子玻尔兹曼求解器。在约1.34亿单元的一个基准测试中，Warp在单个40 GB GPU上比JAX快了约8倍，大致相当于JAX需要8个GPU才能达到的吞吐量。在更大规模上，Warp的内存使用减少了约2.5到3倍，并且成功完成了JAX在同一GPU上因内存不足而无法运行的最大案例。

图7. Warp与JAX的吞吐量和内存使用对比

MuJoCo Warp：大规模多体动力学

有机构最近发布了MuJoCo Warp——一个基于Warp的大规模多体动力学后端。在类似硬件上，Warp后端相比JAX实现了最高252倍和475倍的加速。MJWarp通过利用稀疏矩阵操作和推测执行，更精确地调度计算，同时保持了与JAX训练的即插即用兼容性。

图8. MJWarp与MuJoCo MJX在相应基准上的物理步吞吐量对比

AutoAssembler：AI驱动工业工作流

有公司的AutoAssembler ASI引擎，展示了Warp在AI驱动工业工作流中的价值，其应用范围超越了物理仿真。它通过直接从原始几何体计算接触、干涉和间隙，将全保真CAD装配转换为AI规划的运动约束。该技术使用针对大规模处理优化的Warp内核来实现空间智能。

在特定GPU上，Warp GPU后端相比优化的CPU基线实现了最高669倍的加速。该技术已经在头部原始设备制造商的制造工作流中得到实际应用。

图9. 联络图构建：CPU与AutoAssembler ASI引擎在五个复杂度递增的CAD装配上的对比

开始使用Warp进行计算物理应用

Warp让开发者能够在Python中编写物理和几何GPU内核，而不需要把一切强制塞进基于张量的框架里。在计算流体动力学中，时间推进和可微分求解能够清晰地映射到内核上，保持物理结构的完整性。

这个模型已经在工业AI工作流中得到验证，包括可微分CFD求解器、多体动力学研究以及空间推理引擎。通过与PyTorch和JAX的零拷贝互操作，Warp可以无缝接入机器学习流水线，同时保留这些工作负载所需的控制流——在性能、内存和可扩展性方面都带来了可量化的提升。