fix: generalize stable ABI extension porting for MUSA

README.md

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

README_CN.md

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

benchmarks/benchmark_history.json

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

pyproject.toml

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

src/torchada/_mappings/libtorch_stable.py

@@ -3,5 +3,7 @@
 MAPPING = {
     'aoti_torch_get_current_cuda_stream': 'aoti_torch_get_current_musa_stream',
     'torch_get_current_cuda_blas_handle': 'torch_get_current_musa_blas_handle',
-    'STABLE_TORCH_LIBRARY_IMPL(_C, CUDA': 'STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1',
+    'torch_set_current_cuda_stream': 'torch_set_current_musa_stream',
+    'torch_get_cuda_stream_from_pool': 'torch_get_musa_stream_from_pool',
+    'torch_cuda_stream_synchronize': 'torch_musa_stream_synchronize',
 }

src/torchada/utils/cpp_extension.py

@@ -377,6 +377,20 @@ def _narrow_cuda_header_mapping(mapping_rule):
 
 _INCLUDE_DIRECTIVE_RE = re.compile(r'^\s*#\s*include\s*[<"](?P<header>[^>"]+)[>"]')
 _NVJPEG_PREFIX_RULES = frozenset(("nvjpeg", "NVJPEG"))
+_STABLE_LIBRARY_IMPL_CUDA_RE = re.compile(
+    r"(\bSTABLE_TORCH_LIBRARY_IMPL\s*\(\s*[A-Za-z_]\w*\s*,\s*)CUDA(\s*,)"
+)
+
+
+def _rekey_stable_library_impl(source_code):
+    """Register stable-ABI CUDA implementations on MUSA's dispatch key.
+
+    The dispatch key is a bare macro argument, so mappings such as
+    ``torch::kCUDA -> torch::kPrivateUse1`` cannot match it.  Handle every
+    valid library namespace and whitespace layout, including multiline macro
+    invocations, instead of maintaining project-specific literal mappings.
+    """
+    return _STABLE_LIBRARY_IMPL_CUDA_RE.sub(r"\1PrivateUse1\2", source_code)
 
 
 def _replace_porting_line(line, mapping_rule):
@@ -487,7 +501,8 @@ def modify_file(self, cuda_filepath, musa_filepath):
                 f"Refusing to port CUDA/C++ source outside the in-place root: {cuda_filepath}"
             )
         with open(cuda_filepath, encoding="utf-8", errors="surrogateescape") as f:
-            lines = f.readlines()
+            source_code = _rekey_stable_library_impl(f.read())
+            lines = source_code.splitlines(keepends=True)
 
         def port_line(line):
             if line.startswith("*") or line.startswith("/") or line == "":
@@ -813,6 +828,8 @@ def _port_cuda_source(source_code: str, mapping_rules: Optional[Dict[str, str]]
     if mapping_rules is None:
         mapping_rules = _MAPPING_RULE
 
+    source_code = _rekey_stable_library_impl(source_code)
+
     # Sort rules by length (longest first) to avoid partial replacements
     sorted_rules = sorted(mapping_rules.items(), key=lambda x: len(x[0]), reverse=True)
     return "".join(

tests/test_inplace_porting.py

@@ -23,6 +23,33 @@ def _build_extension_command():
     return BuildExtension(Distribution())
 
 
+def test_stable_library_impl_cuda_key_is_rekeyed_for_any_namespace(tmp_path):
+    source_dir = tmp_path / "csrc"
+    source_dir.mkdir()
+    source = source_dir / "registration.cpp"
+    source.write_text(
+        "STABLE_TORCH_LIBRARY_IMPL(custom_ops, CUDA, m) {}\n"
+        "STABLE_TORCH_LIBRARY_IMPL(\n"
+        "    third_party_ops,\n"
+        "    CUDA,\n"
+        "    m) {}\n"
+        "STABLE_TORCH_LIBRARY_IMPL(custom_ops, CPU, m) {}\n",
+        encoding="utf-8",
+    )
+
+    command = _build_extension_command()
+    command._port_directory(str(source_dir), {})
+
+    assert source.read_text(encoding="utf-8") == (
+        "STABLE_TORCH_LIBRARY_IMPL(custom_ops, PrivateUse1, m) {}\n"
+        "STABLE_TORCH_LIBRARY_IMPL(\n"
+        "    third_party_ops,\n"
+        "    PrivateUse1,\n"
+        "    m) {}\n"
+        "STABLE_TORCH_LIBRARY_IMPL(custom_ops, CPU, m) {}\n"
+    )
+
+
 def test_existing_mirror_sibling_is_preserved(tmp_path):
     source_dir = tmp_path / "csrc"
     mirror_dir = tmp_path / "csrc_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.79"
+        assert version == "0.1.80"
         assert isinstance(version, str)
 
     def test_project_version_matches_runtime_version(self):

tests/test_stable_abi_mappings.py

@@ -28,15 +28,18 @@ def test_stable_blas_handle_symbol_rule(self):
             == "torch_get_current_musa_blas_handle"
         )
 
-    def test_stable_impl_dispatch_key_rekey(self):
-        """STABLE_TORCH_LIBRARY_IMPL registers under the literal dispatch-key
-        token; MUSA tensors are PrivateUse1, so the block must be re-keyed."""
+    @pytest.mark.parametrize(
+        ("cuda_symbol", "musa_symbol"),
+        [
+            ("torch_set_current_cuda_stream", "torch_set_current_musa_stream"),
+            ("torch_get_cuda_stream_from_pool", "torch_get_musa_stream_from_pool"),
+            ("torch_cuda_stream_synchronize", "torch_musa_stream_synchronize"),
+        ],
+    )
+    def test_stable_stream_symbol_rules(self, cuda_symbol, musa_symbol):
         from torchada._mapping import _MAPPING_RULE
 
-        assert (
-            _MAPPING_RULE["STABLE_TORCH_LIBRARY_IMPL(_C, CUDA"]
-            == "STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1"
-        )
+        assert _MAPPING_RULE[cuda_symbol] == musa_symbol
 
     def test_cuda_header_rules_present(self):
         from torchada._mapping import _MAPPING_RULE
@@ -65,6 +68,17 @@ def test_port_rewrites_stable_blas_handle(self):
         assert "torch_get_current_musa_blas_handle" in ported
         assert "torch_get_current_cuda_blas_handle" not in ported
 
+    def test_port_rewrites_stable_stream_symbols(self):
+        ported = self._port(
+            "torch_set_current_cuda_stream(stream, device);\n"
+            "torch_get_cuda_stream_from_pool(false, device, &stream);\n"
+            "torch_cuda_stream_synchronize(stream, device);\n"
+        )
+        assert "torch_set_current_musa_stream" in ported
+        assert "torch_get_musa_stream_from_pool" in ported
+        assert "torch_musa_stream_synchronize" in ported
+        assert "cuda_stream" not in ported
+
     def test_ported_blas_handle_has_torch_29_fallback(self):
         """The symbol emitted by porting must exist on the torch 2.9 path."""
         from torchada.utils.cpp_extension import stable_compat_box_header
@@ -82,10 +96,25 @@ def test_ported_blas_handle_has_torch_29_fallback(self):
         assert f"static inline AOTITorchError {symbol}(void** ret)" in torch_29_branch
         assert f"return {symbol}(ret);" in torch_29_branch
 
-    def test_port_rekeys_stable_impl_block(self):
+    @pytest.mark.parametrize("namespace", ["_C", "custom_ops", "third_party_ops"])
+    def test_port_rekeys_stable_impl_block_for_any_namespace(self, namespace):
         ported = self._port(
-            "STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl(\"x\", f); }")
-        assert "STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1" in ported
+            f"STABLE_TORCH_LIBRARY_IMPL({namespace}, CUDA, ops) {{ ops.impl(\"x\", f); }}")
+        assert f"STABLE_TORCH_LIBRARY_IMPL({namespace}, PrivateUse1" in ported
+
+    def test_port_rekeys_multiline_stable_impl_block(self):
+        ported = self._port(
+            "STABLE_TORCH_LIBRARY_IMPL(\n"
+            "    third_party_ops,\n"
+            "    CUDA,\n"
+            "    ops) { ops.impl(\"x\", f); }"
+        )
+        assert "third_party_ops,\n    PrivateUse1," in ported
+
+    @pytest.mark.parametrize("dispatch_key", ["CPU", "CompositeExplicitAutograd", "PrivateUse1"])
+    def test_port_preserves_non_cuda_stable_impl_key(self, dispatch_key):
+        source = f"STABLE_TORCH_LIBRARY_IMPL(custom_ops, {dispatch_key}, ops) {{}}"
+        assert self._port(source) == source
 
 
 class TestStableCompatHeadersShipped: