Register native pinned host memory with the CANN runtime for NPU by VenusTZZ · Pull Request #8532 · deepspeedai/DeepSpeed

VenusTZZ · 2026-09-16T03:01:11Z

Follow-up to #8283; XPU counterpart tracked in #8315 (attempted in #8322).

Summary

Implement register_host_memory / unregister_host_memory on NPU_Acceleratorvia the NPU runtime API torch_npu ships: torch.npu.npurt() returns the runtime-API module with npuHostRegister / npuHostUnregister, thin wrappers over aclrtHostRegisterV2 / aclrtHostUnregister. Previously NPU_Accelerator inherited the base-class no-ops, so DS_PIN_MEMORY_BACKEND=native buffers stayed mlock-only on Ascend NPUs and asynchronous copies ran at pageable bandwidth.
Degrades to "log once and continue with mlock" on every failure path and keeps honoring the existing DS_PIN_MEMORY_REGISTER_DEVICE opt-out.
Extends tests/unit/v1/pin_memory/test_pin_memory.py with NPU cases mirroring the CUDA cudart-mock tests, including the availability fallbacks.
Makes benchmarks/pin_memory/h2d_d2h_bench.py accelerator-agnostic (it hardcoded a CUDA check and torch's pin_memory=True fast path) so the before/after bandwidth could be measured on NPU.
Documents the NPU registration path in docs/code-docs/source/memory.rst.

Implementation notes

The npurt binding. torch_npu exposes torch.npu.npurt(), which initializes the runtime if needed and returns the runtime-API module. Using it avoids binding libascendcl.so ourselves and raises a clear RuntimeError when the build lacks the module. Two defenses gate the hook, both falling back to mlock-only with a one-time warning:

no torch.npu.npurt attribute (the hook checks for its presence),
npurt() failing to initialize the runtime.

Flag choice: ACL_HOST_REG_MAPPED. The flags parameter of npuHostRegister is the CANN ACL_HOST_REG_* bitmask (per the npurt docstring and acl_rt.h): MAPPED (0x2) page-locks the range and adds a device mapping, PINNED (0x10000000) only page-locks. Measured through the bench methodology at 64 MiB (10 process-level A/B runs): PINNED-only registrations fell back to mlock-only H2D speed (~8-9 GB/s) in 5/5 runs, while MAPPED kept ~23 GB/s (5/6). MAPPED additionally requires 4K-aligned addresses (a 64B-offset probe fails with ACL_ERROR_RT_DRV_INTERNAL_ERROR, 507899), which DeepSpeed's native allocator guarantees (posix_memalign(..., sysconf(_SC_PAGESIZE), ...)); sizes need no page alignment (probed 400 B to 1 MiB+3, rc=0). As a guard against unaligned callers, the hook extends the range down to the page boundary, padding the size accordingly, instead of relying on the driver's opaque internal error.

Failure semantics. A non-zero return code from npuHostRegister logs once and returns False, so NativePinnedMemory keeps the buffer mlock-only and does not track the address. A failed npuHostUnregister raises (like the CUDA hook), which makes NativePinnedMemory keep the allocation alive so the driver never holds a registration for pages later reused by malloc.

Bandwidth on Ascend 910B4

benchmarks/pin_memory/h2d_d2h_bench.py, one full run of the #8283 methodology (--sizes-mib 4 64 --warmup 10 --iters 50), CANN 8.5.0, torch 2.10.0+cpu, torch_npu 2.10.0, Ascend 910B4 (aarch64), shared box:

Arm	Size (MiB)	H2D (GB/s)	D2H (GB/s)	torch.is_pinned
torch	4	10.4	21.1	True
torch	64	16.5	21.0	True
native-unregistered	4	9.7	9.4	False
native-unregistered	64	21.1	19.6	False
native-registered	4	22.7	20.6	True
native-registered	64	23.4	20.9	True

After registration, torch.is_pinned() on a torch.frombuffer view of the range returns True — the torch_npu runtime recognizes the registered range — so the DMA path is confirmed to engage, not just inferred from the numbers.

Registration lifts H2D bandwidth from ~9.7 to ~22.7 GB/s at 4 MiB (2.3x) and holds ~23.4 GB/s at 64 MiB, with D2H improving as well (9.4 → 20.6 and 19.6 → 20.9 GB/s). Registered copies match or exceed the torch pin_memory() arm at every size.

Test plan

pre-commit run --files on touched paths — yapf, flake8, check-torchdist, check-torchcuda, check-license, codespell all pass
pytest tests/unit/v1/pin_memory/test_pin_memory.py — 24 passed on Ascend 910B4 (18 existing incl. the CUDA cudart-mock tests + 6 new NPU cases: npurt register/unregister calls, unaligned address extended to the page boundary, non-zero rc degrades to mlock, unregister failure raises, missing npurt is a no-op, availability gate table)
benchmarks/pin_memory/h2d_d2h_bench.py on Ascend 910B4 — all three arms run to completion and exit 0 at 4/64 MiB

Known limitation

Some torch_npu builds may ship without the npurt module. The hasattr check degrades to mlock-only with a one-time warning, so registration is only available when the binding is present.

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 deepspeedai#8531 Signed-off-by: Zizhao Tang <1601677496@qq.com>

hipudding · 2026-09-16T06:54:05Z

+    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:


Do we need to check torch_npu' version? I think check npurt is enough.

Agreed — the version gate could reject builds that actually ship npurt, and the hasattr check alone covers the case we care about. Removed it in b7e6784; the hook now falls back based only on whether torch.npu.npurt is present (plus npurt() initializing successfully).

… version Signed-off-by: Zizhao Tang <1601677496@qq.com>

hipudding · 2026-09-16T07:42:40Z

+        # 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:


We can align address to 4k, remember ping memory always contains all the buffer, and add num_bytes after align.

Agreed — Done in bb2e596 — the hook rounds the address down to the 4K page boundary and pads num_bytes by the same offset, so the registered range always covers the original buffer exactly. An already-aligned address passes through unchanged (offset 0), and unregister_host_memory applies the same rounding so both sides stay consistent.

Signed-off-by: Zizhao Tang <1601677496@qq.com>

hipudding

Most of these changes already very good, though there are a few minor points I suggest revising.

hipudding · 2026-09-16T09:25:22Z

+        # 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


It's better to make align as a function, register and unregister function can call align func. Avoid inconsistent

Thanks! Done in 8ba87d1 — extracted _align_to_page_boundary(address) and both hooks now call it: register_host_memory uses the aligned address plus the returned offset for the size pad, and unregister_host_memory passes the aligned address back to the driver. The rounding now lives in one place, so the two can't drift apart.

hipudding · 2026-09-16T09:26:08Z

+    assert unregistered == [4096]
+
+
+def test_npu_unaligned_address_is_extended_to_page_boundary(monkeypatch):


Should test unaligned address and aligned address.

Done in 8ba87d1 — the two tests were consolidated into one parametrized case (test_npu_register_aligns_to_page_boundary) covering both addresses: the page-aligned address (4096) registers unchanged, the unaligned one (4096+1234) is extended down to the boundary with a matching size pad, and unregister rounds down identically in both cases.

…ster hooks Signed-off-by: Zizhao Tang <1601677496@qq.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8ba87d196e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-09-16T09:57:47Z

+
+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)


Guard the NPU lookup test on builds without torch.npu

On the regular CPU/CUDA PyTorch builds that run this unmarked unit-test file, torch has no npu attribute. This dereference raises AttributeError before the helper can exercise its intended missing-npurt fallback, so the newly added test fails the non-NPU unit-test suite; install a stub torch.npu for this test or skip it when the backend is absent.

Useful? React with 👍 / 👎.

Done in 91d7f01 — the test now installs a stand-in torch.npu via monkeypatch before touching npurt, so the gate logic also runs on builds without the backend.

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>

VenusTZZ · 2026-09-16T10:06:53Z

Hi @delock, this PR follows up on #8283 by implementing the register_host_memory / unregister_host_memory hooks for the NPU accelerator (the XPU counterpart is tracked in #8315). Would you be able to take a look when you have a chance? Thanks!

hipudding reviewed Sep 16, 2026

View reviewed changes

Gate NPU host registration on npurt availability instead of torch_npu…

b7e6784

… version Signed-off-by: Zizhao Tang <1601677496@qq.com>

hipudding reviewed Sep 16, 2026

View reviewed changes

VenusTZZ added 2 commits September 16, 2026 08:50

Extend unaligned NPU host registrations to the 4K page boundary

bb2e596

Signed-off-by: Zizhao Tang <1601677496@qq.com>

Merge branch 'master' into npu-register-host-memory

09f8306

hipudding reviewed Sep 16, 2026

View reviewed changes

Share page-boundary alignment between NPU host memory register/unregi…

8ba87d1

…ster hooks Signed-off-by: Zizhao Tang <1601677496@qq.com>

VenusTZZ marked this pull request as ready for review September 16, 2026 09:53

VenusTZZ requested review from loadams, tjruwase and tohtana as code owners September 16, 2026 09:53

chatgpt-codex-connector Bot reviewed Sep 16, 2026

View reviewed changes

delock self-requested a review September 16, 2026 13:39

Navigation Menu

Uh oh!

Register native pinned host memory with the CANN runtime for NPU - #8532

Register native pinned host memory with the CANN runtime for NPU#8532
VenusTZZ wants to merge 6 commits into
deepspeedai:masterfrom
VenusTZZ:npu-register-host-memory

VenusTZZ commented Sep 16, 2026 •

edited

Loading

Uh oh!

hipudding Sep 16, 2026

Uh oh!

VenusTZZ Sep 16, 2026 •

edited

Loading

Uh oh!

hipudding Sep 16, 2026

Uh oh!

VenusTZZ Sep 16, 2026

Uh oh!

hipudding left a comment

Uh oh!

hipudding Sep 16, 2026

Uh oh!

VenusTZZ Sep 16, 2026

Uh oh!

hipudding Sep 16, 2026

Uh oh!

VenusTZZ Sep 16, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Sep 16, 2026

Uh oh!

VenusTZZ Sep 16, 2026 •

edited

Loading

Uh oh!

VenusTZZ commented Sep 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		assert unregistered == [4096]


		def test_npu_unaligned_address_is_extended_to_page_boundary(monkeypatch):

Uh oh!

Conversation

VenusTZZ commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation notes

Bandwidth on Ascend 910B4

Test plan

Known limitation

Uh oh!

hipudding Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

VenusTZZ Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hipudding Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

VenusTZZ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

hipudding left a comment

Choose a reason for hiding this comment

Uh oh!

hipudding Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

VenusTZZ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

hipudding Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

VenusTZZ Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

VenusTZZ Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VenusTZZ commented Sep 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

VenusTZZ commented Sep 16, 2026 •

edited

Loading

VenusTZZ Sep 16, 2026 •

edited

Loading

VenusTZZ Sep 16, 2026 •

edited

Loading