Register native pinned host memory with CUDA for GPU DMA

[sfc-gh-truwase]

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

[chatgpt-codex-connector]

💡 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".

[delock]

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.

[sfc-gh-truwase]

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.