jax_callable with WARP_STAGED or WARP_STAGED_EX fails when the same call runs on several GPUs in one process #2164

Bug Description

wp.jax_callable() with graph_mode=wp.JaxCallableGraphMode.WARP_STAGED or WARP_STAGED_EX fails when the same call runs on several GPUs in one process, for example under jax.shard_map or jax.pmap over a process's local devices. The first GPU's calls work, but calls on any other GPU fail:

• WARP_STAGED: Failed to update graph memcpy batch: Warp CUDA error 1: invalid argument (in function wp_cuda_graph_update_memcpy_batch)
• WARP_STAGED_EX: Failed to run input memcpy batch: Warp CUDA error 1: invalid argument (in function wp_memcpy_batch)

In a program with collectives, this can hang rather than raise. The other GPU waits indefinitely in XLA's collective rendezvous. Plain WARP mode works.

Reproduction (needs at least two GPUs):

import jax
import jax.numpy as jp
import numpy as np
import warp as wp


@wp.kernel
def double_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
    i = wp.tid()
    b[i] = 2.0 * a[i]


def double_func(a: wp.array(dtype=float), b: wp.array(dtype=float)):
    wp.launch(double_kernel, dim=a.shape, inputs=[a], outputs=[b])


mesh = jax.sharding.Mesh(np.array(jax.local_devices()), ("devices",))
spec = jax.sharding.PartitionSpec("devices")
f = wp.jax_callable(double_func, graph_mode=wp.JaxCallableGraphMode.WARP_STAGED_EX)
double = jax.jit(jax.shard_map(lambda v: f(v)[0], mesh=mesh, in_specs=spec, out_specs=spec))

x = jp.arange(len(jax.local_devices()) * 64, dtype=jp.float32)
double(x)  # the first call on each GPU captures
double(x)  # replays: fails on the second GPU

Cause: FfiCallable keeps one FfiCallDesc per call site. In the staged modes it holds a single capture, plus staging arrays and memcpy descriptors allocated on whichever GPU runs the call first. The replay paths only check that a capture exists, so other GPUs launch the first GPU's graph and copy through its staging buffers, from their own stream. WARP mode is unaffected because its graph cache key includes the device ordinal. The JAX-mode path has a related issue: it keeps only the last device's capture, which is what holds the capture's modules alive.

Expected: each GPU captures and replays its own graph and staging buffers.

Why it matters: the staged modes avoid re-capturing in WARP mode when XLA passes a call different buffer addresses, which happens when a call's inputs and outputs are XLA temporaries. In an RL rollout where a jax_callable steps the simulation, WARP mode re-captured the step graph once per training iteration, costing about 35 ms of CPU per GPU. With one process driving several GPUs, those re-captures also run back to back. WARP_STAGED_EX avoids the re-capture, but only works with one GPU per process until this is fixed. On two L40s, a patched WARP_STAGED_EX was about 1% faster than WARP in that workload.

I have a fix ready. It keeps the capture and staging state per device: a small per-device descriptor that FfiCallDesc looks up by device ordinal. It also adds a multi-device test covering WARP, WARP_STAGED and WARP_STAGED_EX under jax.shard_map. I'll submit it for review.

System Information

• Warp 1.18.0, and main as of 2026-10-26
• JAX 0.11.0 (CUDA 13), Python 3.12, Ubuntu 22.04
• 2× NVIDIA L40, CUDA driver 13.4