[Newton] Fix multi-GPU initialization on non-default devices

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.

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__)
@@ -1081,13 +1082,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."""

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,

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:

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
+
+"""Internal 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 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"``.
+    """
+    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)

source/isaaclab/test/app/test_kwarg_launch.py

@@ -80,6 +80,18 @@ def __init__(self, _launcher_args):
     assert close_args == {"exit_code": 1}
 
 
+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 = {}

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 == []

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")]

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).

source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py

@@ -1536,7 +1536,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)")

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
 # ---------------------------------------------------------------------------

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.