From ba72facf28a8b80ebadf595add6c40f298da6114 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 2 Jul 2026 03:58:50 +0000
Subject: [PATCH 1/3] Fix Newton initialization on non-default GPUs

Pin Torch and Warp to the simulation device before Newton allocates GPU
state, and bind CUDA graph capture to that same device. This prevents
nonzero distributed ranks from starting capture on Warp's default cuda:0
context.
---
 ...h-fix-newton-mgpu-capture-device-codex.rst |  5 +++
 .../isaaclab_newton/physics/newton_manager.py | 21 +++++++++++-
 .../test_newton_manager_abstraction.py        | 32 +++++++++++++++++++
 3 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst

diff --git a/source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst b/source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
new file mode 100644
index 000000000000..c53da690a8d8
--- /dev/null
+++ b/source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed Newton physics failing to initialize on non-default CUDA devices
+  (``cuda:1`` and higher).
diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
index 7d04c6e0ea5d..de8390b559db 100644
--- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
+++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
@@ -1107,6 +1107,15 @@ def start_simulation(cls) -> None:
         cls._cl_inject_sites_fallback()
 
         device = PhysicsManager._device
+        # Pin torch + Warp to the target device before any Warp/Newton
+        # allocations. Without this, mujoco_warp's collision pipeline later
+        # allocates on cuda:1 against a primary CUDA context that was never
+        # made current, returning null pointers (issue #5132).
+        if device and "cuda" in device:
+            import torch
+
+            torch.cuda.set_device(device)
+            wp.set_device(device)
         logger.info(f"Finalizing model on device: {device}")
         cls._builder.up_axis = Axis.from_string(cls._up_axis)
         # Forward pending extended attribute requests to builder and clear them
@@ -1394,6 +1403,16 @@ def initialize_solver(cls) -> None:
         if cfg is None:
             return
 
+        # Pin torch + Warp to the target device before solver build and
+        # collision-pipeline init.  Mirrors the guard in :meth:`start_simulation`
+        # (issue #5132); idempotent if already pinned.
+        device = PhysicsManager._device
+        if device and "cuda" in device:
+            import torch
+
+            torch.cuda.set_device(device)
+            wp.set_device(device)
+
         with Timer(name="newton_initialize_solver", msg="Initialize solver took:"):
             NewtonManager._num_substeps = cfg.num_substeps  # type: ignore[union-attr]
             NewtonManager._collision_decimation = cfg.collision_decimation  # type: ignore[union-attr]
@@ -1476,7 +1495,7 @@ def _capture_or_defer_graph(cls) -> None:
             with Timer(name="newton_cuda_graph", msg="CUDA graph took:"):
                 if cls._usdrt_stage is None:
                     simulate = cls._simulate_full if cls._is_all_graphable() else cls._simulate_physics_only
-                    with wp.ScopedCapture() as capture:
+                    with wp.ScopedCapture(device=device) as capture:
                         simulate()
                     NewtonManager._graph = capture.graph
                     logger.info("Newton CUDA graph captured (standard Warp mode)")
diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py
index 991a5f63da89..7adb20b58710 100644
--- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py
+++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py
@@ -426,6 +426,38 @@ def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch):
     assert NewtonManager._graph_capture_pending is False
 
 
+def test_cuda_graph_capture_uses_simulation_device(monkeypatch):
+    """CUDA graph capture should use the simulation device instead of Warp's default device."""
+    from isaaclab.physics import PhysicsManager
+
+    captured_devices = []
+    captured_graph = object()
+
+    class FakeScopedCapture:
+        def __init__(self, device=None):
+            captured_devices.append(device)
+            self.graph = captured_graph
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            return False
+
+    monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False)
+    monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_solver", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False))
+    monkeypatch.setattr(NewtonManager, "_simulate_physics_only", classmethod(lambda cls: None))
+    monkeypatch.setattr(wp, "ScopedCapture", FakeScopedCapture)
+
+    NewtonManager._capture_or_defer_graph()
+
+    assert captured_devices == ["cuda:1"]
+    assert NewtonManager._graph is captured_graph
+
+
 # ---------------------------------------------------------------------------
 # Manager class hierarchy and factory contracts
 # ---------------------------------------------------------------------------

From f20302ce2c27fb93b1f07edf56f6e654b611bd54 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 2 Jul 2026 05:11:57 +0000
Subject: [PATCH 2/3] Synchronize Torch and Warp CUDA devices

Select the resolved CUDA device in both runtimes at shared launcher and simulation initialization boundaries. Keep Newton graph capture explicitly bound to the simulation device without backend-specific default-device setup.
---
 ...h-fix-newton-mgpu-capture-device-codex.rst |  5 +++
 source/isaaclab/isaaclab/app/app_launcher.py  |  7 ++-
 source/isaaclab/isaaclab/app/sim_launcher.py  |  3 +-
 .../isaaclab/sim/simulation_context.py        |  7 +++
 source/isaaclab/isaaclab/utils/device.py      | 27 ++++++++++++
 source/isaaclab/test/utils/test_device.py     | 37 ++++++++++++++++
 .../isaaclab_newton/physics/newton_manager.py | 19 --------
 .../test_distributed_device_resolution.py     | 44 +++++++++----------
 8 files changed, 103 insertions(+), 46 deletions(-)
 create mode 100644 source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
 create mode 100644 source/isaaclab/isaaclab/utils/device.py
 create mode 100644 source/isaaclab/test/utils/test_device.py

diff --git a/source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst b/source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
new file mode 100644
index 000000000000..ba2c7b78ec0d
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed PyTorch and Warp selecting different CUDA devices during simulation
+  initialization.
diff --git a/source/isaaclab/isaaclab/app/app_launcher.py b/source/isaaclab/isaaclab/app/app_launcher.py
index 99215b696afd..0fee720a3c3f 100644
--- a/source/isaaclab/isaaclab/app/app_launcher.py
+++ b/source/isaaclab/isaaclab/app/app_launcher.py
@@ -29,6 +29,7 @@
     from isaacsim import SimulationApp
 
 from isaaclab.app.settings_manager import get_settings_manager, initialize_carb_settings
+from isaaclab.utils.device import set_cuda_device
 
 # import logger
 logger = logging.getLogger(__name__)
@@ -1073,13 +1074,11 @@ def _resolve_device_settings(self, launcher_args: dict):
         logger.info("Using device: %s", device)
 
     def _set_deferred_cuda_device(self) -> None:
-        """Set the current torch CUDA device after Kit startup."""
+        """Set the current CUDA device after Kit startup."""
         if self._deferred_cuda_device_id is None:
             return
 
-        import torch
-
-        torch.cuda.set_device(self._deferred_cuda_device_id)
+        set_cuda_device(self._deferred_cuda_device_id)
 
     def _resolve_experience_file(self, launcher_args: dict):
         """Resolve experience file related settings."""
diff --git a/source/isaaclab/isaaclab/app/sim_launcher.py b/source/isaaclab/isaaclab/app/sim_launcher.py
index 0a5ac08a1101..0aaf2fb04cc4 100644
--- a/source/isaaclab/isaaclab/app/sim_launcher.py
+++ b/source/isaaclab/isaaclab/app/sim_launcher.py
@@ -30,6 +30,7 @@
 from isaaclab.physics.physics_manager_cfg import PhysicsCfg
 from isaaclab.renderers.renderer_cfg import RendererCfg
 from isaaclab.sensors.camera.camera_cfg import CameraCfg
+from isaaclab.utils.device import set_cuda_device
 
 logger = logging.getLogger(__name__)
 
@@ -379,7 +380,7 @@ def _resolve_distributed_device(cfg, launcher_args: argparse.Namespace | dict |
     sim_cfg = getattr(cfg, "sim", None)
     if sim_cfg is not None:
         sim_cfg.device = device_str
-    torch.cuda.set_device(device_str)
+    set_cuda_device(device_str)
     logger.info(
         "Distributed device resolved to %s (local_rank=%d, visible_gpus=%d)",
         device_str,
diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py
index ef0a43d77ccc..78516ea0a3cd 100644
--- a/source/isaaclab/isaaclab/sim/simulation_context.py
+++ b/source/isaaclab/isaaclab/sim/simulation_context.py
@@ -29,6 +29,7 @@
 from isaaclab.scene_data import SceneDataProvider
 from isaaclab.sim.service_locator import ServiceLocator
 from isaaclab.sim.utils import create_new_stage
+from isaaclab.utils.device import set_cuda_device
 from isaaclab.utils.string import clear_resolve_matching_names_cache
 from isaaclab.utils.version import has_kit
 from isaaclab.visualizers.base_visualizer import BaseVisualizer
@@ -160,6 +161,12 @@ def __init__(self, cfg: SimulationCfg | None = None):
             device_id = max(0, int(cuda_device) if cuda_device is not None else 0)
             self.cfg.device = f"cuda:{device_id}"
 
+        # Synchronize the process-wide CUDA device before any physics backend
+        # initializes or allocates state. PyTorch must select the device before
+        # Warp so that both runtimes retain the same primary CUDA context.
+        if "cuda" in self.cfg.device:
+            set_cuda_device(self.cfg.device)
+
         # Set default physics backend if not specified
         if self.cfg.physics is None:
             from isaaclab_physx.physics import PhysxCfg
diff --git a/source/isaaclab/isaaclab/utils/device.py b/source/isaaclab/isaaclab/utils/device.py
new file mode 100644
index 000000000000..2ca0f85bd489
--- /dev/null
+++ b/source/isaaclab/isaaclab/utils/device.py
@@ -0,0 +1,27 @@
+# Copyright (c) 2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Utilities for selecting compute devices."""
+
+from __future__ import annotations
+
+
+def set_cuda_device(device: str | int) -> None:
+    """Set the process-wide CUDA device for both PyTorch and Warp.
+
+    PyTorch must select the device before Warp so that Warp initializes and
+    retains the primary CUDA context for the intended device.
+
+    Args:
+        device: CUDA device index or a device string such as ``"cuda:1"``.
+    """
+    import torch
+
+    torch.cuda.set_device(device)
+
+    import warp as wp
+
+    warp_device = f"cuda:{device}" if isinstance(device, int) else device
+    wp.set_device(warp_device)
diff --git a/source/isaaclab/test/utils/test_device.py b/source/isaaclab/test/utils/test_device.py
new file mode 100644
index 000000000000..457417e1d687
--- /dev/null
+++ b/source/isaaclab/test/utils/test_device.py
@@ -0,0 +1,37 @@
+# Copyright (c) 2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for compute-device selection utilities."""
+
+import sys
+from types import SimpleNamespace
+
+from isaaclab.utils.device import set_cuda_device
+
+
+def test_set_cuda_device_sets_torch_before_warp(monkeypatch):
+    """Setting a CUDA device must update PyTorch before Warp."""
+    calls = []
+    torch = SimpleNamespace(cuda=SimpleNamespace(set_device=lambda device: calls.append(("torch", device))))
+    warp = SimpleNamespace(set_device=lambda device: calls.append(("warp", device)))
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.setitem(sys.modules, "warp", warp)
+
+    set_cuda_device("cuda:2")
+
+    assert calls == [("torch", "cuda:2"), ("warp", "cuda:2")]
+
+
+def test_set_cuda_device_normalizes_integer_for_warp(monkeypatch):
+    """An integer device index must be converted to a Warp CUDA alias."""
+    calls = []
+    torch = SimpleNamespace(cuda=SimpleNamespace(set_device=lambda device: calls.append(("torch", device))))
+    warp = SimpleNamespace(set_device=lambda device: calls.append(("warp", device)))
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.setitem(sys.modules, "warp", warp)
+
+    set_cuda_device(3)
+
+    assert calls == [("torch", 3), ("warp", "cuda:3")]
diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
index de8390b559db..f4698375713e 100644
--- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
+++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
@@ -1107,15 +1107,6 @@ def start_simulation(cls) -> None:
         cls._cl_inject_sites_fallback()
 
         device = PhysicsManager._device
-        # Pin torch + Warp to the target device before any Warp/Newton
-        # allocations. Without this, mujoco_warp's collision pipeline later
-        # allocates on cuda:1 against a primary CUDA context that was never
-        # made current, returning null pointers (issue #5132).
-        if device and "cuda" in device:
-            import torch
-
-            torch.cuda.set_device(device)
-            wp.set_device(device)
         logger.info(f"Finalizing model on device: {device}")
         cls._builder.up_axis = Axis.from_string(cls._up_axis)
         # Forward pending extended attribute requests to builder and clear them
@@ -1403,16 +1394,6 @@ def initialize_solver(cls) -> None:
         if cfg is None:
             return
 
-        # Pin torch + Warp to the target device before solver build and
-        # collision-pipeline init.  Mirrors the guard in :meth:`start_simulation`
-        # (issue #5132); idempotent if already pinned.
-        device = PhysicsManager._device
-        if device and "cuda" in device:
-            import torch
-
-            torch.cuda.set_device(device)
-            wp.set_device(device)
-
         with Timer(name="newton_initialize_solver", msg="Initialize solver took:"):
             NewtonManager._num_substeps = cfg.num_substeps  # type: ignore[union-attr]
             NewtonManager._collision_decimation = cfg.collision_decimation  # type: ignore[union-attr]
diff --git a/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py b/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py
index 67268cb816c4..33de34289b59 100644
--- a/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py
+++ b/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py
@@ -15,8 +15,8 @@
 - Non-distributed: no device override applied
 - launch_simulation device propagation from AppLauncher
 
-No actual GPUs required — ``torch.cuda.device_count`` and
-``torch.cuda.set_device`` are mocked throughout.
+No actual GPUs required — ``torch.cuda.device_count`` and the shared CUDA
+device-selection helper are mocked throughout.
 """
 
 from __future__ import annotations
@@ -95,7 +95,7 @@ def _make_env_vars(
 class TestResolveDistributedDeviceNamespace:
     """Tests for _resolve_distributed_device with argparse.Namespace args."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_normal_multi_gpu_rank0(self, mock_count, mock_set_device):
         """4 visible GPUs, world_size=4, rank 0 → cuda:0."""
@@ -109,7 +109,7 @@ def test_normal_multi_gpu_rank0(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_normal_multi_gpu_rank3(self, mock_count, mock_set_device):
         """4 visible GPUs, world_size=4, rank 3 → cuda:3."""
@@ -123,7 +123,7 @@ def test_normal_multi_gpu_rank3(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:3"
         mock_set_device.assert_called_once_with("cuda:3")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_cuda_visible_devices_restricted_rank0(self, mock_count, mock_set_device):
         """1 visible GPU (CUDA_VISIBLE_DEVICES set), world_size=2, rank 0 → cuda:0."""
@@ -137,7 +137,7 @@ def test_cuda_visible_devices_restricted_rank0(self, mock_count, mock_set_device
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_cuda_visible_devices_restricted_rank1(self, mock_count, mock_set_device):
         """1 visible GPU, world_size=2, rank 1 → falls back to cuda:0 (not cuda:1)."""
@@ -151,7 +151,7 @@ def test_cuda_visible_devices_restricted_rank1(self, mock_count, mock_set_device
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=2)
     def test_jax_local_rank_added(self, mock_count, mock_set_device):
         """JAX_LOCAL_RANK is added to LOCAL_RANK for correct device mapping."""
@@ -166,7 +166,7 @@ def test_jax_local_rank_added(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:1"
         mock_set_device.assert_called_once_with("cuda:1")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_jax_local_rank_with_restricted_gpus(self, mock_count, mock_set_device):
         """JAX_LOCAL_RANK + restricted GPUs → fallback to cuda:0."""
@@ -189,7 +189,7 @@ def test_jax_local_rank_with_restricted_gpus(self, mock_count, mock_set_device):
 class TestResolveDistributedDeviceDict:
     """Tests for _resolve_distributed_device with dict-style args."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_dict_args_distributed(self, mock_count, mock_set_device):
         """Dict launcher_args with distributed=True should work identically."""
@@ -203,7 +203,7 @@ def test_dict_args_distributed(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:2"
         mock_set_device.assert_called_once_with("cuda:2")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_dict_args_restricted(self, mock_count, mock_set_device):
         """Dict args with restricted GPUs should fall back to cuda:0."""
@@ -226,7 +226,7 @@ def test_dict_args_restricted(self, mock_count, mock_set_device):
 class TestResolveDistributedDeviceNoop:
     """Tests that non-distributed runs skip device resolution."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_not_distributed_namespace(self, mock_set_device):
         """distributed=False → device unchanged, set_device not called."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -237,7 +237,7 @@ def test_not_distributed_namespace(self, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_not_called()
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_not_distributed_dict(self, mock_set_device):
         """Dict with distributed=False → no-op."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -248,7 +248,7 @@ def test_not_distributed_dict(self, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_not_called()
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_no_distributed_key(self, mock_set_device):
         """Dict without 'distributed' key → no-op."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -259,7 +259,7 @@ def test_no_distributed_key(self, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_not_called()
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_none_launcher_args(self, mock_set_device):
         """launcher_args=None → no-op."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -278,7 +278,7 @@ def test_none_launcher_args(self, mock_set_device):
 class TestResolveDistributedDeviceEdgeCases:
     """Edge cases for device resolution."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=2)
     def test_env_cfg_without_sim(self, mock_count, mock_set_device):
         """env_cfg with no 'sim' attribute → set_device still called, no crash."""
@@ -296,7 +296,7 @@ class _BareEnvCfg:
         # Should still call set_device even without sim_cfg
         mock_set_device.assert_called_once_with("cuda:1")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=2)
     def test_world_size_equals_visible_gpus(self, mock_count, mock_set_device):
         """Exact match: 2 visible GPUs, world_size=2 → use local_rank directly."""
@@ -309,7 +309,7 @@ def test_world_size_equals_visible_gpus(self, mock_count, mock_set_device):
 
         assert env_cfg.sim.device == "cuda:1"
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=8)
     def test_more_gpus_than_world_size(self, mock_count, mock_set_device):
         """8 visible GPUs but only 2 ranks → use local_rank directly."""
@@ -322,7 +322,7 @@ def test_more_gpus_than_world_size(self, mock_count, mock_set_device):
 
         assert env_cfg.sim.device == "cuda:1"
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=0)
     def test_zero_visible_gpus(self, mock_count, mock_set_device):
         """0 visible GPUs → fallback to cuda:0 (will fail later at CUDA init)."""
@@ -335,7 +335,7 @@ def test_zero_visible_gpus(self, mock_count, mock_set_device):
 
         assert env_cfg.sim.device == "cuda:0"
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_missing_env_vars_default_to_zero(self, mock_count, mock_set_device):
         """Missing LOCAL_RANK/WORLD_SIZE → defaults to 0/1."""
@@ -364,7 +364,7 @@ def test_missing_env_vars_default_to_zero(self, mock_count, mock_set_device):
 class TestResolveDistributedDeviceMultiNode:
     """Tests for multi-node setups where WORLD_SIZE > local GPU count."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_multi_node_rank3_sees_4_gpus(self, mock_count, mock_set_device):
         """2 nodes × 4 GPUs, WORLD_SIZE=8, local_rank=3, 4 visible → cuda:3.
@@ -382,7 +382,7 @@ def test_multi_node_rank3_sees_4_gpus(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:3"
         mock_set_device.assert_called_once_with("cuda:3")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_multi_node_rank0_sees_4_gpus(self, mock_count, mock_set_device):
         """2 nodes × 4 GPUs, WORLD_SIZE=8, local_rank=0 → cuda:0."""
@@ -396,7 +396,7 @@ def test_multi_node_rank0_sees_4_gpus(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_multi_node_restricted_gpus(self, mock_count, mock_set_device):
         """Multi-node with CUDA_VISIBLE_DEVICES=<one GPU per rank>, local_rank=1 → cuda:0."""

From 64b5ec8617878caeb074c7746a4738dae7165f13 Mon Sep 17 00:00:00 2001
From: jichuanh <jichuanh@nvidia.com>
Date: Thu, 2 Jul 2026 19:07:21 +0000
Subject: [PATCH 3/3] Harden CUDA device initialization

Synchronize Torch and Warp in the base physics manager before backend-specific setup. Keep the helper internal and cover the launcher and backend initialization boundaries directly.
---
 source/isaaclab/isaaclab/app/app_launcher.py  |  2 +-
 source/isaaclab/isaaclab/app/sim_launcher.py  |  2 +-
 .../isaaclab/physics/physics_manager.py       |  8 ++++
 .../isaaclab/sim/simulation_context.py        |  7 ---
 .../isaaclab/utils/{device.py => _device.py}  |  6 +--
 source/isaaclab/test/app/test_kwarg_launch.py | 12 +++++
 .../devices/test_physics_manager_device.py    | 46 +++++++++++++++++++
 source/isaaclab/test/utils/test_device.py     |  2 +-
 ...-fix-newton-mgpu-capture-device-codex.skip |  1 +
 9 files changed, 73 insertions(+), 13 deletions(-)
 rename source/isaaclab/isaaclab/utils/{device.py => _device.py} (74%)
 create mode 100644 source/isaaclab/test/devices/test_physics_manager_device.py
 create mode 100644 source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip

diff --git a/source/isaaclab/isaaclab/app/app_launcher.py b/source/isaaclab/isaaclab/app/app_launcher.py
index 0fee720a3c3f..cf9ea81c0509 100644
--- a/source/isaaclab/isaaclab/app/app_launcher.py
+++ b/source/isaaclab/isaaclab/app/app_launcher.py
@@ -29,7 +29,7 @@
     from isaacsim import SimulationApp
 
 from isaaclab.app.settings_manager import get_settings_manager, initialize_carb_settings
-from isaaclab.utils.device import set_cuda_device
+from isaaclab.utils._device import set_cuda_device
 
 # import logger
 logger = logging.getLogger(__name__)
diff --git a/source/isaaclab/isaaclab/app/sim_launcher.py b/source/isaaclab/isaaclab/app/sim_launcher.py
index 0aaf2fb04cc4..06d6dba408f5 100644
--- a/source/isaaclab/isaaclab/app/sim_launcher.py
+++ b/source/isaaclab/isaaclab/app/sim_launcher.py
@@ -30,7 +30,7 @@
 from isaaclab.physics.physics_manager_cfg import PhysicsCfg
 from isaaclab.renderers.renderer_cfg import RendererCfg
 from isaaclab.sensors.camera.camera_cfg import CameraCfg
-from isaaclab.utils.device import set_cuda_device
+from isaaclab.utils._device import set_cuda_device
 
 logger = logging.getLogger(__name__)
 
diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py
index 5255af873604..5c1a5d4403f1 100644
--- a/source/isaaclab/isaaclab/physics/physics_manager.py
+++ b/source/isaaclab/isaaclab/physics/physics_manager.py
@@ -14,6 +14,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, ClassVar
 
+from isaaclab.utils._device import set_cuda_device
+
 if TYPE_CHECKING:
     from isaaclab.scene_data import SceneDataBackend
     from isaaclab.sim.simulation_context import SimulationContext
@@ -271,6 +273,12 @@ def initialize(cls, sim_context: SimulationContext) -> None:
         PhysicsManager._device = sim_context.cfg.device
         PhysicsManager._sim_time = 0.0
 
+        # Synchronize the process-wide CUDA device before backend-specific
+        # initialization allocates state. PyTorch must select the device before
+        # Warp so that both runtimes retain the same primary CUDA context.
+        if "cuda" in PhysicsManager._device:
+            set_cuda_device(PhysicsManager._device)
+
     @classmethod
     @abstractmethod
     def reset(cls, soft: bool = False) -> None:
diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py
index 78516ea0a3cd..ef0a43d77ccc 100644
--- a/source/isaaclab/isaaclab/sim/simulation_context.py
+++ b/source/isaaclab/isaaclab/sim/simulation_context.py
@@ -29,7 +29,6 @@
 from isaaclab.scene_data import SceneDataProvider
 from isaaclab.sim.service_locator import ServiceLocator
 from isaaclab.sim.utils import create_new_stage
-from isaaclab.utils.device import set_cuda_device
 from isaaclab.utils.string import clear_resolve_matching_names_cache
 from isaaclab.utils.version import has_kit
 from isaaclab.visualizers.base_visualizer import BaseVisualizer
@@ -161,12 +160,6 @@ def __init__(self, cfg: SimulationCfg | None = None):
             device_id = max(0, int(cuda_device) if cuda_device is not None else 0)
             self.cfg.device = f"cuda:{device_id}"
 
-        # Synchronize the process-wide CUDA device before any physics backend
-        # initializes or allocates state. PyTorch must select the device before
-        # Warp so that both runtimes retain the same primary CUDA context.
-        if "cuda" in self.cfg.device:
-            set_cuda_device(self.cfg.device)
-
         # Set default physics backend if not specified
         if self.cfg.physics is None:
             from isaaclab_physx.physics import PhysxCfg
diff --git a/source/isaaclab/isaaclab/utils/device.py b/source/isaaclab/isaaclab/utils/_device.py
similarity index 74%
rename from source/isaaclab/isaaclab/utils/device.py
rename to source/isaaclab/isaaclab/utils/_device.py
index 2ca0f85bd489..020be59dfec2 100644
--- a/source/isaaclab/isaaclab/utils/device.py
+++ b/source/isaaclab/isaaclab/utils/_device.py
@@ -3,7 +3,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Utilities for selecting compute devices."""
+"""Internal utilities for selecting compute devices."""
 
 from __future__ import annotations
 
@@ -11,8 +11,8 @@
 def set_cuda_device(device: str | int) -> None:
     """Set the process-wide CUDA device for both PyTorch and Warp.
 
-    PyTorch must select the device before Warp so that Warp initializes and
-    retains the primary CUDA context for the intended device.
+    PyTorch must select the device before Warp so that a newly imported Warp
+    runtime initializes on that device, or an existing runtime switches to it.
 
     Args:
         device: CUDA device index or a device string such as ``"cuda:1"``.
diff --git a/source/isaaclab/test/app/test_kwarg_launch.py b/source/isaaclab/test/app/test_kwarg_launch.py
index 1eb5b59764ac..53e1326b2064 100644
--- a/source/isaaclab/test/app/test_kwarg_launch.py
+++ b/source/isaaclab/test/app/test_kwarg_launch.py
@@ -40,6 +40,18 @@ def test_livestream_rejects_disabled_visualizers():
         _ensure_livestream_kit_visualizer(args)
 
 
+def test_deferred_cuda_device_synchronizes_torch_and_warp(monkeypatch: pytest.MonkeyPatch):
+    """The post-Kit device hook must synchronize both CUDA runtimes."""
+    devices = []
+    monkeypatch.setattr(app_launcher_module, "set_cuda_device", devices.append)
+    launcher = AppLauncher.__new__(AppLauncher)
+    launcher._deferred_cuda_device_id = 2
+
+    launcher._set_deferred_cuda_device()
+
+    assert devices == [2]
+
+
 class _DummySettings:
     def __init__(self):
         self.values = {}
diff --git a/source/isaaclab/test/devices/test_physics_manager_device.py b/source/isaaclab/test/devices/test_physics_manager_device.py
new file mode 100644
index 000000000000..6b8face7f304
--- /dev/null
+++ b/source/isaaclab/test/devices/test_physics_manager_device.py
@@ -0,0 +1,46 @@
+# Copyright (c) 2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for physics-manager CUDA device initialization."""
+
+from types import SimpleNamespace
+
+from isaaclab.physics import PhysicsManager
+from isaaclab.physics import physics_manager as physics_manager_module
+
+
+def test_initialize_synchronizes_cuda_before_backend_setup(monkeypatch):
+    """The base manager must synchronize CUDA before backend-specific setup."""
+    calls = []
+    monkeypatch.setattr(physics_manager_module, "set_cuda_device", lambda device: calls.append(("cuda", device)))
+    monkeypatch.setattr(PhysicsManager, "_sim", None)
+    monkeypatch.setattr(PhysicsManager, "_cfg", None)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:0")
+
+    class _TestManager(PhysicsManager):
+        @classmethod
+        def initialize(cls, sim_context):
+            super().initialize(sim_context)
+            calls.append(("backend", PhysicsManager._device))
+
+    sim_context = SimpleNamespace(cfg=SimpleNamespace(physics=object(), device="cuda:2"))
+
+    _TestManager.initialize(sim_context)
+
+    assert calls == [("cuda", "cuda:2"), ("backend", "cuda:2")]
+
+
+def test_initialize_does_not_synchronize_cpu_device(monkeypatch):
+    """CPU simulation must not initialize or select CUDA runtimes."""
+    devices = []
+    monkeypatch.setattr(physics_manager_module, "set_cuda_device", devices.append)
+    monkeypatch.setattr(PhysicsManager, "_sim", None)
+    monkeypatch.setattr(PhysicsManager, "_cfg", None)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:0")
+    sim_context = SimpleNamespace(cfg=SimpleNamespace(physics=object(), device="cpu"))
+
+    PhysicsManager.initialize(sim_context)
+
+    assert devices == []
diff --git a/source/isaaclab/test/utils/test_device.py b/source/isaaclab/test/utils/test_device.py
index 457417e1d687..7bfeacec1339 100644
--- a/source/isaaclab/test/utils/test_device.py
+++ b/source/isaaclab/test/utils/test_device.py
@@ -8,7 +8,7 @@
 import sys
 from types import SimpleNamespace
 
-from isaaclab.utils.device import set_cuda_device
+from isaaclab.utils._device import set_cuda_device
 
 
 def test_set_cuda_device_sets_torch_before_warp(monkeypatch):
diff --git a/source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip b/source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip
new file mode 100644
index 000000000000..f03e45edd95d
--- /dev/null
+++ b/source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip
@@ -0,0 +1 @@
+Test-only updates for shared CUDA device-selection coverage.