FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Register native pinned host memory with CUDA for GPU DMA by sfc-gh-truwase · Pull Request #8283 · deepspeedai/DeepSpeed · GitHub

/ DeepSpeed Public

Register native pinned host memory with CUDA for GPU DMA - #8283

Merged
sfc-gh-truwase merged 3 commits into
masterfrom
tjruwase/native-pin-device-register
Aug 23, 2026
ConversationCommits (3)ChecksFiles changed
Merged

Register native pinned host memory with CUDA for GPU DMA#8283
sfc-gh-truwase merged 3 commits into
masterfrom
tjruwase/native-pin-device-register

Conversation

sfc-gh-truwase commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Native pin (posix_memalign + mlock) is device-independent and sufficient for DeepNVMe, but the GPU DMA engine does not treat mlock pages as pinned. After allocation, optionally call cudaHostRegister (and cudaHostUnregister before free) via accelerator hooks so copy_(..., non_blocking=True) can DMA.
  • Default-on via DS_PIN_MEMORY_REGISTER_DEVICE (1/true/yes/on; 0 keeps mlock-only). Registration failure logs once and continues with mlock. CPU accelerators no-op.
  • Adds benchmarks/pin_memory/h2d_d2h_bench.py (torch vs native-unregistered vs native-registered).

H200 H2D/D2H (1 GPU)

Host tunji-h200-n1g2-ds-pin-0, NVIDIA H200, --sizes-mib 4 64 256 --warmup 10 --iters 50, autorun job-20260820T155549Z.

Arm Size (MiB) H2D (GB/s) D2H (GB/s) torch.is_pinned
torch 4 53.20 52.96 True
torch 64 55.32 52.49 True
torch 256 55.41 52.30 True
native-unregistered 4 16.95 14.05 False
native-unregistered 64 16.45 15.45 False
native-unregistered 256 9.98 16.42 False
native-registered 4 53.16 52.78 True
native-registered 64 55.29 54.79 True
native-registered 256 55.40 54.94 True

Registered native matches torch bandwidth (~53–55 GB/s). mlock without cudaHostRegister stays in the ~10–17 GB/s pageable range and reports torch.is_pinned=False.

Test plan

  • pre-commit run --files on touched paths
  • pytest tests/unit/v1/pin_memory/test_pin_memory.py 16 passed (register on/off, CPU no-op, CUDA cudart mock, register failure keeps mlock, GC unregister, invalid env)
  • GPU autorun job-20260820T155549Z: pin_memory + accelerator tests 25 passed, H2D/D2H bench exit 0

Made with Cursor

Register native mlock pages with CUDA so GPU DMA can use them.
e6bea83 
mlock alone is page-locked for AIO but not device-pinned; cudaHostRegister after alloc restores torch-class H2D/D2H bandwidth.

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
chatgpt-codex-connector Bot reviewed Aug 20, 2026
View reviewed changes

chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e6bea83b5d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread deepspeed/utils/pin_memory.py Outdated
Do not free native pin pages if device unregister fails.
bbfd83a 
cudaHostUnregister must succeed before returning the allocation, or the driver can keep a registration on recycled host pages.

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
sfc-gh-truwase requested a review from delock August 20, 2026 16:32
sfc-gh-truwase mentioned this pull request Aug 20, 2026
Merged
4 tasks

delock commented Aug 22, 2026

Copy link
Copy Markdown
Collaborator

Hi @sfc-gh-truwase I'm generally okay with this PR. The only thing I want to add is XPU implementation of register_host_memory, which should be added from a seperate PR. I'll approve your PR.

PKUWZP self-requested a review August 22, 2026 04:51

sfc-gh-truwase commented Aug 23, 2026

Copy link
Copy Markdown
Collaborator Author

Hi @sfc-gh-truwase I'm generally okay with this PR. The only thing I want to add is XPU implementation of register_host_memory, which should be added from a seperate PR. I'll approve your PR.

@delock, good call. Please let me know when this is done.

sfc-gh-truwase added this pull request to the merge queue Aug 23, 2026
Merged via the queue into master with commit cb26080 Aug 23, 2026
14 checks passed
sfc-gh-truwase deleted the tjruwase/native-pin-device-register branch August 23, 2026 23:04
delock mentioned this pull request Aug 25, 2026
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Show hidden characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

chatgpt-codex-connector[bot] chatgpt-codex-connector[bot] left review comments
PKUWZP PKUWZP approved these changes
delock delock approved these changes
tjruwase Awaiting requested review from tjruwase tjruwase is a code owner
tohtana Awaiting requested review from tohtana tohtana is a code owner
loadams Awaiting requested review from loadams loadams is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

Back | FazBrowse Home | New Git URL