From efc68451f5080465e4b59c779011674890a3be6c Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Wed, 16 Sep 2026 02:32:18 +0000
Subject: [PATCH 1/8] Register native pinned host memory with the CANN runtime
 for NPU

Implement register_host_memory / unregister_host_memory on NPU_Accelerator
via torch.npu.npurt() (npuHostRegister / npuHostUnregister over
aclrtHostRegisterV2 / aclrtHostUnregister), gated on torch_npu >= 2.9.0
with mlock-only fallbacks. Extend the pin-memory tests with NPU cases,
make the H2D/D2H benchmark accelerator-agnostic, and document the NPU
registration path.

Fixes #8531

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 accelerator/npu_accelerator.py              |  80 +++++++++++++
 benchmarks/pin_memory/h2d_d2h_bench.py      |  12 +-
 docs/code-docs/source/memory.rst            |   6 +-
 tests/unit/v1/pin_memory/test_pin_memory.py | 117 ++++++++++++++++++++
 4 files changed, 208 insertions(+), 7 deletions(-)

diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py
index fbcc10172ac8..d10e3c3e0615 100644
--- a/accelerator/npu_accelerator.py
+++ b/accelerator/npu_accelerator.py
@@ -13,6 +13,49 @@
 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).
+ACL_HOST_REG_MAPPED = 0x2
+_MIN_TORCH_NPU_REGISTER_VERSION = (2, 9)
+
+
+def _torch_npu_version():
+    """Return the installed torch_npu version as (major, minor), or None."""
+    try:
+        import torch_npu
+        major, minor, *_ = torch_npu.__version__.split("+")[0].split(".")
+        return (int(major), int(minor))
+    except Exception:
+        return None
+
+
+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. Host registration is
+    gated on torch_npu >= 2.9.0, the minimum supported version.
+    """
+    version = _torch_npu_version()
+    if version is None:
+        return None, "unable to determine the installed torch_npu version"
+    if version < _MIN_TORCH_NPU_REGISTER_VERSION:
+        return None, (f"torch_npu {version[0]}.{version[1]} is older than the minimum supported version "
+                      f"{_MIN_TORCH_NPU_REGISTER_VERSION[0]}.{_MIN_TORCH_NPU_REGISTER_VERSION[1]}")
+    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 +195,43 @@ 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.
+        if address % 4096:
+            # MAPPED registration requires 4K-aligned addresses and the driver
+            # only reports an opaque internal error otherwise; fail fast here.
+            from deepspeed.utils import logger
+            logger.warning_once(
+                f"Native pinned buffer is not 4K-aligned (address={address:#x}); skipping host-memory registration.")
+            return False
+        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(address, num_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
+        rc = unregister(address)
+        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()
diff --git a/benchmarks/pin_memory/h2d_d2h_bench.py b/benchmarks/pin_memory/h2d_d2h_bench.py
index 0cad917cdbb9..8211e351ea41 100644
--- a/benchmarks/pin_memory/h2d_d2h_bench.py
+++ b/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)
diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst
index 1a52b04cc3f2..d8d24865f775 100644
--- a/docs/code-docs/source/memory.rst
+++ b/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
 
diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index d30f38a37706..ceacfeb96856 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/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,118 @@ 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: the hook requires 4K-aligned addresses for MAPPED registration.
+    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_skipped(monkeypatch):
+    # MAPPED registration requires 4K-aligned addresses and the driver reports
+    # only an opaque internal error; the hook must fail fast without resolving
+    # or calling the npurt functions.
+    def fail_lookup():
+        raise AssertionError("npurt must not be resolved for unaligned addresses")
+
+    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", fail_lookup)
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+
+    assert accelerator.register_host_memory(1234, 4096) is False
+
+
+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):
+    # Below the minimum supported torch_npu version the functions must not
+    # resolve, with a reason pointing at the version gap.
+    monkeypatch.setattr(npu_accelerator, "_torch_npu_version", lambda: (2, 8))
+    funcs, reason = npu_accelerator._npu_host_copy_funcs()
+    assert funcs is None
+    assert "older than" in reason
+
+    # A version that cannot be parsed must fail closed with an explanation.
+    monkeypatch.setattr(npu_accelerator, "_torch_npu_version", lambda: None)
+    funcs, reason = npu_accelerator._npu_host_copy_funcs()
+    assert funcs is None
+    assert "unable to determine" in reason
+
+    # A supported version whose build lacks npurt must not resolve either.
+    monkeypatch.setattr(npu_accelerator, "_torch_npu_version", lambda: (2, 9))
+    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

From b7e6784d614501616788fdf9d870893cdf03e434 Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Wed, 16 Sep 2026 07:27:08 +0000
Subject: [PATCH 2/8] Gate NPU host registration on npurt availability instead
 of torch_npu version

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 accelerator/npu_accelerator.py              | 20 +-------------------
 tests/unit/v1/pin_memory/test_pin_memory.py | 16 +---------------
 2 files changed, 2 insertions(+), 34 deletions(-)

diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py
index d10e3c3e0615..a5baaea691bc 100644
--- a/accelerator/npu_accelerator.py
+++ b/accelerator/npu_accelerator.py
@@ -21,17 +21,6 @@
 # (64 MiB), while MAPPED keeps full DMA bandwidth (~23 GB/s). MAPPED requires
 # 4K-aligned addresses, which the native allocator guarantees (posix_memalign).
 ACL_HOST_REG_MAPPED = 0x2
-_MIN_TORCH_NPU_REGISTER_VERSION = (2, 9)
-
-
-def _torch_npu_version():
-    """Return the installed torch_npu version as (major, minor), or None."""
-    try:
-        import torch_npu
-        major, minor, *_ = torch_npu.__version__.split("+")[0].split(".")
-        return (int(major), int(minor))
-    except Exception:
-        return None
 
 
 def _npu_host_copy_funcs():
@@ -39,15 +28,8 @@ def _npu_host_copy_funcs():
 
     torch.npu.npurt() returns the runtime-API module exposing
     npuHostRegister/npuHostUnregister; the binding is maintained with
-    torch_npu and initializes the runtime itself. Host registration is
-    gated on torch_npu >= 2.9.0, the minimum supported version.
+    torch_npu and initializes the runtime itself.
     """
-    version = _torch_npu_version()
-    if version is None:
-        return None, "unable to determine the installed torch_npu version"
-    if version < _MIN_TORCH_NPU_REGISTER_VERSION:
-        return None, (f"torch_npu {version[0]}.{version[1]} is older than the minimum supported version "
-                      f"{_MIN_TORCH_NPU_REGISTER_VERSION[0]}.{_MIN_TORCH_NPU_REGISTER_VERSION[1]}")
     if not hasattr(torch, "npu") or not hasattr(torch.npu, "npurt"):
         return None, "torch.npu.npurt is unavailable in this torch_npu build"
     try:
diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index ceacfeb96856..b231ffe55fb7 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/tests/unit/v1/pin_memory/test_pin_memory.py
@@ -363,21 +363,7 @@ def test_npu_missing_npurt_is_noop(monkeypatch):
 
 
 def test_npu_host_copy_lookup_gates(monkeypatch):
-    # Below the minimum supported torch_npu version the functions must not
-    # resolve, with a reason pointing at the version gap.
-    monkeypatch.setattr(npu_accelerator, "_torch_npu_version", lambda: (2, 8))
-    funcs, reason = npu_accelerator._npu_host_copy_funcs()
-    assert funcs is None
-    assert "older than" in reason
-
-    # A version that cannot be parsed must fail closed with an explanation.
-    monkeypatch.setattr(npu_accelerator, "_torch_npu_version", lambda: None)
-    funcs, reason = npu_accelerator._npu_host_copy_funcs()
-    assert funcs is None
-    assert "unable to determine" in reason
-
-    # A supported version whose build lacks npurt must not resolve either.
-    monkeypatch.setattr(npu_accelerator, "_torch_npu_version", lambda: (2, 9))
+    # 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

From bb2e5967cbc152db9fcce9f4886c864085311ce4 Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Wed, 16 Sep 2026 08:50:27 +0000
Subject: [PATCH 3/8] Extend unaligned NPU host registrations to the 4K page
 boundary

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 accelerator/npu_accelerator.py              | 29 ++++++++++++++-------
 tests/unit/v1/pin_memory/test_pin_memory.py | 29 ++++++++++++++-------
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py
index a5baaea691bc..af48300465d8 100644
--- a/accelerator/npu_accelerator.py
+++ b/accelerator/npu_accelerator.py
@@ -19,7 +19,9 @@
 # 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).
+# 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
 
 
@@ -182,13 +184,17 @@ 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.
-        if address % 4096:
-            # MAPPED registration requires 4K-aligned addresses and the driver
-            # only reports an opaque internal error otherwise; fail fast here.
-            from deepspeed.utils import logger
-            logger.warning_once(
-                f"Native pinned buffer is not 4K-aligned (address={address:#x}); skipping host-memory registration.")
-            return False
+        # 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
@@ -196,7 +202,7 @@ def register_host_memory(self, address, num_bytes):
                                 "native pinned memory stays mlock-only.")
             return False
         register, _ = funcs
-        rc = register(address, num_bytes, ACL_HOST_REG_MAPPED)
+        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.")
@@ -208,7 +214,10 @@ def unregister_host_memory(self, address):
         if funcs is None:
             return None
         _, unregister = funcs
-        rc = unregister(address)
+        # 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.
diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index b231ffe55fb7..029ed2d97191 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/tests/unit/v1/pin_memory/test_pin_memory.py
@@ -303,24 +303,35 @@ def unregister(address):
     monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None))
     accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
 
-    # 4096: the hook requires 4K-aligned addresses for MAPPED registration.
+    # 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_skipped(monkeypatch):
-    # MAPPED registration requires 4K-aligned addresses and the driver reports
-    # only an opaque internal error; the hook must fail fast without resolving
-    # or calling the npurt functions.
-    def fail_lookup():
-        raise AssertionError("npurt must not be resolved for unaligned addresses")
+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 = []
 
-    monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", fail_lookup)
+    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(1234, 4096) is False
+    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):

From 8ba87d196e57cee6ed57f3c2ba4c54f52f2f4381 Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Wed, 16 Sep 2026 09:42:32 +0000
Subject: [PATCH 4/8] Share page-boundary alignment between NPU host memory
 register/unregister hooks

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 accelerator/npu_accelerator.py              | 22 +++++++--
 tests/unit/v1/pin_memory/test_pin_memory.py | 53 ++++++++-------------
 2 files changed, 36 insertions(+), 39 deletions(-)

diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py
index af48300465d8..cae0a27fde1c 100644
--- a/accelerator/npu_accelerator.py
+++ b/accelerator/npu_accelerator.py
@@ -41,6 +41,18 @@ def _npu_host_copy_funcs():
         return None, "torch.npu.npurt() failed to initialize the NPU runtime"
 
 
+def _align_to_page_boundary(address):
+    """Round an address down to the 4K page boundary, returning (aligned address, offset).
+
+    MAPPED registration requires 4K-aligned addresses (per the
+    aclrtHostRegisterV2 API reference cited on ACL_HOST_REG_MAPPED). The
+    returned offset lets callers pad the size so the aligned range still
+    covers the original request.
+    """
+    offset = address % 4096
+    return address - offset, offset
+
+
 class NPU_Accelerator(DeepSpeedAccelerator):
 
     def __init__(self):
@@ -190,9 +202,9 @@ def register_host_memory(self, address, num_bytes):
         # 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
+        # (offset 0), and unregister_host_memory rounds down via the same
+        # helper.
+        aligned_address, offset = _align_to_page_boundary(address)
         # The pad keeps the registered range covering the original request.
         padded_bytes = num_bytes + offset
         funcs, reason = _npu_host_copy_funcs()
@@ -216,8 +228,8 @@ def unregister_host_memory(self, address):
         _, 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)
+        aligned_address, _ = _align_to_page_boundary(address)
+        rc = unregister(aligned_address)
         if rc != ACL_SUCCESS:
             # Raise so NativePinnedMemory keeps the allocation alive: the driver
             # must never hold a registration for pages later reused by malloc.
diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index 029ed2d97191..e38b322c5c97 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/tests/unit/v1/pin_memory/test_pin_memory.py
@@ -288,50 +288,35 @@ def unregister_host_memory(self, address):
     assert begin not in native_pins._device_registered
 
 
-def test_npu_device_registration_calls_npurt(monkeypatch):
+@pytest.mark.parametrize(
+    "address, expected_address, expected_bytes",
+    [
+        (4096, 4096, 4096),  # page-aligned: registers unchanged
+        (4096 + 1234, 4096, 4096 + 1234),  # unaligned: extended down with a matching size pad
+    ])
+def test_npu_register_aligns_to_page_boundary(monkeypatch, address, expected_address, expected_bytes):
+    # MAPPED registration requires 4K-aligned addresses; the hook rounds the
+    # address down to the page boundary and pads the size so the registered
+    # range still covers the original request. An already-aligned address
+    # passes through unchanged, 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)
-
-    # 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))
+    def register(addr, num_bytes, flag):
+        registered.append((addr, num_bytes, flag))
         return 0
 
-    def unregister(address):
-        unregistered.append(address)
+    def unregister(addr):
+        unregistered.append(addr)
         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]
+    assert accelerator.register_host_memory(address, 4096) is True
+    assert registered == [(expected_address, expected_bytes, npu_accelerator.ACL_HOST_REG_MAPPED)]
+    accelerator.unregister_host_memory(address)
+    assert unregistered == [expected_address]
 
 
 def test_npu_device_registration_failure_returns_false(monkeypatch):

From 91d7f015a9c62d1d76c5400186b55c98328479b3 Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Wed, 16 Sep 2026 10:04:35 +0000
Subject: [PATCH 5/8] Guard NPU host copy lookup test against builds without
 torch.npu

torch.npu only exists in torch_npu builds, so test_npu_host_copy_lookup_gates
dereferenced a missing attribute on the regular CPU/CUDA unit-test suite.
Install a stand-in torch.npu for the test so the gate logic is exercised
everywhere.

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 tests/unit/v1/pin_memory/test_pin_memory.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index e38b322c5c97..7c6de7cca239 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/tests/unit/v1/pin_memory/test_pin_memory.py
@@ -359,6 +359,12 @@ def test_npu_missing_npurt_is_noop(monkeypatch):
 
 
 def test_npu_host_copy_lookup_gates(monkeypatch):
+    # torch.npu only exists in torch_npu builds; install a stand-in so this
+    # pure-Python gate logic also runs where the NPU backend is absent.
+    class _StubNpu:
+        pass
+
+    monkeypatch.setattr(torch, "npu", _StubNpu(), raising=False)
     # 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()

From ff9e91b921e6f5176ae515bd6da144ba20e8b278 Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Fri, 18 Sep 2026 06:34:42 +0000
Subject: [PATCH 6/8] Expose pin memory alignment on the accelerator interface

The 4K page-boundary rounding moves out of the NPU accelerator into NativePinnedMemory, driven by a new pin_memory_alignment() query on the abstract interface (default 1 = no requirement; NPU declares 4096). The NPU-specific alignment test becomes a generic one parametrized over the declared alignment, exercised through the public pin/unpin flow.

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 accelerator/abstract_accelerator.py         |  4 +
 accelerator/npu_accelerator.py              | 37 +++------
 deepspeed/utils/pin_memory.py               | 22 +++++-
 tests/unit/v1/pin_memory/test_pin_memory.py | 85 +++++++++++++++++----
 4 files changed, 103 insertions(+), 45 deletions(-)

diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py
index 96a671e7f3c8..9f33b4f75bbc 100644
--- a/accelerator/abstract_accelerator.py
+++ b/accelerator/abstract_accelerator.py
@@ -286,6 +286,10 @@ def unregister_host_memory(self, address):
         """Unregister host memory previously registered with the device runtime."""
         return None
 
+    def pin_memory_alignment(self):
+        """Byte alignment required for device-registered host memory; 1 means none."""
+        return 1
+
     # CPU torch pinning is a historical no-op; subclasses that really page-lock
     # keep the default True so tracker accounting matches the docs.
     _torch_pins_host_memory = True
diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py
index cae0a27fde1c..2b9abfab9507 100644
--- a/accelerator/npu_accelerator.py
+++ b/accelerator/npu_accelerator.py
@@ -41,18 +41,6 @@ def _npu_host_copy_funcs():
         return None, "torch.npu.npurt() failed to initialize the NPU runtime"
 
 
-def _align_to_page_boundary(address):
-    """Round an address down to the 4K page boundary, returning (aligned address, offset).
-
-    MAPPED registration requires 4K-aligned addresses (per the
-    aclrtHostRegisterV2 API reference cited on ACL_HOST_REG_MAPPED). The
-    returned offset lets callers pad the size so the aligned range still
-    covers the original request.
-    """
-    offset = address % 4096
-    return address - offset, offset
-
-
 class NPU_Accelerator(DeepSpeedAccelerator):
 
     def __init__(self):
@@ -192,21 +180,17 @@ def available_memory(self, device_index=None):
         return self.total_memory(device_index) - self.memory_allocated(device_index)
 
     # Host memory registration
+    def pin_memory_alignment(self):
+        # MAPPED registration requires 4K-aligned addresses (per the
+        # aclrtHostRegisterV2 API reference cited on ACL_HOST_REG_MAPPED).
+        # NativePinnedMemory rounds the range down to this alignment before
+        # calling register_host_memory/unregister_host_memory.
+        return 4096
+
     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 via the same
-        # helper.
-        aligned_address, offset = _align_to_page_boundary(address)
-        # 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
@@ -214,7 +198,7 @@ def register_host_memory(self, address, num_bytes):
                                 "native pinned memory stays mlock-only.")
             return False
         register, _ = funcs
-        rc = register(aligned_address, padded_bytes, ACL_HOST_REG_MAPPED)
+        rc = register(address, num_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.")
@@ -226,10 +210,7 @@ def unregister_host_memory(self, address):
         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.
-        aligned_address, _ = _align_to_page_boundary(address)
-        rc = unregister(aligned_address)
+        rc = unregister(address)
         if rc != ACL_SUCCESS:
             # Raise so NativePinnedMemory keeps the allocation alive: the driver
             # must never hold a registration for pages later reused by malloc.
diff --git a/deepspeed/utils/pin_memory.py b/deepspeed/utils/pin_memory.py
index fa992df5c965..5099f6d43587 100644
--- a/deepspeed/utils/pin_memory.py
+++ b/deepspeed/utils/pin_memory.py
@@ -64,7 +64,13 @@ def _new_locked(self, example, out_shape):
         if base.nbytes and self._device_registration_enabled():
             from deepspeed.accelerator import get_accelerator
             try:
-                if get_accelerator().register_host_memory(begin, base.nbytes):
+                accelerator = get_accelerator()
+                # Some device runtimes only register page-aligned ranges; round
+                # the address down to the accelerator-declared alignment and pad
+                # the size so the registered range still covers the request.
+                registered_begin = self._align_host_address(accelerator, begin)
+                padded_bytes = base.nbytes + (begin - registered_begin)
+                if accelerator.register_host_memory(registered_begin, padded_bytes):
                     self._device_registered.add(begin)
             except Exception as e:
                 logger.warning_once(
@@ -131,12 +137,24 @@ def _release(handle, begin, ranges, finalizers, device_registered):
             # during interpreter shutdown.
             pass
 
+    @staticmethod
+    def _align_host_address(accelerator, address):
+        # Device runtimes may require page-aligned registration ranges; round
+        # down to the alignment the accelerator declares (1 = no requirement).
+        alignment = accelerator.pin_memory_alignment()
+        if alignment <= 1:
+            return address
+        return address - (address % alignment)
+
     @staticmethod
     def _unregister_device(begin, device_registered):
         if begin not in device_registered:
             return
         from deepspeed.accelerator import get_accelerator
-        get_accelerator().unregister_host_memory(begin)
+        accelerator = get_accelerator()
+        # Same rounding as at registration time, so the driver releases exactly
+        # the range it was given.
+        accelerator.unregister_host_memory(NativePinnedMemory._align_host_address(accelerator, begin))
         device_registered.discard(begin)
 
     @staticmethod
diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index 7c6de7cca239..3eb9b5f4b017 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/tests/unit/v1/pin_memory/test_pin_memory.py
@@ -142,6 +142,9 @@ def register_host_memory(self, address, num_bytes):
     def unregister_host_memory(self, address):
         self.unregistered.append(address)
 
+    def pin_memory_alignment(self):
+        return 1
+
 
 def test_native_device_registration_and_unpin(monkeypatch, native_pins):
     accelerator = _RegisteringAccelerator()
@@ -227,6 +230,9 @@ def register_host_memory(self, address, num_bytes):
         def unregister_host_memory(self, address):
             raise AssertionError("unregister must not run when register failed")
 
+        def pin_memory_alignment(self):
+            return 1
+
     monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: _FailingAccelerator())
     monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1")
     pinned = native_pins.pin(torch.empty(32), make_copy=False)
@@ -288,17 +294,66 @@ def unregister_host_memory(self, address):
     assert begin not in native_pins._device_registered
 
 
-@pytest.mark.parametrize(
-    "address, expected_address, expected_bytes",
-    [
-        (4096, 4096, 4096),  # page-aligned: registers unchanged
-        (4096 + 1234, 4096, 4096 + 1234),  # unaligned: extended down with a matching size pad
-    ])
-def test_npu_register_aligns_to_page_boundary(monkeypatch, address, expected_address, expected_bytes):
-    # MAPPED registration requires 4K-aligned addresses; the hook rounds the
-    # address down to the page boundary and pads the size so the registered
-    # range still covers the original request. An already-aligned address
-    # passes through unchanged, and unregister rounds down identically.
+class _AlignedAccelerator(_RegisteringAccelerator):
+
+    def __init__(self, alignment):
+        super().__init__()
+        self._alignment = alignment
+
+    def pin_memory_alignment(self):
+        return self._alignment
+
+
+class _OffsetHandle:
+    """pin_memory-op stand-in whose buffer base sits at a fixed byte offset."""
+
+    def __init__(self, offset):
+        self._offset = offset
+
+    def new_cpu_locked_tensor(self, numel, example):
+        storage = torch.empty(numel * example.element_size() + self._offset, dtype=torch.uint8)
+        return storage[self._offset:].view(example.dtype)
+
+    def free_cpu_locked_tensor_by_ptr(self, address):
+        return True
+
+
+@pytest.mark.parametrize("alignment, offset", [(1, 64), (2048, 1232), (4096, 64), (4096, 0)])
+def test_device_registration_aligns_to_declared_alignment(native_pins, monkeypatch, alignment, offset):
+    # Accelerators declare the alignment their device runtime requires for
+    # host-memory registration. NativePinnedMemory must round the registered
+    # range down to it, pad the size so the full request is covered, and
+    # unregister the same aligned address. Alignment 1 means no requirement,
+    # so the request passes through unchanged.
+    accelerator = _AlignedAccelerator(alignment)
+    monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: accelerator)
+    monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1")
+    monkeypatch.setattr(native_pins, "_handle", _OffsetHandle(offset))
+
+    pinned = native_pins.pin(torch.empty(32), make_copy=False)
+    begin = pinned.data_ptr()
+    registered_address, registered_bytes = accelerator.registered[0]
+    if alignment == 1:
+        assert registered_address == begin
+    else:
+        assert registered_address % alignment == 0
+        assert registered_address <= begin < registered_address + alignment
+    assert registered_bytes == pinned.nbytes + (begin - registered_address)
+
+    assert native_pins.unpin(pinned) is True
+    assert accelerator.unregistered == [registered_address]
+
+
+def test_npu_declares_page_alignment():
+    # MAPPED registration rejects non-4K-aligned addresses; the declared
+    # alignment is what makes NativePinnedMemory round ranges down for NPU.
+    accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
+    assert accelerator.pin_memory_alignment() == 4096
+
+
+def test_npu_register_uses_mapped_flag(monkeypatch):
+    # MAPPED is deliberate: PINNED-only registrations fall back to mlock-speed
+    # copies on this platform (measured ~9 GB/s vs ~23 GB/s for 64 MiB buffers).
     registered = []
     unregistered = []
 
@@ -313,10 +368,10 @@ def unregister(addr):
     monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None))
     accelerator = NPU_Accelerator.__new__(NPU_Accelerator)
 
-    assert accelerator.register_host_memory(address, 4096) is True
-    assert registered == [(expected_address, expected_bytes, npu_accelerator.ACL_HOST_REG_MAPPED)]
-    accelerator.unregister_host_memory(address)
-    assert unregistered == [expected_address]
+    assert accelerator.register_host_memory(4096, 4096) is True
+    assert registered == [(4096, 4096, npu_accelerator.ACL_HOST_REG_MAPPED)]
+    accelerator.unregister_host_memory(4096)
+    assert unregistered == [4096]
 
 
 def test_npu_device_registration_failure_returns_false(monkeypatch):

From bca3d15beac49bd5d3b65cebe180d33a4ff340fe Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Fri, 18 Sep 2026 06:34:47 +0000
Subject: [PATCH 7/8] Describe native device registration in
 accelerator-neutral terms

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 docs/code-docs/source/memory.rst | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst
index d8d24865f775..614824e0c59a 100644
--- a/docs/code-docs/source/memory.rst
+++ b/docs/code-docs/source/memory.rst
@@ -411,12 +411,13 @@ Example:
 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. 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:
+Native allocations are device-independent ``mlock`` buffers. DeepSpeed
+additionally registers them with the device through the accelerator's
+host-memory registration hook, so PyTorch can use them for asynchronous
+H2D/D2H DMA. Accelerators declare the alignment required for registered host
+memory via ``pin_memory_alignment``; native buffers are rounded down and
+size-padded to it before registration. Device registration is enabled by
+default and can be disabled for comparison or debugging:
 
 .. code-block:: bash
 

From bba5a0f816c4dfc81780333d74c13d99744d4ee3 Mon Sep 17 00:00:00 2001
From: Zizhao Tang <1601677496@qq.com>
Date: Fri, 18 Sep 2026 07:21:20 +0000
Subject: [PATCH 8/8] Skip NPU host copy lookup test on builds without
 torch_npu

Stubbing torch.npu proved unreliable across torch versions (it broke on the cpu-torch-latest CI), so skip the gate test where torch_npu is absent instead of running it with a stand-in.

Signed-off-by: Zizhao Tang <1601677496@qq.com>
---
 tests/unit/v1/pin_memory/test_pin_memory.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py
index 3eb9b5f4b017..20895b0ed554 100644
--- a/tests/unit/v1/pin_memory/test_pin_memory.py
+++ b/tests/unit/v1/pin_memory/test_pin_memory.py
@@ -413,9 +413,11 @@ def test_npu_missing_npurt_is_noop(monkeypatch):
     assert accelerator.unregister_host_memory(4096) is None
 
 
+@pytest.mark.skipif(not hasattr(torch, "npu") or not hasattr(torch.npu, "npurt"), reason="torch_npu is not installed")
 def test_npu_host_copy_lookup_gates(monkeypatch):
-    # torch.npu only exists in torch_npu builds; install a stand-in so this
-    # pure-Python gate logic also runs where the NPU backend is absent.
+    # Exercise each npurt resolution path (missing, init failure, success)
+    # by stubbing torch.npu; where torch_npu is absent the whole test is
+    # skipped because stubbing torch.npu is not reliable across torch versions.
     class _StubNpu:
         pass