fix(musa): expose device index helper for TorchDynamo

src/torchada/_patch.py

@@ -775,6 +775,10 @@ class _CudaModuleWrapper(ModuleType):
     _SPECIAL_ATTRS = {
         "StreamContext": "core.stream.StreamContext",
         "streams": "core.stream",
+        # torch._dynamo models ``torch.cuda.device`` through this private CUDA
+        # helper.  torch_musa exposes the equivalent helper below ``core``
+        # instead of at the module top level.
+        "_get_device_index": "core._utils._get_musa_device_index",
     }
 
     # Attribute name remappings (CUDA name -> MUSA name)

tests/test_cuda_patching.py

@@ -727,6 +727,17 @@ def test_streams_module(self):
             assert sys.modules["torch.cuda.streams"] is torch_musa.core.stream
             assert torch.cuda.streams.Stream is torch_musa.core.stream.Stream
 
+    def test_get_device_index_helper(self):
+        """Test the private CUDA device helper used by PyTorch Dynamo."""
+        import torch
+
+        import torchada
+
+        if torchada.is_musa_platform():
+            from torch_musa.core._utils import _get_musa_device_index
+
+            assert torch.cuda._get_device_index is _get_musa_device_index
+
     def test_stream_cuda_stream_property(self):
         """Test stream.cuda_stream returns same value as stream.musa_stream on MUSA."""
         import torch