fix(patch): keep *_like factories out of _device_constructors()

README.md

@@ -407,7 +407,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain.
 
 ```
 # pyproject.toml or requirements.txt
-torchada>=0.1.89
+torchada>=0.1.90
 ```
 
 ### Step 2: Conditional Import

README_CN.md

@@ -386,7 +386,7 @@ if torchada.is_gpu_device(device):  # 在 CUDA 和 MUSA 上都能工作
 
 ```
 # pyproject.toml 或 requirements.txt
-torchada>=0.1.89
+torchada>=0.1.90
 ```
 
 ### 步骤 2：条件导入

benchmarks/benchmark_history.json

@@ -3,7 +3,7 @@
   "description": "Historical benchmark results for torchada performance tracking",
   "results": [
     {
-      "version": "0.1.89",
+      "version": "0.1.90",
       "date": "2026-01-29",
       "platform": "MUSA",
       "pytorch_version": "2.7.1",

examples/extension_setup.py

@@ -69,7 +69,7 @@ def get_extensions():
     if extensions:
         setup(
             name="my_cuda_extension",
-            version="0.1.89",
+            version="0.1.90",
             ext_modules=extensions,
             cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
             python_requires=">=3.8",

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchada"
-version = "0.1.89"
+version = "0.1.90"
 description = "Adapter package for torch_musa to act exactly like PyTorch CUDA"
 readme = "README.md"
 license = {text = "MIT"}

src/torchada/__init__.py

@@ -24,7 +24,7 @@
     from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME
 """
 
-__version__ = "0.1.89"
+__version__ = "0.1.90"
 
 from . import cuda, utils
 

src/torchada/_patch.py

@@ -30,7 +30,7 @@
 import time
 import warnings
 from types import ModuleType, SimpleNamespace
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, List, Optional, Tuple
 
 import torch
 
@@ -1140,7 +1140,7 @@ def wrapped_script(*args, **kwargs):
 _EXTRA_FACTORY_FUNCTIONS = ("from_file", "normal")
 
 
-def _discover_factory_functions() -> List[str]:
+def _discover_factory_functions() -> Tuple[List[str], List[str]]:
     """Resolve, at runtime, the torch tensor-factory names whose ``device=``
     kwarg we translate — instead of hand-maintaining a static list.
 
@@ -1155,23 +1155,36 @@ def _discover_factory_functions() -> List[str]:
     Falls back to ``_FALLBACK_FACTORY_FUNCTIONS`` if the private registry is gone.
     """
     names = set()
+    # Subset of ``names`` that torch itself device-injects into, i.e. the
+    # original members of ``_device_constructors()``.  Only those may be added
+    # back to that registry: ``*_like`` factories derive their device from the
+    # input tensor, and registering them there makes ``torch.zeros_like(x)``
+    # return a *cpu* tensor as soon as any device context is active
+    # (``torch.set_default_device`` pushes one even for ``cpu``), which then
+    # feeds host pointers into device kernels.
+    device_injectable = set()
     try:
         from torch.utils._device import _device_constructors
 
         for fn in _device_constructors():
             name = getattr(fn, "__name__", None)
             if name and getattr(torch, name, None) is fn:
                 names.add(name)
+                device_injectable.add(name)
     except Exception:
         pass
     if not names:
         names.update(_FALLBACK_FACTORY_FUNCTIONS)
+        # torch does not device-inject into the ``*_like`` family.
+        device_injectable.update(
+            n for n in _FALLBACK_FACTORY_FUNCTIONS if not n.endswith("_like")
+        )
     # ``*_like`` variants accept device= but are not device-injected by torch.
     names |= {n + "_like" for n in tuple(names) if callable(getattr(torch, n + "_like", None))}
     for extra in _EXTRA_FACTORY_FUNCTIONS:
         if callable(getattr(torch, extra, None)):
             names.add(extra)
-    return sorted(names)
+    return sorted(names), sorted(device_injectable)
 
 
 # Salt mixed into the AOT-autograd cache key for the wrapped factories; bump to
@@ -2629,24 +2642,32 @@ def apply_patches():
     # keeps torch.compile AOT caching working.
     wrapped_names: List[str] = []
     original_fns = []
-    for fn_name in _discover_factory_functions():
+    device_injectable_fns = set()
+    factory_names, device_injectable_names = _discover_factory_functions()
+    for fn_name in factory_names:
         original_fn = getattr(torch, fn_name, None)
         if original_fn is None or hasattr(original_fn, "__wrapped__"):
             continue  # missing on this torch, or already wrapped
         original_fns.append(original_fn)
+        if fn_name in device_injectable_names:
+            device_injectable_fns.add(original_fn)
         wrapped_fn = _wrap_factory_function(original_fn)
         _register_jit_builtin_alias(original_fn, wrapped_fn)
         setattr(torch, fn_name, wrapped_fn)
         wrapped_names.append(fn_name)
 
     # PyTorch's __torch_function__ dispatch (e.g. the ``with torch.device(...):``
     # context manager) receives the original C functions, so the
-    # device-constructor set must include the unwrapped originals.
+    # device-constructor set must include the unwrapped originals -- but only the
+    # ones torch itself device-injects into.  Re-adding the ``*_like`` family
+    # silently overrode their device inheritance (see
+    # ``_discover_factory_functions``), turning device tensors into host tensors
+    # that sgl-kernel then dereferenced as device pointers.
     try:
         from torch.utils._device import _device_constructors
 
         constructors = _device_constructors()
-        for orig_fn in original_fns:
+        for orig_fn in device_injectable_fns:
             constructors.add(orig_fn)
     except (ImportError, AttributeError):
         pass

src/torchada/csrc/ops.h

@@ -40,7 +40,7 @@
 namespace torchada {
 
 // Version information
-constexpr const char* VERSION = "0.1.89";
+constexpr const char* VERSION = "0.1.90";
 
 // Check if operator override is enabled via environment variable
 inline bool is_override_enabled(const char* op_name) {

tests/test_device_strings.py

@@ -562,3 +562,105 @@ def test_original_functions_in_device_constructors(self):
                 assert torch.ones.__wrapped__ in constructors
         except ImportError:
             pytest.skip("torch.utils._device not available")
+
+    def test_like_factories_are_not_device_injected(self):
+        """The ``*_like`` family must stay out of ``_device_constructors()``.
+
+        torch never lists them there: ``*_like`` inherits its device from the
+        input tensor, so device-injecting them makes
+        ``torch.zeros_like(x)`` return a *cpu* tensor as soon as any device
+        context is active -- including the one ``torch.set_default_device(...)``
+        pushes even when the value does not change (it is not a no-op).  Those
+        host pointers then reach device kernels (illegal memory access).
+
+        Regression test: torchada used to re-add the unwrapped ``*_like``
+        factories to that registry, which silently broke device inheritance.
+        """
+        import torch
+
+        import torchada
+
+        if not torchada.is_musa_platform():
+            pytest.skip("MUSA platform required")
+
+        try:
+            from torch.utils._device import _device_constructors
+        except ImportError:
+            pytest.skip("torch.utils._device not available")
+
+        constructors = _device_constructors()
+        for name in (
+            "zeros_like",
+            "ones_like",
+            "empty_like",
+            "full_like",
+            "rand_like",
+            "randn_like",
+        ):
+            wrapped = getattr(torch, name, None)
+            if wrapped is None or not hasattr(wrapped, "__wrapped__"):
+                continue
+            assert wrapped.__wrapped__ not in constructors, (
+                f"torch.{name} (unwrapped) must not be in _device_constructors(): "
+                "torch does not device-inject the *_like family"
+            )
+
+        # torch *does* device-inject these, and they must keep working.
+        for name in ("empty", "zeros", "ones"):
+            wrapped = getattr(torch, name, None)
+            if wrapped is not None and hasattr(wrapped, "__wrapped__"):
+                assert wrapped.__wrapped__ in constructors
+
+    @pytest.mark.gpu
+    def test_like_factories_keep_input_device_in_device_context(self):
+        """Behavioural counterpart: ``with torch.device(cpu)`` must not move a
+        ``torch.zeros_like(<musa tensor>)`` onto cpu."""
+        import torch
+
+        import torchada
+
+        if not torchada.is_musa_platform():
+            pytest.skip("MUSA platform required")
+
+        src = torch.zeros(4, device="musa")
+        with torch.device("cpu"):
+            assert torch.zeros_like(src).device.type == "musa"
+            assert torch.ones_like(src).device.type == "musa"
+            assert torch.empty_like(src).device.type == "musa"
+            assert torch.full_like(src, 0).device.type == "musa"
+
+    @pytest.mark.gpu
+    def test_like_factories_still_translate_an_explicit_cuda_device(self):
+        """An explicit ``device="cuda"`` must still be translated to ``"musa"``.
+
+        Keeping the ``*_like`` family out of the device-*injection* registry
+        must not stop them from honouring a device the caller passed on
+        purpose: the wrapping is what makes CUDA-authored code such as
+        ``torch.ones_like(x, device="cuda")`` run on MUSA.  Guards both the
+        wrap list and the translation, inside and outside a device context.
+        """
+        import torch
+
+        import torchada
+
+        if not torchada.is_musa_platform():
+            pytest.skip("MUSA platform required")
+
+        src = torch.zeros(4, device="musa")
+        calls = (
+            ("zeros_like", (src,)),
+            ("ones_like", (src,)),
+            ("empty_like", (src,)),
+            ("full_like", (src, 1)),
+            ("rand_like", (src,)),
+            ("randn_like", (src,)),
+            ("randint_like", (src, 4)),
+        )
+
+        for name, args in calls:
+            factory = getattr(torch, name)
+            assert hasattr(factory, "__wrapped__"), f"torch.{name} must stay wrapped"
+            assert factory(*args, device="cuda").device.type == "musa"
+            # A device context must not override a device the caller passed.
+            with torch.device("cpu"):
+                assert factory(*args, device="cuda").device.type == "musa"

tests/test_platform.py

@@ -59,7 +59,7 @@ def test_get_version(self):
 
         version = torchada.get_version()
         assert version == torchada.__version__
-        assert version == "0.1.89"
+        assert version == "0.1.90"
         assert isinstance(version, str)
 
     def test_project_version_matches_runtime_version(self):