Register native pinned host memory with the CANN runtime for NPU

accelerator/npu_accelerator.py

@@ -13,6 +13,33 @@
 except ImportError:
     pass
 
+ACL_SUCCESS = 0
+# The host-registration flag is an ACL_HOST_REG_* bitmask. MAPPED page-locks the
+# range and adds a device mapping (which DeepSpeed never reads through); it is
+# chosen over PINNED because measured H2D on this platform falls back to
+# mlock-only speed (~8-9 GB/s) for PINNED-only registrations of larger buffers
+# (64 MiB), while MAPPED keeps full DMA bandwidth (~23 GB/s). MAPPED requires
+# 4K-aligned addresses, which the native allocator guarantees (posix_memalign);
+# see the aclrtHostRegisterV2 API reference:
+# https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/latest/API/runtimeapi/aclcppdevg_03_2128.html
+ACL_HOST_REG_MAPPED = 0x2
+
+
+def _npu_host_copy_funcs():
+    """Resolve the npurt host-registration functions, or (None, reason).
+
+    torch.npu.npurt() returns the runtime-API module exposing
+    npuHostRegister/npuHostUnregister; the binding is maintained with
+    torch_npu and initializes the runtime itself.
+    """
+    if not hasattr(torch, "npu") or not hasattr(torch.npu, "npurt"):
+        return None, "torch.npu.npurt is unavailable in this torch_npu build"
+    try:
+        npurt = torch.npu.npurt()
+        return (npurt.npuHostRegister, npurt.npuHostUnregister), None
+    except RuntimeError:
+        return None, "torch.npu.npurt() failed to initialize the NPU runtime"
+
 
 class NPU_Accelerator(DeepSpeedAccelerator):
 
@@ -152,6 +179,50 @@ def total_memory(self, device_index=None):
     def available_memory(self, device_index=None):
         return self.total_memory(device_index) - self.memory_allocated(device_index)
 
+    # Host memory registration
+    def register_host_memory(self, address, num_bytes):
+        # Register natively pinned (posix_memalign + mlock) host memory with the
+        # ACL runtime so torch's async copies can use the DMA engine. npurt
+        # initializes the runtime itself, so no set_device ordering is needed.
+        # MAPPED registration requires 4K-aligned addresses (per the
+        # aclrtHostRegisterV2 API reference cited on ACL_HOST_REG_MAPPED). The
+        # native allocator always yields aligned addresses; if a caller passes
+        # an unaligned one, extend the range down to the page boundary so
+        # registration still succeeds instead of failing on the driver's opaque
+        # internal error. An already-aligned address passes through unchanged
+        # (offset 0), and unregister_host_memory rounds down identically.
+        offset = address % 4096
+        aligned_address = address - offset
+        # The pad keeps the registered range covering the original request.
+        padded_bytes = num_bytes + offset
+        funcs, reason = _npu_host_copy_funcs()
+        if funcs is None:
+            from deepspeed.utils import logger
+            logger.warning_once(f"Host-memory registration is unavailable ({reason}); "
+                                "native pinned memory stays mlock-only.")
+            return False
+        register, _ = funcs
+        rc = register(aligned_address, padded_bytes, ACL_HOST_REG_MAPPED)
+        if rc != ACL_SUCCESS:
+            from deepspeed.utils import logger
+            logger.warning_once(f"npuHostRegister failed with rc={rc}; native pinned memory stays mlock-only.")
+            return False
+        return True
+
+    def unregister_host_memory(self, address):
+        funcs, _ = _npu_host_copy_funcs()
+        if funcs is None:
+            return None
+        _, unregister = funcs
+        # Same page-boundary rounding as register_host_memory, so the driver
+        # releases exactly the range it was given.
+        offset = address % 4096
+        rc = unregister(address - offset)
+        if rc != ACL_SUCCESS:
+            # Raise so NativePinnedMemory keeps the allocation alive: the driver
+            # must never hold a registration for pages later reused by malloc.
+            raise RuntimeError(f"npuHostUnregister failed with rc={rc}")
+
     # Data types
     def is_bf16_supported(self):
         return torch.npu.is_bf16_supported()

benchmarks/pin_memory/h2d_d2h_bench.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
-"""Compare torch and native pinned-memory H2D/D2H bandwidth on one CUDA GPU."""
+"""Compare torch and native pinned-memory H2D/D2H bandwidth on one accelerator device."""
 
 import argparse
 import json
@@ -55,7 +55,9 @@ def _time_copy(accelerator, copy_fn, stream, warmup, iters):
 
 def _allocate_host(accelerator, numel, arm):
     if arm == "torch":
-        return torch.empty(numel, dtype=torch.float32, pin_memory=True)
+        # Go through the accelerator so every backend pins against its own
+        # device rather than relying on torch's ``pin_memory=True`` fast path.
+        return accelerator._torch_pin_memory(torch.empty(numel, dtype=torch.float32))
 
     return accelerator.pin_memory(torch.empty(numel, dtype=torch.float32), make_copy=False)
 
@@ -65,8 +67,8 @@ def _run_arm(args):
         os.environ[key] = value
 
     accelerator = get_accelerator()
-    if accelerator.device_name() != "cuda" or not accelerator.is_available():
-        raise RuntimeError("CUDA GPU is required")
+    if not accelerator.is_available():
+        raise RuntimeError(f"No {accelerator.device_name()} device is available")
     accelerator.set_device(0)
     stream = accelerator.Stream()
 
@@ -86,7 +88,7 @@ def _run_arm(args):
             "size_mib": size_mib,
             "h2d_gbps": num_bytes / h2d_seconds / 1e9,
             "d2h_gbps": num_bytes / d2h_seconds / 1e9,
-            "torch_is_pinned": host.is_pinned(),
+            "torch_is_pinned": accelerator._torch_is_pinned(host),
             "accelerator_is_pinned": accelerator.is_pinned(host),
         }
         print(f"RESULT={json.dumps(result, sort_keys=True)}", flush=True)

docs/code-docs/source/memory.rst

@@ -413,8 +413,10 @@ Native device registration
 
 Native allocations are device-independent ``mlock`` buffers. On CUDA systems,
 DeepSpeed additionally calls ``cudaHostRegister`` so PyTorch can use them for
-asynchronous H2D/D2H DMA. Device registration is enabled by default and can be
-disabled for comparison or debugging:
+asynchronous H2D/D2H DMA. On Ascend NPU systems, DeepSpeed additionally calls
+``aclrtHostRegisterV2`` / ``aclrtHostUnregister`` for the same purpose. Device
+registration is enabled by default and can be disabled for comparison or
+debugging:
 
 .. code-block:: bash
 

tests/unit/v1/pin_memory/test_pin_memory.py

@@ -5,8 +5,10 @@
 import pytest
 import torch
 
+from deepspeed.accelerator import npu_accelerator
 from deepspeed.accelerator.cpu_accelerator import CPU_Accelerator
 from deepspeed.accelerator.cuda_accelerator import CUDA_Accelerator
+from deepspeed.accelerator.npu_accelerator import NPU_Accelerator
 from deepspeed.utils.pin_memory import NativePinnedMemory
 
 
@@ -284,3 +286,115 @@ def unregister_host_memory(self, address):
     assert native_pins.unpin(pinned) is True
     assert accelerator.unregistered == [begin]
     assert begin not in native_pins._device_registered
+
+
+def test_npu_device_registration_calls_npurt(monkeypatch):
+    registered = []
+    unregistered = []
+
+    def register(address, num_bytes, flag):
+        registered.append((address, num_bytes, flag))
+        return 0
+
+    def unregister(address):
+        unregistered.append(address)
+        return 0
+
+    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None))
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+
+    # 4096: page-aligned addresses register without any range extension.
+    assert accelerator.register_host_memory(4096, 4096) is True
+    accelerator.unregister_host_memory(4096)
+    assert registered == [(4096, 4096, npu_accelerator.ACL_HOST_REG_MAPPED)]
+    assert unregistered == [4096]
+
+
+def test_npu_unaligned_address_is_extended_to_page_boundary(monkeypatch):
+    # MAPPED registration requires 4K-aligned addresses; unaligned addresses
+    # are extended down to the page boundary with a matching size pad so the
+    # registration still succeeds, and unregister rounds down identically.
+    registered = []
+    unregistered = []
+
+    def register(address, num_bytes, flag):
+        registered.append((address, num_bytes, flag))
+        return 0
+
+    def unregister(address):
+        unregistered.append(address)
+        return 0
+
+    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None))
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+
+    assert accelerator.register_host_memory(4096 + 1234, 4096) is True
+    assert registered == [(4096, 4096 + 1234, npu_accelerator.ACL_HOST_REG_MAPPED)]
+    accelerator.unregister_host_memory(4096 + 1234)
+    assert unregistered == [4096]
+
+
+def test_npu_device_registration_failure_returns_false(monkeypatch):
+    # A non-zero npuHostRegister return code must degrade to mlock-only, not
+    # raise: the NativePinnedMemory caller only tracks the address on True.
+    def register(address, num_bytes, flag):
+        return 107000
+
+    def unregister(address):
+        raise AssertionError("unregister must not run when register failed")
+
+    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None))
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+
+    assert accelerator.register_host_memory(4096, 4096) is False
+
+
+def test_npu_unregister_failure_raises(monkeypatch):
+    # Raising keeps the allocation alive in NativePinnedMemory so the driver
+    # never holds a registration for pages later reused by malloc.
+    def register(address, num_bytes, flag):
+        return 0
+
+    def unregister(address):
+        return 107000
+
+    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None))
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+
+    with pytest.raises(RuntimeError, match="npuHostUnregister"):
+        accelerator.unregister_host_memory(4096)
+
+
+def test_npu_missing_npurt_is_noop(monkeypatch):
+    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: (None, "test"))
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+
+    assert accelerator.register_host_memory(4096, 4096) is False
+    assert accelerator.unregister_host_memory(4096) is None
+
+
+def test_npu_host_copy_lookup_gates(monkeypatch):
+    # A build lacking npurt must not resolve, with a reason saying so.
+    monkeypatch.delattr(torch.npu, "npurt", raising=False)
+    funcs, reason = npu_accelerator._npu_host_copy_funcs()
+    assert funcs is None
+    assert "npurt is unavailable" in reason
+
+    # npurt() failing to initialize the runtime resolves to None with a reason.
+    def fail_npurt():
+        raise RuntimeError("init failed")
+
+    monkeypatch.setattr(torch.npu, "npurt", fail_npurt, raising=False)
+    funcs, reason = npu_accelerator._npu_host_copy_funcs()
+    assert funcs is None
+    assert "initialize" in reason
+
+    # A working npurt module resolves to its host copy functions.
+    class _Npurt:
+        npuHostRegister = "register"
+        npuHostUnregister = "unregister"
+
+    monkeypatch.setattr(torch.npu, "npurt", lambda: _Npurt(), raising=False)
+    funcs, reason = npu_accelerator._npu_host_copy_funcs()
+    assert funcs == ("register", "unregister")
+    assert reason is None