Fast, GPU-native differentiable physics and beyond with Tesseract

Tesseract Core v1.15.0, Tesseract-JAX v0.5.0 and Tesseract-Torch v0.2.0 add GPU transport that keeps arrays on the GPU across process boundaries, sending only a 64-byte handle instead of copying data through host memory, Pasteur Labs' Tesseract team reported. Combined with better batching and a forward-pass cache, the change brings the Multi-Agent-DPC swarm-heater training step within 20% of importing the solver directly, and for the large arrays common in 3D simulation transport is now two orders of magnitude faster than the host-memory round trip. The team said the earlier path, which encoded each array as text in a JSON body over HTTP, made GPU execution slower than running the solver on CPU.

← Tesseract Blog ../ Fast, GPU-native differentiable physics and beyond with Tesseract ¶ fast-gpu-native-differentiable-physics-and-beyond-with-tesseract Tesseract lets you put a differentiable physics simulation inside a JAX or PyTorch training loop. Paired with Tesseract-JAX https://github.com/pasteurlabs/tesseract-jax and Tesseract-Torch https://github.com/pasteurlabs/tesseract-torch , the simulator runs in its own process or container, while the training loop treats it like any other differentiable function. Until now, every GPU array passed into or out of a Tesseract was copied through host memory, which for large arrays can take longer than the computation itself. The new GPU transport ../../content/how-to/gpu-transport/ removes this detour, so arrays stay on the GPU the whole way. This post shows what that unlocks on a real application, how to turn it on, and what happens under the hood when passing GPU memory up and down the call stack and over HTTP between processes . The problem ¶ the-problem It all started with the winning entry of our first hackathon, Multi-Agent-DPC https://github.com/SOLARIS-JHU/Multi-Agent-DPC by Pietro Zanotta, Dibakar Roy, and Honghui Zheng, which later grew into the ICML 2026 paper CINOC https://arxiv.org/abs/2605.25867 . It learns an optimal control policy for a swarm of mobile heaters that shape the temperature along a rod. The policy is a small neural network that decides how hard each heater heats or cools and where it moves, and it is trained by backpropagating through 300 steps of a differentiable PDE solver. The team wrapped the solver as a Tesseract and called it from their JAX training loop through Tesseract-JAX. It worked but was too slow to train with, so the team’s training script ended up calling the solver directly with the Tesseract path switched off. They published a branch that reproduces the slowdown https://github.com/SOLARIS-JHU/Multi-Agent-DPC/tree/repro-performance-issue , and over the following months we used it as a benchmark to make Tesseract faster for differentiable physics applications, one release at a time. A deeper problem was the data communication overhead. Each array left the GPU, was encoded as text into a JSON body, sent over HTTP, and was decoded and copied back onto the GPU on the other side, once for the forward pass and once for the gradient. Through that path, running the solver on the GPU gained so little that it was faster to run it on CPU instead. With the new GPU transport capabilities in Tesseract Core v1.15.0 , Tesseract-JAX v0.5.0 , and Tesseract-Torch v0.2.0 , arrays now stay where they are. Only a 64-byte handle is sent across the process boundary, and the data moves only within the memory of the GPU, which is shared by both processes. Together with better batching support and a cache for the solver’s forward pass, this brings Multi-Agent-DPC’s training step within 20% of importing the solver directly. For the much larger arrays common in 3D simulation, transport is now two orders of magnitude faster than the trip through host memory. Why not just import the solver? If you can import a component, import it. The Multi-Agent-DPC team came to the same conclusion, and the code for CINOC imports its solvers directly. But often you can’t. Many entries to this year’s Tesseract Hackathon ../2026-09-30-tesseract-hackathon-winners/ were built around components that did not want to be imported into Python directly. δSnowSac https://github.com/kasProg/dSnowSac trains a PyTorch network through NOAA’s Fortran streamflow models, PRISMO https://github.com/benvial/prismo takes one JAX gradient through a Julia semiconductor solver and a FEniCS optics solver, and Differentiable Silicon https://github.com/hozaifa1/differentiable-silicon dispatches to a proprietary TCAD binary over SSH. Mosaic https://github.com/pasteurlabs/mosaic puts solvers from deal.II and OpenFOAM to JAX-FEM and XLB behind the same Tesseract interface, each with its own dependencies. Tesseract puts such a component in its own process or container and still allows it to participate in differentiable pipelines. See also When to use Tesseract ../../content/concepts/when-to-use/ . A faster, GPU-native Multi-Agent-DPC ¶ a-faster-gpu-native-multi-agent-dpc Let’s turn on GPU transport for the Multi-Agent-DPC training loss, with the solver served from the team’s tesseract api.py https://github.com/dionhaefner/Multi-Agent-DPC/blob/dca8d074eacc9ea97a1cd43740f115340c437dff/tesseracts/solverHeat decentralized/tesseract api.py . Passing gpu transport="cuda ipc" to the Tesseract constructor is all you need see also Keep GPU arrays on the device with cuda ipc ../../content/how-to/gpu-transport/ , and Tesseract-JAX picks it up from there: python import os NEW: two JAX processes now share one GPU, so neither should claim 75% of it at startup. os.environ "XLA PYTHON CLIENT PREALLOCATE" = "false" import jax from jax.flatten util import ravel pytree from tesseract core import Tesseract from tesseract jax import apply tesseract solver = Tesseract.from source "tesseracts/solverHeat decentralized/tesseract api.py", gpu transport="cuda ipc", NEW: pass GPU arrays by handle 2 def loss params, z init, xi init, z target : inputs = { "z init": z init, "xi init": xi init, "z target": z target, "flat params": ravel pytree params 0 , } NEW: one call for the whole batch instead of one per rollout 1 traj = apply tesseract solver, inputs, vmap method="broadcast all" return tracking error traj "z trajectory" , z target + penalties traj def batch loss params, z init, xi init, z target : losses = jax.vmap loss, in axes= None, 0, 0, 0 params, z init, xi init, z target return losses.mean with solver: grads = jax.jit jax.grad batch loss params, z init, xi init, z target This deviates from the original script in two ways: 1. apply tesseract ..., vmap method="broadcast all" sends the whole batch of rollouts to the Tesseract as one call, which the solver supports by accepting an optional leading batch dimension. 2. Tesseract.from source ..., gpu transport="cuda ipc" enables the GPU transport for a source-built Tesseract a containerized Tesseract works the same way through Tesseract.from image ..., gpus= "all" , gpu transport="cuda ipc" . For applications that use torch instead of JAX, Tesseract-Torch works in exactly the same way with an identical apply tesseract function. Turns out it’s pretty good ¶ turns-out-it-s-pretty-good Here is one training step of Multi-Agent-DPC, that is, a batch of 32 rollouts forwards and backwards, as it went from the team’s setup to the GPU path benchmark script https://github.com/dionhaefner/Multi-Agent-DPC/blob/dca8d074eacc9ea97a1cd43740f115340c437dff/examples/heat1d/decentralized/benchmark.py : | Setup | Time per training step | |---|---| | The original setup, on current Tesseract releases | 1.5 s | | Batched, solver on the GPU, arrays through host memory | 261 ms | | Same, with cuda ipc | 209 ms | | Same, with the solver’s VJP cache | 127 ms | | Solver imported directly, one process baseline | 106 ms | The GPU transport removes the trip through host memory, which saves about 50 ms per step. The VJP cache ../../content/reference/tesseract-runtime-api/ tesseract core.runtime.experimental.set jax vjp cache size , an experimental option in Tesseract’s JAX recipe, saves another 80 ms by keeping the residuals of the forward pass, so the gradient call doesn’t rerun all 300 steps. What remains, about 20 ms per step, is mostly the fixed cost of two HTTP requests and the handle exchange. A Multi-Agent-DPC call moves about 4.8 MB in and out combined, which is modest. How much the transport matters depends mostly on how many bytes are sent across the boundary. The figure below isolates this with a Tesseract that only doubles its input. Below about 100 kB per call, the trip through host memory is the fastest option, because the GPU transport has a slightly higher fixed cost. Above that, the host path grows linearly by about 8 ms per megabyte, spent on copies and text encoding on the CPU, while cuda ipc stays nearly flat. At 100 million values, or 400 MB, the GPU transport is 190 times faster than Tesseract’s default host path, and still almost 60 times faster than exchanging raw binary data via shared memory ../../content/how-to/fast-local-runs/ . Beyond Multi-Agent-DPC ¶ beyond-multi-agent-dpc So far, using Tesseract hasn’t really been feasible for GPU-based training through a differentiable solver. The new GPU transport changes that, making it practical for scenarios such as: - Solver in the loop. A network learns to correct a coarse simulation, and its loss is computed after unrolling the solver + network system for many steps Um et al., 2020 https://arxiv.org/abs/2007.00016 ; Kochkov et al., 2021 https://arxiv.org/abs/2102.01010 . - Learned closures. A network supplies a term the simulator can’t compute efficiently, such as subgrid-scale stresses in turbulence, and is trained through the simulation List et al., 2022 https://arxiv.org/abs/2202.06988 ; NeuralGCM https://arxiv.org/abs/2311.07222 for weather and climate . Idealized example: Our learned closure demo ../../content/demo/learned-closure/ does this for a viscosity model in Burgers’ equation. - Control and reinforcement learning. A network is learning to interact with a physical system by training against a differentiable simulator Holl et al., 2020 https://arxiv.org/abs/2001.07457 , as in Multi-Agent-DPC, which builds on differentiable predictive control https://arxiv.org/abs/2011.03699 . - Served surrogates. A trained surrogate runs as a service on the GPU, and an optimizer or pipeline calls it over and over. In all of these, the network and the solver both run on the GPU, and training alternates between them thousands of times. Most move far bigger arrays than Multi-Agent-DPC. Sending a single 256³ float32 field 16.8 million values through the host to a Tesseract and getting one back takes over half a second. With the new GPU transport, it takes about 6 ms. How the gradient flows up and down the stack ¶ how-the-gradient-flows-up-and-down-the-stack You don’t need to know any of this to use the GPU transport, but to us, how it works is the most remarkable part of the story. Specifically, a single gradient call drops out of compiled XLA code into C++, goes back into the Python interpreter, crosses HTTP into another process, descends into that process’s own XLA, and returns the same way, while the arrays never leave the GPU Most of that path is specific to Tesseract-JAX. Tesseract-Torch reaches the same place without a C++ bridge, for reasons we come back to at the end. Take these five lines, a stripped-down version of the training loss above: python def loss theta : traj = apply tesseract solver, {"theta": theta} "trajectory" return jnp.mean traj - target 2 with Tesseract.from source "./tesseract api.py" as solver: grads = jax.jit jax.grad loss theta When jax.grad runs, JAX asks Tesseract-JAX for a primal and a vector-Jacobian product VJP . Tesseract-JAX turns this into two requests to the solver, one to apply for the trajectory and one to vector jacobian product for the gradient. Here is how the gradient with respect to the trajectory, 3.84 MB of float32 for a batch of 32 rollouts, is represented in the body of the second request: { "object type": "array", "shape": 32, 300, 100 , "dtype": "float32", "data": { "buffer": "slPUPa5mfD5CBcs9OirIvgsQiz5TIAk+uPMkvv2EMj5I/t89C7e0PXWzPRUWPdbiFD2WZBI9SKoPPeagDD1KNAk96E8FPZzfAD3UoPc8KCPsPKgq3zwIo9A8EIHAPEDDrjyocps8AKOGPFDlYDyAFDI8sDYBPEB9nTsAXtk6AMzHukBHmruQwwC8gEwzvGA3ZLwIg4m8AKOfvABJtLzwSce8CIPYvNjZ57w4PPW8PFAAvXYCBb2otwi9OHYLvcdHDb3AOA69C1gOvaO2Db0gZwy9VH0KvdANCL1xLQW95PABvWDY/LyIZPW8BKntvLTE5bz40d28EObVvJAQzrwwW8a8yMm+vGBat7xYBbC8AL6ovBBzobyoD5q8CHySvGCeirxgXIK84DdzvCCIYLwgfEy8AO42vIC/H7yg3Aa8AHjYuwDBn7uAZUe7AA6SugDgazoAzkM7QEOoO0AF8DsAShw8YKZAPIDEZDywKIQ8IH6VPNA8pjxwQrY8cHHFPOCx0zxA8eA8qCPtPMBD+DwoKQE9/KoFPWytCT3YOA09blcQPcwUEz14fRU9Up4XPS+EGT1uOxs95gwcPXezFj3IfxU9...", "encoding": "base64" } } { "object type": "array", "shape": 32, 300, 100 , "dtype": "float32", "data": { "buffer": "80c3660b-f7b0-48f9-8ad6-7f36c1e7e9bf:xAz7Rvl/rnXSpn869f++dNC9NWxplFo3qLs+qce5gyNA1eMWkRWGatgyGZfLq4e001iiEs3oW+hxjxR7Q+f2cg==:0:3840000", "encoding": "cuda ipc" } } The base64 buffer goes on for another 5.1 million characters, while the cuda ipc buffer is complete. It contains the GPU’s UUID, a 64-byte CUDA IPC https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/inter-process-communication.html handle that another process on the same machine can use to map the same allocation, and the position of the array within that allocation. Shapes and dtypes stay in the JSON payload, so the server still validates every request against its schema. Getting from those five lines to that request and back takes the path below, spanning Python, C++, CUDA, and HTTP across process boundaries. The C++ bridge of Tesseract-JAX is needed because under jit , XLA needs to know how to run the primitive without copying data. On the CPU everything runs as a host callback into Python which receive NumPy arrays, so XLA copies every input off the device before Python sees it, and no code downstream could undo that. JAX has no Python-level API for handing GPU buffers out of a compiled program, so the fix had to be native. On CUDA, Tesseract-JAX now lowers the primitive to an XLA FFI https://docs.jax.dev/en/latest/ffi.html custom call into a small C++ shim, which receives raw device pointers from XLA and waits for XLA’s stream so the inputs are ready. The shim loads the CUDA runtime only when first used, so the same wheel installs and runs on machines without a GPU. The next step is the craziest one: While XLA is in the middle of running a compiled program, the shim called from one of XLA’s threads takes Python’s global interpreter lock and calls back into the Python interpreter . Each call site keeps track of an integer token, baked into the compiled program as an FFI attribute, which the shim uses to look up the Python closure for that particular call. The closure receives the device pointers as minimal objects that expose cuda array interface , passes them to the Tesseract’s Python client, and waits for the solver to answer. Only then does control return to C++, and from there to XLA. Another complication is that XLA allocates GPU memory through CUDA’s virtual memory management API https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/virtual-memory-management.html , which does not support legacy CUDA IPC handles https://developer.nvidia.com/blog/introducing-low-level-gpu-virtual-memory-management/ . So the client first copies each array into a buffer it can export and reuses those buffers from call to call . For PyTorch, Tesseract-Torch needs none of this. Its tensors already expose their device pointers to Python, so Tesseract-Torch calls the same client without native code, and with PyTorch’s default allocator also without the extra copy. Most mistakes in code like this return wrong numbers or quietly fall back to the host instead of crashing, so there was a lot of stress-testing involved. We had to set up GPU tests to reject any pointer that isn’t in GPU memory, so an accidental trip through the host fails loudly. And a Tesseract created with gpu transport="cuda ipc" checks before its first call that the transport works between the two processes, and raises with the reason instead of falling back to the host. The troubleshooting section of the how-to ../../content/how-to/gpu-transport/ gpu-transport-troubleshooting covers more common failure modes, such as JAX itself running on the CPU. What’s next ¶ what-s-next Today, Tesseract’s GPU transport is built on CUDA IPC, which works great and is fast, but it means the caller and the Tesseract must have direct memory access to the same GPU. However, GPU transport selection is already fully generic in Tesseract, so adding support for other transport mechanisms can be done without changing the user experience or the core logic of Tesseract. A logical next step is supporting communication across multiple machines. An experimental nixl transport https://github.com/pasteurlabs/tesseract-core/pull/723 built on NVIDIA’s NIXL https://github.com/ai-dynamo/nixl transfer library moves GPU memory through CUDA IPC on the same host and through UCX between hosts. This includes RDMA over NVLink or Infiniband reading directly from and to device memory where the hardware supports it. Triggering a request still works with a short descriptor instead of the full array bytes, so with GPUDirect RDMA, a training loop on one machine could call a GPU solver on another without the arrays passing through host memory on either side. Try it ¶ try-it The GPU transport has a few requirements on both sides of the boundary, from the GPU and CUDA setup to a Tesseract that can consume and return GPU arrays. Start with the GPU transport how-to ../../content/how-to/gpu-transport/ , which walks through them and covers every way to serve and call a Tesseract with the transport. Once those are in place, enabling it takes one extra argument: python from tesseract core import Tesseract from tesseract jax import apply tesseract solver = Tesseract.from image "my-solver", gpus= "all" , gpu transport="cuda ipc" with solver: outputs = apply tesseract solver, inputs GPU arrays in and out, no host copies python from tesseract core import Tesseract from tesseract torch import apply tesseract solver = Tesseract.from image "my-solver", gpus= "all" , gpu transport="cuda ipc" with solver: outputs = apply tesseract solver, inputs CUDA tensors in and out, no host copies python from tesseract core import Tesseract solver = Tesseract.from image "my-solver", gpus= "all" , gpu transport="cuda ipc" with solver: outputs = solver.apply inputs CuPy, JAX, or PyTorch GPU arrays in, device arrays out bash $ tesseract serve my-solver --gpus all -e TESSERACT GPU TRANSPORT=cuda ipc Tesseract-JAX and Tesseract-Torch use the transport for a Tesseract created with Tesseract.from url , as long as it works from the calling process. From there: - Tesseract-JAX https://github.com/pasteurlabs/tesseract-jax and Tesseract-Torch https://github.com/pasteurlabs/tesseract-torch are the places to start if your training loop lives in one of those frameworks. - New to Tesseract? The getting started guide ../../content/tutorials/get-started/ covers building and serving your first Tesseract. If you have a GPU component that lives in another stack, or a training loop that currently bounces arrays through the CPU, we’d like to hear how the new GPU transport works for you