deepspeedai / deepspeedai/DeepSpeed
[REQUEST] XPU: register native pinned host memory with the device runtime (follow-up to #8283)
@delock is already working on this.
Since Aug 26, 2026.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Is your feature request related to a problem? Please describe.
#8283 taught the native pinned-memory allocator (DS_PIN_MEMORY_BACKEND=native, posix_memalign + mlock) to register host memory with the device runtime so that copy_(..., non_blocking=True) can go through the DMA engine. On CUDA this closed a large bandwidth gap (H200: ~10-17 GB/s with mlock only vs. ~53-55 GB/s after cudaHostRegister).
Registration goes through the new register_host_memory / unregister_host_memory accelerator hooks, but only CUDA_Accelerator overrides them. XPU_Accelerator inherits the base-class no-op, so on Intel GPUs native pinned buffers stay mlock-only and H2D/D2H transfers run at pageable bandwidth even with DS_PIN_MEMORY_REGISTER_DEVICE=1. The degradation is silent: registration just returns False and nothing logs.
Describe the solution you'd like
- Implement
register_host_memory/unregister_host_memoryinaccelerator/xpu_accelerator.py, e.g. via Level ZerozeMemRegisterHostMemory/zeMemUnregisterMemoryon the allocation range, honoring the sameDS_PIN_MEMORY_REGISTER_DEVICEopt-out and the graceful "log once and continue with mlock" fallback on failure. - Extend
tests/unit/v1/pin_memory/test_pin_memory.pywith an XPU equivalent of the CUDA cudart-mock test. - Optionally extend
benchmarks/pin_memory/h2d_d2h_bench.pyso the before/after H2D/D2H bandwidth can be quantified on XPU.
Describe alternatives you've considered
- Keep mlock-only on XPU (status quo): functionally correct, but leaves H2D/D2H at pageable speeds for native-pin users on Intel GPUs.
Additional context
- Hooks added in #8283: no-op defaults in
accelerator/abstract_accelerator.py, CUDA implementation inaccelerator/cuda_accelerator.py(torch.cuda.cudart().cudaHostRegister/cudaHostUnregister). - Call sites:
NativePinnedMemory.pin/_releaseindeepspeed/utils/pin_memory.py.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.