Support vLLM v0.28 accelerator memory APIs

README.md

@@ -289,11 +289,11 @@ with torch.accelerator.stream(torch.musa.Stream()):
     ...
 ```
 
-**Forward compatibility:** The wrapper always prefers the real
-`torch.accelerator` implementation and only falls back to `torch.musa` when an
-attribute is missing, so upgrading to a future PyTorch release that ships
-official implementations requires no changes on your side — you will
-automatically get the upstream version.
+**Forward compatibility:** The wrapper prefers the real `torch.accelerator`
+implementation and only falls back to `torch.musa` when an attribute is
+missing. The exception is torch_musa releases before `2.11.0.post2`, where
+known-broken accelerator memory APIs are forced through `torch.musa`. Starting
+with `2.11.0.post2`, the fixed official implementations are used automatically.
 
 ## Platform Detection
 
@@ -392,7 +392,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain.
 
 ```
 # pyproject.toml or requirements.txt
-torchada>=0.1.82
+torchada>=0.1.83
 ```
 
 ### Step 2: Conditional Import

README_CN.md

@@ -275,8 +275,9 @@ with torch.accelerator.stream(torch.musa.Stream()):
     ...
 ```
 
-**前向兼容性：** 包装器始终优先使用真正的 `torch.accelerator` 实现，只有在缺少属性时才回退到
-`torch.musa`，因此升级到提供官方实现的未来 PyTorch 版本时无需任何更改 —— 您将自动获得上游版本。
+**前向兼容性：** 包装器优先使用真正的 `torch.accelerator` 实现，只有在缺少属性时才回退到
+`torch.musa`。唯一例外是 `2.11.0.post2` 之前的 torch_musa：已知有问题的 accelerator 内存 API
+会被强制转发到 `torch.musa`。从 `2.11.0.post2` 开始，将自动使用已修复的官方实现。
 
 ## 平台检测
 
@@ -375,7 +376,7 @@ if torchada.is_gpu_device(device):  # 在 CUDA 和 MUSA 上都能工作
 
 ```
 # pyproject.toml 或 requirements.txt
-torchada>=0.1.82
+torchada>=0.1.83
 ```
 
 ### 步骤 2：条件导入

benchmarks/benchmark_history.json

@@ -3,7 +3,7 @@
   "description": "Historical benchmark results for torchada performance tracking",
   "results": [
     {
-      "version": "0.1.82",
+      "version": "0.1.83",
       "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.82"
+version = "0.1.83"
 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.82"
+__version__ = "0.1.83"
 
 from . import cuda, utils
 

src/torchada/_patch.py

@@ -331,9 +331,7 @@ def _patch_torch_device():
 
         _register_builtin(DeviceFactoryWrapper, "aten::device")
     except (AttributeError, ImportError, RuntimeError, TypeError):
-        logger.debug(
-            "Unable to register torch.device as a TorchScript builtin", exc_info=True
-        )
+        logger.debug("Unable to register torch.device as a TorchScript builtin", exc_info=True)
 
 
 # Store original torch.Generator for patching
@@ -619,9 +617,7 @@ def _wrap_jit_script(original_script: Callable) -> Callable:
 
     @functools.wraps(original_script)
     def wrapped_script(*args, **kwargs):
-        return _rewrite_scripted_object_device_constants(
-            original_script(*args, **kwargs)
-        )
+        return _rewrite_scripted_object_device_constants(original_script(*args, **kwargs))
 
     return wrapped_script
 
@@ -1725,6 +1721,48 @@ def __getitem__(self, name: str):
 _original_ctypes_CDLL = None
 
 
+_TORCH_MUSA_ACCELERATOR_FIX_VERSION = "2.11.0.post2"
+
+
+def _musa_accelerator_overrides_required(version) -> bool:
+    """Return whether torch.accelerator still needs MUSA memory overrides.
+
+    torch_musa 2.11.0.post2 fixes the unified accelerator memory APIs.  Older
+    releases still need torchada to force those calls through torch.musa.
+    Ignore the local version suffix (for example ``+musa5.2.0``), because it
+    identifies the MUSA stack build rather than the torch_musa fix level.
+
+    If a torch_musa build does not expose a version, retain the compatibility
+    overrides rather than risking the known runtime failure.
+    """
+    if version is None:
+        return True
+
+    public_version = str(version).split("+", 1)[0]
+    try:
+        # Use PyTorch's vendored PEP 440 parser so post releases compare
+        # semantically (post10 > post2) without adding a torchada dependency.
+        from torch._vendor.packaging.version import InvalidVersion, Version
+    except ImportError:
+        # If the parser is unavailable, keep the workaround enabled: disabling
+        # it could re-expose the allocator failure this gate fixes.
+        logger.warning(
+            "Unable to parse torch_musa version %r; retaining accelerator memory overrides",
+            version,
+        )
+        return True
+    try:
+        return Version(public_version) < Version(_TORCH_MUSA_ACCELERATOR_FIX_VERSION)
+    except InvalidVersion:
+        # An unknown or malformed version must keep the workaround enabled:
+        # disabling it could re-expose the allocator failure this gate fixes.
+        logger.warning(
+            "Unable to parse torch_musa version %r; retaining accelerator memory overrides",
+            version,
+        )
+        return True
+
+
 class _AcceleratorModuleWrapper(ModuleType):
     """
     Wrapper module that extends torch.accelerator with fallbacks to torch.musa.
@@ -1774,10 +1812,11 @@ class _AcceleratorModuleWrapper(ModuleType):
         "StreamContext": "core.stream.StreamContext",
     }
 
-    # Memory APIs that exist on torch.accelerator (PyTorch 2.9+) but internally
-    # call torch._C._accelerator_* C++ functions which fail on MUSA because the
-    # MUSA allocator is not a CUDA DeviceAllocator. These are overridden to
-    # delegate to torch.musa, following the same pattern as synchronize().
+    # Before torch_musa 2.11.0.post2, memory APIs that exist on
+    # torch.accelerator internally call torch._C._accelerator_* C++ functions
+    # which fail on MUSA because the MUSA allocator is not a CUDA
+    # DeviceAllocator. On those releases, delegate to torch.musa following the
+    # same pattern as synchronize().
     # When an API in this list exists on the original torch.accelerator AND on
     # torch.musa, we install an override that prefers torch.musa over the
     # upstream implementation.
@@ -1800,12 +1839,15 @@ def __init__(self, original_accel, musa_module):
         self._musa_module = musa_module
         self._overrides = {}
 
-        # Apply MUSA overrides for memory APIs that exist upstream but are
-        # broken on MUSA (they route through torch._C._accelerator_* which
-        # doesn't dispatch to the MUSA allocator).
-        for name in self._MUSA_OVERRIDES:
-            if hasattr(original_accel, name) and hasattr(musa_module, name):
-                self._set_override(name, getattr(musa_module, name))
+        # torch_musa versions before 2.11.0.post2 route these APIs through
+        # torch._C._accelerator_* without dispatching to the MUSA allocator.
+        # Newer versions provide working unified accelerator implementations,
+        # so preserve those instead of forcing the torch.musa compatibility path.
+        if _musa_accelerator_overrides_required(getattr(musa_module, "__version__", None)):
+            for name in self._MUSA_OVERRIDES:
+                musa_name = self._REMAP_ATTRS.get(name, name)
+                if hasattr(original_accel, name) and hasattr(musa_module, musa_name):
+                    self._set_override(name, getattr(musa_module, musa_name))
 
     def _set_override(self, name, value):
         """Install an override that takes precedence over the wrapped modules."""
@@ -1937,10 +1979,10 @@ def _patch_torch_accelerator():
        delegates to torch.musa.synchronize().
 
     2. Overrides for memory APIs that exist on torch.accelerator (PyTorch 2.9+)
-       but are broken on MUSA because they route through torch._C._accelerator_*
-       C++ functions that don't dispatch to the MUSA allocator. These are
-       redirected to torch.musa implementations (see _AcceleratorModuleWrapper
-       ._MUSA_OVERRIDES).
+       but are broken before torch_musa 2.11.0.post2 because they route through
+       torch._C._accelerator_* C++ functions that don't dispatch to the MUSA
+       allocator. On affected versions, these are redirected to torch.musa
+       implementations (see _AcceleratorModuleWrapper._MUSA_OVERRIDES).
 
     3. Forward compatibility for APIs that PyTorch is expected to add to
        torch.accelerator in future releases but are not yet present (Stream,
@@ -1950,12 +1992,6 @@ def _patch_torch_accelerator():
     4. device_index(idx) and stream(s) context managers, which are not yet
        present on torch.accelerator in torch 2.7.
 
-    TODO(torchada): README.md / README_CN.md claim "the wrapper always prefers
-    the real torch.accelerator implementation and only falls back to torch.musa
-    when an attribute is missing". That is no longer accurate after adding the
-    memory API overrides (point 2 above). Update those documents to describe
-    the actual resolution order: (1) torchada overrides, (2) real torch.accelerator,
-    (3) fallback to torch.musa.
     """
     global _original_torch_accelerator
 

src/torchada/utils/cpp_extension.py

@@ -1274,6 +1274,15 @@ def _dir_has_portable_sources(path):
 
                 @staticmethod
                 def _is_system_include_dir(path):
+                    # The stable-ABI compatibility headers are torchada-owned
+                    # inputs, not downstream project sources.  This matters for
+                    # editable/source installs where the package lives outside
+                    # site-packages: porting this directory in place rewrites
+                    # the CUDA-named forwarding shims into self-recursive MUSA
+                    # declarations and breaks every subsequent native build.
+                    if _path_overlaps_any(path, [stable_compat_include_dir()]):
+                        return True
+
                     is_system_path = (
                         path.startswith("/usr/")
                         or path.startswith("/opt/")

tests/test_cuda_patching.py

@@ -715,6 +715,7 @@ def test_stream_context_class(self):
     def test_streams_module(self):
         """Test the torch.cuda.streams module path used by PyTorch Dynamo."""
         import sys
+
         import torch
 
         import torchada
@@ -2887,7 +2888,35 @@ class TestAcceleratorModuleWrapper:
     official implementations of APIs that currently fall back to torch.musa).
     """
 
-    def _make_wrapper(self, accel_attrs=None, musa_attrs=None):
+    @pytest.mark.parametrize(
+        ("musa_version", "expected"),
+        (
+            ("2.10.0", True),
+            ("2.11.0", True),
+            ("2.11.0.post1+musa5.2.0", True),
+            ("2.11.0.post1+musa5.3.0", True),
+            ("2.11.0.post2", False),
+            ("2.11.0.post2+musa5.2.0", False),
+            ("2.11.0.post2+musa5.3.0", False),
+            ("2.11.0.post2+future/musa/build", False),
+            ("2.11.0.post10+musa5.2.0", False),
+            ("2.12.0+musa6.0.0", False),
+            ("not-a-version+musa5.2.0", True),
+            ("", True),
+            (None, True),
+        ),
+    )
+    def test_musa_accelerator_override_version_boundary(self, musa_version, expected):
+        from torchada._patch import _musa_accelerator_overrides_required
+
+        assert _musa_accelerator_overrides_required(musa_version) is expected
+
+    def _make_wrapper(
+        self,
+        accel_attrs=None,
+        musa_attrs=None,
+        musa_version="2.11.0.post1+musa5.2.0",
+    ):
         from types import ModuleType
 
         from torchada._patch import _AcceleratorModuleWrapper
@@ -2897,6 +2926,8 @@ def _make_wrapper(self, accel_attrs=None, musa_attrs=None):
             setattr(accel, k, v)
 
         musa = ModuleType("fake_torch_musa")
+        if musa_version is not None:
+            musa.__version__ = musa_version
         for k, v in (musa_attrs or {}).items():
             setattr(musa, k, v)
 
@@ -2917,11 +2948,12 @@ def test_original_accelerator_takes_precedence_over_musa(self):
         assert wrapper.manual_seed == "official_impl"
 
     def test_musa_overrides_take_precedence_when_both_exist(self):
-        """Memory APIs in _MUSA_OVERRIDES use torch.musa even when torch.accelerator has them.
+        """Affected torch_musa versions override official accelerator memory APIs.
 
         Starting in PyTorch 2.9+, torch.accelerator.empty_cache() exists but
         routes through torch._C._accelerator_* which doesn't work with the MUSA
-        allocator. The wrapper must override it to use torch.musa.empty_cache().
+        allocator before torch_musa 2.11.0.post2. The wrapper must override it
+        to use torch.musa.empty_cache() on those releases.
         """
         wrapper, _, _ = self._make_wrapper(
             accel_attrs={"empty_cache": "official_impl"},
@@ -2930,6 +2962,41 @@ def test_musa_overrides_take_precedence_when_both_exist(self):
         # empty_cache is in _MUSA_OVERRIDES, so torch.musa wins
         assert wrapper.empty_cache == "musa_fallback"
 
+    def test_remapped_musa_override_takes_precedence(self):
+        """Overrides must resolve torch.musa APIs through _REMAP_ATTRS.
+
+        PyTorch 2.11 exposes torch.accelerator.get_memory_info(), but its
+        implementation does not support the MUSA allocator. torch.musa keeps
+        the equivalent API under the older mem_get_info name, so the override
+        must use that remapped attribute instead of leaving the broken official
+        implementation in place.
+        """
+        musa_mem_get_info = object()
+        wrapper, _, _ = self._make_wrapper(
+            accel_attrs={"get_memory_info": "official_but_broken"},
+            musa_attrs={"mem_get_info": musa_mem_get_info},
+        )
+        assert wrapper.get_memory_info is musa_mem_get_info
+
+    def test_fixed_torch_musa_keeps_official_accelerator_api(self):
+        """torch_musa post2+ must use its fixed torch.accelerator implementation."""
+        wrapper, _, _ = self._make_wrapper(
+            accel_attrs={"get_memory_info": "official_fixed_impl"},
+            musa_attrs={"mem_get_info": "legacy_musa_impl"},
+            musa_version="2.11.0.post2+musa5.2.0",
+        )
+        assert wrapper.get_memory_info == "official_fixed_impl"
+        assert "get_memory_info" not in wrapper._overrides
+
+    def test_fixed_version_still_remaps_when_official_api_is_missing(self):
+        """The version gate must not disable the normal torch.musa fallback."""
+        musa_mem_get_info = object()
+        wrapper, _, _ = self._make_wrapper(
+            musa_attrs={"mem_get_info": musa_mem_get_info},
+            musa_version="2.11.0.post2+musa5.2.0",
+        )
+        assert wrapper.get_memory_info is musa_mem_get_info
+
     def test_fallback_to_musa_when_accelerator_missing(self):
         """Attributes absent from torch.accelerator must fall back to torch.musa."""
         wrapper, _, _ = self._make_wrapper(
@@ -3071,8 +3138,8 @@ def test_existing_accelerator_apis_preserved(self):
         # Function objects should come from the real torch.accelerator module
         assert torch.accelerator.is_available.__module__ == "torch.accelerator"
 
-    def test_empty_cache_falls_back_to_musa(self):
-        """torch.accelerator.empty_cache() must work via torch.musa fallback.
+    def test_empty_cache_uses_version_appropriate_implementation(self):
+        """torch.accelerator.empty_cache() must use the compatible implementation.
 
         Regression test for the user-reported AttributeError:
             >>> torch.accelerator.empty_cache()
@@ -3085,12 +3152,17 @@ def test_empty_cache_falls_back_to_musa(self):
         if not torchada.is_musa_platform():
             pytest.skip("Only applicable on MUSA platform")
 
-        # Must not raise
+        from torchada._patch import _musa_accelerator_overrides_required
+
+        # Must not raise on either side of the torch_musa post2 boundary.
         torch.accelerator.empty_cache()
-        assert torch.accelerator.empty_cache.__module__.startswith("torch_musa")
+        if _musa_accelerator_overrides_required(getattr(torch.musa, "__version__", None)):
+            assert torch.accelerator.empty_cache.__module__.startswith("torch_musa")
+        else:
+            assert torch.accelerator.empty_cache is torch.accelerator._original_accel.empty_cache
 
-    def test_memory_apis_fall_back_to_musa(self):
-        """Memory query APIs missing from torch.accelerator must work via fallback."""
+    def test_memory_apis_work(self):
+        """Memory APIs must work through either the compatibility or native path."""
         import torch
 
         import torchada
@@ -3103,6 +3175,10 @@ def test_memory_apis_fall_back_to_musa(self):
         assert isinstance(torch.accelerator.memory_reserved(), int)
         assert isinstance(torch.accelerator.max_memory_reserved(), int)
         assert isinstance(torch.accelerator.memory_stats(), dict)
+        memory_info = torch.accelerator.get_memory_info()
+        assert isinstance(memory_info, tuple)
+        assert len(memory_info) == 2
+        assert all(isinstance(value, int) for value in memory_info)
         torch.accelerator.reset_peak_memory_stats()
 
     def test_rng_apis_fall_back_to_musa(self):