fix: protect dependency headers during extension builds

Author: fengchao.wei

README.md

@@ -201,9 +201,12 @@ torch_musa 2.9. The patched `torch.utils.cpp_extension.include_paths()` exposes
 the compatibility include directory on MUSA. Custom stable-ABI builds should
 add `stable_compat_include_dir()` explicitly, and kernels that use `TORCH_BOX`
 must force-include the path returned by `stable_compat_box_header()`. Both
-helpers are available from `torchada.utils.cpp_extension`. The torch_musa 2.9
-header backport runs lazily and best-effort at extension-build time; read-only
-headers are left unchanged. A plain `import torchada` does not modify PyTorch or
+helpers are available from `torchada.utils.cpp_extension`. Mark custom stable
+ABI extensions explicitly with `CUDAExtension(..., torchada_stable_abi=True)`.
+The torch_musa 2.9 header backport then runs lazily and best-effort only for
+stable-ABI extensions; read-only headers are left unchanged. Existing vLLM and
+SGLang extension names/source layouts are recognized for backward compatibility.
+Ordinary extension builds and a plain `import torchada` do not modify PyTorch or
 torch_musa headers.
 
 ### Custom Ops

README_CN.md

@@ -192,9 +192,11 @@ torchada 还为近期 vLLM 和 SGLang 在 torch_musa 2.9 上使用的 libtorch 
 `torch.utils.cpp_extension.include_paths()` 会返回该兼容 include 目录。自定义
 稳定 ABI 构建应显式加入 `stable_compat_include_dir()`；使用 `TORCH_BOX` 的内核
 还必须通过编译器强制 include `stable_compat_box_header()` 返回的头文件。这两个
-辅助函数都位于 `torchada.utils.cpp_extension`。torch_musa 2.9 头文件回补只会在
-扩展构建时延迟、尽力执行；只读头文件保持不变。单纯 `import torchada` 不会修改
-PyTorch 或 torch_musa 头文件。
+辅助函数都位于 `torchada.utils.cpp_extension`。自定义 stable ABI 扩展应通过
+`CUDAExtension(..., torchada_stable_abi=True)` 显式标记。torch_musa 2.9 头文件
+回补只会针对 stable ABI 扩展延迟、尽力执行；为保持兼容，现有 vLLM/SGLang 的
+扩展名称和源码目录结构也会被自动识别。普通扩展构建以及单纯
+`import torchada` 都不会修改 PyTorch 或 torch_musa 头文件。
 
 ### 自定义算子
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchada"
-version = "0.1.76"
+version = "0.1.77"
 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.76"
+__version__ = "0.1.77"
 
 from . import cuda, utils
 

src/torchada/utils/cpp_extension.py

@@ -26,6 +26,7 @@
 """
 
 import functools
+import importlib
 import inspect
 import logging
 import os
@@ -243,6 +244,34 @@ def _coalesce_port_roots(paths):
     return result
 
 
+def _runtime_dependency_roots() -> List[str]:
+    """Return package roots that may be used as extension include paths.
+
+    ``torch_musa`` is often installed in editable mode during development. In
+    that case ``torch_musa.__file__`` resolves into the checkout, even though
+    the directory is a dependency of the extension being built. Such a path
+    must remain available to the compiler, but it must never become an
+    in-place CUDA-to-MUSA porting root.
+    """
+    roots = []
+    for module_name in ("torch", "torch_musa"):
+        try:
+            module = importlib.import_module(module_name)
+            module_file = getattr(module, "__file__", None)
+            if module_file:
+                roots.append(os.path.realpath(os.path.dirname(module_file)))
+        except (ImportError, AttributeError, OSError):
+            continue
+    return roots
+
+
+def _path_overlaps_any(path: str, roots: List[str]) -> bool:
+    """Return whether ``path`` contains or is contained by a protected root."""
+    return any(
+        _path_is_within(path, root) or _path_is_within(root, path) for root in roots
+    )
+
+
 def _validate_portable_symlinks(source_dir: str) -> None:
     """Reject portable symlink files before upstream ``realpath`` can escape.
 
@@ -926,6 +955,38 @@ def library_paths(cuda: Optional[bool] = None, device_type: Optional[str] = None
             return torch_library_paths(cuda=include_device)
 
 
+def _extension_requires_stable_header_backport(extension) -> bool:
+    """Return whether an extension needs the torch_musa stable-ABI backport.
+
+    Callers can set ``torchada_stable_abi=True`` on ``CUDAExtension`` for an
+    explicit decision. The name/source fallback preserves compatibility with
+    existing vLLM and SGLang extensions that predate that keyword.
+    """
+    explicit = getattr(extension, "_torchada_stable_abi", None)
+    if explicit is not None:
+        return explicit
+
+    markers = ("libtorch_stable", "stable_libtorch", "stable_abi")
+    candidates = [getattr(extension, "name", "")]
+    candidates.extend(getattr(extension, "sources", None) or [])
+    if any(any(marker in str(value).lower() for marker in markers) for value in candidates):
+        return True
+
+    compile_args = getattr(extension, "extra_compile_args", None) or []
+    if isinstance(compile_args, dict):
+        compile_args = [arg for args in compile_args.values() for arg in args]
+    elif isinstance(compile_args, str):
+        compile_args = [compile_args]
+    if any("torchada_stable_box.h" in str(arg) for arg in compile_args):
+        return True
+
+    return any(
+        macro[0] == "TORCHADA_STABLE_ACCESSORS"
+        for macro in (getattr(extension, "define_macros", None) or [])
+        if macro
+    )
+
+
 class CUDAExtension:
     """
     A wrapper that creates either a torch CUDAExtension or MUSA MUSAExtension.
@@ -944,12 +1005,21 @@ def __new__(cls, name: str, sources: List[str], *args, **kwargs):
             *args: Additional positional arguments
             **kwargs: Additional keyword arguments
         """
-        platform = detect_platform()
+        stable_abi = kwargs.pop("torchada_stable_abi", None)
+        if stable_abi is not None and not isinstance(stable_abi, bool):
+            raise TypeError("torchada_stable_abi must be a bool")
 
+        platform = detect_platform()
         if platform == Platform.MUSA:
-            return _create_musa_extension(name, sources, *args, **kwargs)
+            extension = _create_musa_extension(name, sources, *args, **kwargs)
         else:
-            return _create_cuda_extension(name, sources, *args, **kwargs)
+            extension = _create_cuda_extension(name, sources, *args, **kwargs)
+
+        if stable_abi is not None:
+            extension._torchada_stable_abi = stable_abi
+        if platform == Platform.MUSA and _extension_requires_stable_header_backport(extension):
+            _ensure_stable_headers_patched()
+        return extension
 
 
 class CppExtension:
@@ -1036,9 +1106,6 @@ def _create_musa_extension(name: str, sources: List[str], *args, **kwargs):
     """
     # Ensure patches are applied
     _apply_musa_patches()
-    # Building a MUSA extension: apply the (on-disk) libtorch-stable header
-    # backport now so csrc/libtorch_stable/*.cu can compile.
-    _ensure_stable_headers_patched()
 
     # Translate CUDA compiler, library, and feature-macro names to MUSA.
     kwargs = _translate_compile_args(kwargs)
@@ -1131,9 +1198,14 @@ def get_mapping_rule(self):
                     return _MAPPING_RULE.copy()
 
                 def build_extensions(self):
-                    # Building now: apply the (on-disk) libtorch-stable header
-                    # backport before the compiler reads the torch headers.
-                    _ensure_stable_headers_patched()
+                    # Only stable-ABI extensions need the legacy torch_musa
+                    # header backport. Ordinary extensions (for example,
+                    # torchvision) must not mutate dependency headers.
+                    if any(
+                        _extension_requires_stable_header_backport(extension)
+                        for extension in self.extensions
+                    ):
+                        _ensure_stable_headers_patched()
                     # Register .cu, .cuh as valid source extensions
                     self.compiler.src_extensions += [".cu", ".cuh"]
                     super().build_extensions()
@@ -1180,12 +1252,20 @@ def _dir_has_portable_sources(path):
 
                 @staticmethod
                 def _is_system_include_dir(path):
-                    return (
+                    is_system_path = (
                         path.startswith("/usr/")
                         or path.startswith("/opt/")
                         or "site-packages" in path
                         or "dist-packages" in path
                     )
+                    if is_system_path:
+                        return True
+
+                    # Include roots belonging to the installed PyTorch or
+                    # torch_musa dependency are compiler inputs, not project
+                    # sources. This also covers editable installs whose real
+                    # path is a development checkout such as /home/torch_musa.
+                    return _path_overlaps_any(path, _runtime_dependency_roots())
 
                 def run(self):
                     """Port each project-local include root's CUDA sources to MUSA

tests/test_cpp_extension.py

@@ -145,6 +145,108 @@ def test_porting_translates_torch_cuda_header(self):
         assert _replace_porting_line(project_include, rules) == project_include
 
 
+class TestStableHeaderBackportSelection:
+    """Stable header mutation is limited to extensions that opt in or match."""
+
+    @staticmethod
+    def _requires(extension):
+        from torchada.utils.cpp_extension import (
+            _extension_requires_stable_header_backport,
+        )
+
+        return _extension_requires_stable_header_backport(extension)
+
+    def test_ordinary_extension_does_not_require_backport(self):
+        extension = SimpleNamespace(
+            name="torchvision_image",
+            sources=["torchvision/csrc/io/image/decode_jpegs_cuda.cpp"],
+            extra_compile_args={"cxx": ["-O3"], "mcc": ["-O3"]},
+            define_macros=[("NVJPEG_FOUND", 1)],
+        )
+
+        assert self._requires(extension) is False
+
+    def test_explicit_stable_abi_marker_takes_precedence(self):
+        extension = SimpleNamespace(
+            name="ordinary_name",
+            sources=["kernel.cu"],
+            _torchada_stable_abi=True,
+        )
+        assert self._requires(extension) is True
+
+        extension._torchada_stable_abi = False
+        extension.name = "_C_stable_libtorch"
+        assert self._requires(extension) is False
+
+    def test_existing_stable_extension_names_and_sources_are_detected(self):
+        by_name = SimpleNamespace(name="_C_stable_libtorch", sources=[])
+        by_source = SimpleNamespace(
+            name="vllm_extension",
+            sources=["vllm/csrc/libtorch_stable/gemm_kernels.cu"],
+        )
+
+        assert self._requires(by_name) is True
+        assert self._requires(by_source) is True
+
+    def test_stable_box_force_include_is_detected(self):
+        extension = SimpleNamespace(
+            name="custom_ops",
+            sources=["ops.cu"],
+            extra_compile_args={
+                "cxx": ["-include", "/tmp/torchada_stable_box.h"],
+            },
+        )
+
+        assert self._requires(extension) is True
+
+    def test_cuda_extension_only_patches_headers_for_stable_abi(self, monkeypatch):
+        from torchada.utils import cpp_extension
+        from torchada.utils.cpp_extension import Platform
+
+        patched = []
+        monkeypatch.setattr(cpp_extension, "detect_platform", lambda: Platform.MUSA)
+        monkeypatch.setattr(
+            cpp_extension,
+            "_create_musa_extension",
+            lambda name, sources, *args, **kwargs: SimpleNamespace(
+                name=name,
+                sources=sources,
+            ),
+        )
+        monkeypatch.setattr(
+            cpp_extension,
+            "_ensure_stable_headers_patched",
+            lambda: patched.append(True),
+        )
+
+        cpp_extension.CUDAExtension("torchvision_image", ["decode_jpegs_cuda.cpp"])
+        assert patched == []
+
+        extension = cpp_extension.CUDAExtension(
+            "custom_stable_ops",
+            ["ops.cu"],
+            torchada_stable_abi=True,
+        )
+        assert extension._torchada_stable_abi is True
+        assert patched == [True]
+
+
+class TestDependencyIncludeProtection:
+    """Editable dependency paths are never treated as project porting roots."""
+
+    def test_dependency_package_and_checkout_paths_overlap(self):
+        from torchada.utils.cpp_extension import _path_overlaps_any
+
+        dependency_package = "/home/torch_musa/torch_musa"
+
+        assert _path_overlaps_any("/home/torch_musa", [dependency_package])
+        assert _path_overlaps_any(
+            "/home/torch_musa/torch_musa/share/generated_cuda_compatible",
+            [dependency_package],
+        )
+        assert not _path_overlaps_any("/home/torchvision", [dependency_package])
+
+
 class TestMusaPatches:
     """Test patches applied to torch_musa for extension building."""
 

tests/test_inplace_porting.py

@@ -244,3 +244,46 @@ def get_mapping_rule(self):
 
     assert source.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n"
     assert command._ported_dirs == {os.path.realpath(source_dir)}
+
+
+def test_editable_dependency_include_root_is_not_ported(tmp_path, monkeypatch):
+    """An editable torch_musa checkout stays an include path, not a port root."""
+    from torchada.utils import cpp_extension
+    from torchada.utils.cpp_extension import BuildExtension
+
+    project_dir = tmp_path / "vision"
+    dependency_dir = tmp_path / "torch_musa"
+    project_dir.mkdir()
+    dependency_package = dependency_dir / "torch_musa"
+    dependency_package.mkdir(parents=True)
+
+    source = project_dir / "kernel.cu"
+    dependency_header = dependency_package / "dependency.h"
+    source.write_text(f"{TOKEN}\n", encoding="utf-8")
+    dependency_header.write_text(f"{TOKEN}\n", encoding="utf-8")
+
+    monkeypatch.setattr(
+        cpp_extension,
+        "_runtime_dependency_roots",
+        lambda: [str(dependency_package)],
+    )
+
+    class CustomBuildExtension(BuildExtension):
+        def get_mapping_rule(self):
+            return MAPPING
+
+    command = CustomBuildExtension(Distribution())
+    command.extensions = [
+        Extension(
+            "test_editable_dependency",
+            sources=[str(source)],
+            include_dirs=[str(dependency_dir)],
+        )
+    ]
+    monkeypatch.setattr(BuildExtension.__mro__[1], "run", lambda self: None)
+
+    command.run()
+
+    assert source.read_text(encoding="utf-8") == f"{PORTED_TOKEN}\n"
+    assert dependency_header.read_text(encoding="utf-8") == f"{TOKEN}\n"
+    assert command._ported_dirs == {os.path.realpath(project_dir)}