Tesseract lets you put a differentiable physics simulation inside a JAX or PyTorch training loop. Paired with Tesseract-JAX and 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 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¶ #
It all started with the winning entry of our first hackathon, Multi-Agent-DPC by Pietro Zanotta, Dibakar Roy, and Honghui Zheng, which later grew into the ICML 2026 paper CINOC. 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, 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 were built around components that did not want to be imported into Python directly. δSnowSac trains a PyTorch network through NOAA’s Fortran streamflow models, PRISMO takes one JAX gradient through a Julia semiconductor solver and a FEniCS optics solver, and Differentiable Silicon dispatches to a proprietary TCAD binary over SSH. 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.
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. Passing gpu_transport="cuda_ipc" to the Tesseract constructor is all you need (see also Keep GPU arrays on the device with cuda_ipc), and Tesseract-JAX picks it up from there:
import os
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],
}
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:
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.Tesseract.from_source(..., gpu_transport="cuda_ipc")enables the GPU transport for a source-built Tesseract (a containerized Tesseract works the same way throughTesseract.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¶ #
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):
| 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, 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.
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 ;Kochkov et al., 2021 ).
- 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 ;NeuralGCM for weather and climate). Idealized example: Ourlearned closure demo 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 ), as in Multi-Agent-DPC, which builds ondifferentiable predictive control .
- 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¶ #
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:
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 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 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, which does not support legacy CUDA IPC handles. 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 covers more common failure modes, such as JAX itself running on the CPU.
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 built on NVIDIA’s 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¶ #
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, 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:
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 andTesseract-Torch are the places to start if your training loop lives in one of those frameworks.
- New to Tesseract? The getting started guide 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!