From a6eb83fe07f9224eb7db4c306e27c0d5db7dc0a3 Mon Sep 17 00:00:00 2001 From: popsiclexu Date: Thu, 24 Sep 2026 17:17:46 +0800 Subject: [PATCH 1/2] fix(patch): keep *_like factories out of _device_constructors() --- src/torchada/_patch.py | 33 +++++++++--- tests/test_device_strings.py | 102 +++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 6 deletions(-) diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index a919c4e..f906200 100644 --- a/src/torchada/_patch.py +++ b/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,6 +1155,14 @@ 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 @@ -1162,16 +1170,21 @@ def _discover_factory_functions() -> List[str]: 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,11 +2642,15 @@ 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) @@ -2641,12 +2658,16 @@ def apply_patches(): # 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 diff --git a/tests/test_device_strings.py b/tests/test_device_strings.py index bff0463..4177cb9 100644 --- a/tests/test_device_strings.py +++ b/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()`` 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" From feb197705e263a8fea45b54c9cea7a00fc19fe7d Mon Sep 17 00:00:00 2001 From: popsiclexu Date: Thu, 24 Sep 2026 19:21:42 +0800 Subject: [PATCH 2/2] fix: update version to 0.1.90 --- README.md | 2 +- README_CN.md | 2 +- benchmarks/benchmark_history.json | 2 +- examples/extension_setup.py | 2 +- pyproject.toml | 2 +- src/torchada/__init__.py | 2 +- src/torchada/csrc/ops.h | 2 +- tests/test_platform.py | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index dca8a68..ba6a2e2 100644 --- a/README.md +++ b/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 diff --git a/README_CN.md b/README_CN.md index 2607378..21cfdb4 100644 --- a/README_CN.md +++ b/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：条件导入 diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index 2228eb7..34d7123 100644 --- a/benchmarks/benchmark_history.json +++ b/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", diff --git a/examples/extension_setup.py b/examples/extension_setup.py index 3b173af..a1821ba 100755 --- a/examples/extension_setup.py +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index eb8cc70..b1db576 100644 --- a/pyproject.toml +++ b/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"} diff --git a/src/torchada/__init__.py b/src/torchada/__init__.py index 09c1c16..b9890e6 100644 --- a/src/torchada/__init__.py +++ b/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 diff --git a/src/torchada/csrc/ops.h b/src/torchada/csrc/ops.h index 9b2669d..5f7154e 100644 --- a/src/torchada/csrc/ops.h +++ b/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) { diff --git a/tests/test_platform.py b/tests/test_platform.py index 52de101..30ff24c 100644 --- a/tests/test_platform.py +++ b/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):