From 893882171eb0b50ec9c23d704cf1c5bca87cfc94 Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Tue, 25 Nov 2025 06:25:49 -0800
Subject: [PATCH] MuJoCo Warp documentation: Does MJWarp work with multiple
 GPUs?.

PiperOrigin-RevId: 836642204
Change-Id: I559642b2cb902a704257e7e0344f768b690678a1
---
 doc/mjwarp/index.rst | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst
index dbf68489..a4289930 100644
--- a/doc/mjwarp/index.rst
+++ b/doc/mjwarp/index.rst
@@ -383,6 +383,32 @@ Warp's `automatic differentiation <https://nvidia.github.io/warp/modules/differe
 functionality. Updates from the team related to enabling automatic differentiation for MJWarp are tracked in this
 `GitHub issue <https://github.com/google-deepmind/mujoco_warp/issues/500>`__.
 
+**Does MJWarp work with multiple GPUs?**
+
+Yes. Warp's ``wp.ScopedDevice`` enables multi-GPU computation
+
+.. code-block:: python
+
+   # create a graph for each device
+   graph = {}
+   for device in wp.get_cuda_devices():
+     with wp.ScopedDevice(device):
+       m = mjw.put_model(mjm)
+       d = mjw.make_data(mjm)
+       with wp.ScopedCapture(device) as capture:
+         mjw.step(m, d)
+       graph[device] = capture.graph
+
+   # launch a graph on each device
+   for device in wp.get_cuda_devices():
+     wp.capture_launch(graph[device])
+
+Please see the
+`Warp documentation <https://nvidia.github.io/modules/devices.html#example-using-wp-scopeddevice-with-multiple-gpus>`__
+for details and
+`mjlab distributed training <https://github.com/mujocolab/mjlab/tree/main/docs/api/distributed_training.md>`__ for a
+reinforcement learning example.
+
 Orientation representation
 --------------------------