deepseek-ai/DeepEP

Fork 0

mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00

[PR #750] [MERGED] [HybridEP] Add DMA-BUF fallback for GPU RDMA memory registration #754

New issue

Closed

opened 2026-10-27 21:05:54 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:05:54 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/750
Author: @janbernloehr
Created: 9/12/2026
Status: ✅ Merged
Merged: 10/16/2026
Merged by: @jershi425

Base: hybrid-ep ← Head: jbernloehr/hybridep-dmabuf-registration

📝 Commits (3)

6e2c53a fix(hybrid_ep): add DMA-BUF registration fallback
6e6017c fix(hybrid_ep): harden DMA-BUF registration fallback
4242f70 refactor(hybrid_ep): keep GPU registration state in coordinator

📊 Changes

2 files changed (+297 additions, -62 deletions)

View changed files

📝 csrc/hybrid_ep/buffer/internode_doca.cu (+294 -60)
📝 csrc/hybrid_ep/buffer/internode_doca.cuh (+3 -2)

📄 Description

Summary

Add DMA-BUF support to HybridEP's direct verbs-based GPU memory registration
path.

HybridEP currently detects CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, but the
DOCA/RDMA coordinator still registers every cudaMalloc buffer exclusively
with ibv_reg_mr. On systems where nvidia_peermem is unavailable,
ibv_reg_mr returns nullptr. The returned memory-region pointer is then
dereferenced without validation, causing a segmentation fault while
constructing the HybridEP buffer.

This change:

preserves ibv_reg_mr as the preferred path;
falls back to cuMemGetHandleForAddressRange and ibv_reg_dmabuf_mr when
legacy registration fails;
supports non-page-aligned cudaMalloc pointers by exporting their containing
page-aligned range;
uses the original CUDA virtual addresses in RDMA work requests and exchanged
remote-address metadata;
adds explicit CUDA and verbs error reporting instead of dereferencing a null
ibv_mr;
applies the same checked registration path to all dispatch and combine
buffers.

Closes #367.

Problem

The failure reproduces when all of the following are true:

HybridEP is built with multinode support enabled;
the workload spans multiple nodes;
CUDA reports DMA-BUF support;
nvidia_peermem is not loaded;
the HybridEP DOCA backend attempts to register its cudaMalloc-backed
dispatch and combine buffers.

A representative configuration was:

Setting	Value
Architecture	x86_64
GPUs	64 B200 GPUs
Nodes	8
GPUs per node	8
HybridEP group size	64
Ranks per node	8
Expert parallel size	64
Sequence length	8192
Data	Synthetic
HybridEP multinode support	Enabled
CUDA DMA-BUF support	Available
`nvidia_peermem`	Not loaded

The failure occurs immediately when HybridEP initializes its RDMA buffers,
before the first training iteration completes:

Starting training loop at iteration 0
!!!!!!! Segfault encountered !!!!!!!
  File "<unknown>", line 0, in RDMACoordinator::allocate_combine_buffers()
  File "<unknown>", line 0, in RDMACoordinator::allocate_buffers()
  File "<unknown>", line 0, in HybridEPBuffer::allocate_buffer()
  File "<unknown>", line 0, in HybridEPBuffer::HybridEPBuffer(...)
  File "<unknown>", line 0, in pybind11::cpp_function::dispatcher(...)

A similar public failure signature is reported in
NVIDIA-NeMo/Megatron-Bridge#5051.

Root cause

HybridEP's DOCA backend allocates the communication buffers with cudaMalloc
and registers them directly with libibverbs:

ibv_mr* mr = ibv_reg_mr(pd, address, length, access);

This path relies on legacy peer-memory support for GPU allocations. Without
nvidia_peermem, the call fails and returns nullptr.

Although HybridEP already queries CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, the
result was not used by the buffer registration code. NVSHMEM's own DMA-BUF
support does not cover these allocations because the HybridEP RDMA coordinator
registers these buffers directly rather than through NVSHMEM's
memory-registration path.

The subsequent HybridEP code accesses fields such as mr->lkey, mr->rkey,
and mr->addr without checking the result, producing the segmentation fault.

DMA-BUF alignment requirement

The first checked implementation attempted to export the exact CUDA pointer
and size. That replaced the segmentation fault with a useful error and exposed
a second issue:

HybridEP cannot DMA-BUF register GPU buffer 'combine.attn_output_flags':
address=0xfa19e11800, length=3696, page_size=4096

Ordinary cudaMalloc addresses are not guaranteed to be aligned to the host
page size. cuMemGetHandleForAddressRange requires the exported address range
to satisfy the DMA-BUF alignment requirements.

The final implementation therefore exports the containing page-aligned range:

original buffer:  [address, address + length)
exported range:   [align_down(address), align_up(address + length))

The implementation also checks for arithmetic overflow before calculating the
aligned range.

Implementation

A shared register_gpu_memory helper now performs the following sequence for
every dispatch and combine buffer:

Call ibv_reg_mr first.

This preserves existing behavior and avoids changing systems where legacy
peer-memory registration already works.
If ibv_reg_mr fails, save its errno.

The error is preserved before any CUDA or diagnostic calls can overwrite
it.
Verify that CUDA reports DMA-BUF support.

If DMA-BUF is unavailable, raise an error containing the buffer name, the
legacy verbs error, and the fact that CUDA does not support DMA-BUF on the
device.
Determine the host page size with sysconf(_SC_PAGESIZE).
Align the start address down and the end address up to cover the entire CUDA
allocation range.

Export the aligned range using:

cuMemGetHandleForAddressRange(
    &dmabuf_fd,
    aligned_address,
    aligned_length,
    CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD,
    0);

ibv_reg_dmabuf_mr(
    pd,
    0,
    aligned_length,
    aligned_address,
    dmabuf_fd,
    access);

Close the exported file descriptor after registration.
Use the resulting memory region's lkey and rkey, while retaining the
original CUDA buffer addresses for work requests and remote-address
exchange.

The last point is important for subranges of the exported allocation.
ibv_mr::addr must not be assumed to represent the original, potentially
unaligned CUDA buffer pointer. Dispatch and combine metadata now explicitly
uses the corresponding CUDA allocation pointers.

The fallback covers:

Dispatch buffers

attention input tokens;
inter-node group tokens;
attention input flags;
inter-node group flags;
attention input probabilities;
inter-node group probabilities;
attention input scaling factors;
inter-node group scaling factors.

Combine buffers

intra-node reduction tokens;
inter-node group tokens;
attention output flags;
inter-node group flags;
intra-node reduction probabilities;
inter-node group probabilities.

When the fallback is selected, rank 0 of each node reports:

[HybridEP] legacy GPU memory registration is unavailable; using DMA-BUF registration

CUDA export failures and DMA-BUF verbs registration failures include the
buffer name, CUDA error name and description or verbs errno, and the original
ibv_reg_mr error where applicable.

Validation

The implementation was tested with three multinode MoE training
configurations:

BF16;
FP8;
NVFP4.

Each configuration used 64 B200 GPUs across eight nodes, eight ranks per node,
and EP64. Each test ran 30 training iterations with synthetic data.

DMA-BUF environment

The environment had CUDA DMA-BUF support and did not load nvidia_peermem.

Before this change:

all three configurations failed during HybridEP initialization;
the failure occurred at iteration 0;
the stack terminated in RDMACoordinator::allocate_combine_buffers;
no training metrics were produced.

After this change:

Precision	Result	Iterations
BF16	Passed	30/30
FP8	Passed	30/30
NVFP4	Passed	30/30

The DMA-BUF fallback banner appeared on all eight nodes. Every rank exited
successfully.

Legacy peer-memory environment

The same three configurations were also run on a second 8-node B200
environment where ibv_reg_mr can register GPU memory through the legacy
peer-memory path.

Precision	Result	Iterations
BF16	Passed	30/30
FP8	Passed	30/30
NVFP4	Passed	30/30

The DMA-BUF fallback banner did not appear, confirming that the existing
ibv_reg_mr path remained active.

Performance sanity check

Across the matched configurations, throughput in the DMA-BUF environment was
within +/-2.4% of the legacy peer-memory environment:

Precision	DMA-BUF vs. legacy
BF16	+2.39%
FP8	-1.19%
NVFP4	+2.31%

The geometric mean of the matched ratios was +1.16%. These were single runs on
two independent systems rather than a controlled performance benchmark, but
they show no obvious throughput regression from the DMA-BUF registration path.

Tested revisions

The full multinode runtime matrix was executed with HybridEP revision:

58397eac493b075b0baea5861ac6e9793d8d7a21

The same patch applies cleanly to the current hybrid-ep tip used as the base
of this change:

10d4dd7377d5bce900fbb4b80cce863764892b95

The current-tip branch passes git diff --check. The full 64-GPU runtime
validation described above was performed on the earlier tested revision.

Compatibility

Systems where ibv_reg_mr succeeds continue using the existing path.
DMA-BUF is attempted only after legacy registration fails.
Systems supporting neither registration mechanism now receive a descriptive
exception instead of a null-pointer segmentation fault.
No Python API or configuration changes are required.
Intranode-only behavior is unchanged.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/750 **Author:** [@janbernloehr](https://github.com/janbernloehr) **Created:** 9/12/2026 **Status:** ✅ Merged **Merged:** 10/16/2026 **Merged by:** [@jershi425](https://github.com/jershi425) **Base:** `hybrid-ep` ← **Head:** `jbernloehr/hybridep-dmabuf-registration` --- ### 📝 Commits (3) - [`6e2c53a`](https://github.com/deepseek-ai/DeepEP/commit/6e2c53ace2dc86cab4849211db08014a34aae5a0) fix(hybrid_ep): add DMA-BUF registration fallback - [`6e6017c`](https://github.com/deepseek-ai/DeepEP/commit/6e6017c62ea466b5c51deea10d7a803d0958b54d) fix(hybrid_ep): harden DMA-BUF registration fallback - [`4242f70`](https://github.com/deepseek-ai/DeepEP/commit/4242f706f25adc34acd8dfd002ed83dd3ac5b6ad) refactor(hybrid_ep): keep GPU registration state in coordinator ### 📊 Changes **2 files changed** (+297 additions, -62 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/buffer/internode_doca.cu` (+294 -60) 📝 `csrc/hybrid_ep/buffer/internode_doca.cuh` (+3 -2) </details> ### 📄 Description ## Summary Add DMA-BUF support to HybridEP's direct verbs-based GPU memory registration path. HybridEP currently detects `CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED`, but the DOCA/RDMA coordinator still registers every `cudaMalloc` buffer exclusively with `ibv_reg_mr`. On systems where `nvidia_peermem` is unavailable, `ibv_reg_mr` returns `nullptr`. The returned memory-region pointer is then dereferenced without validation, causing a segmentation fault while constructing the HybridEP buffer. This change: - preserves `ibv_reg_mr` as the preferred path; - falls back to `cuMemGetHandleForAddressRange` and `ibv_reg_dmabuf_mr` when legacy registration fails; - supports non-page-aligned `cudaMalloc` pointers by exporting their containing page-aligned range; - uses the original CUDA virtual addresses in RDMA work requests and exchanged remote-address metadata; - adds explicit CUDA and verbs error reporting instead of dereferencing a null `ibv_mr`; - applies the same checked registration path to all dispatch and combine buffers. Closes #367. ## Problem The failure reproduces when all of the following are true: - HybridEP is built with multinode support enabled; - the workload spans multiple nodes; - CUDA reports DMA-BUF support; - `nvidia_peermem` is not loaded; - the HybridEP DOCA backend attempts to register its `cudaMalloc`-backed dispatch and combine buffers. A representative configuration was: | Setting | Value | | --- | --- | | Architecture | x86_64 | | GPUs | 64 B200 GPUs | | Nodes | 8 | | GPUs per node | 8 | | HybridEP group size | 64 | | Ranks per node | 8 | | Expert parallel size | 64 | | Sequence length | 8192 | | Data | Synthetic | | HybridEP multinode support | Enabled | | CUDA DMA-BUF support | Available | | `nvidia_peermem` | Not loaded | The failure occurs immediately when HybridEP initializes its RDMA buffers, before the first training iteration completes: ```text Starting training loop at iteration 0 !!!!!!! Segfault encountered !!!!!!! File "<unknown>", line 0, in RDMACoordinator::allocate_combine_buffers() File "<unknown>", line 0, in RDMACoordinator::allocate_buffers() File "<unknown>", line 0, in HybridEPBuffer::allocate_buffer() File "<unknown>", line 0, in HybridEPBuffer::HybridEPBuffer(...) File "<unknown>", line 0, in pybind11::cpp_function::dispatcher(...) ``` A similar public failure signature is reported in [NVIDIA-NeMo/Megatron-Bridge#5051](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/5051). ## Root cause HybridEP's DOCA backend allocates the communication buffers with `cudaMalloc` and registers them directly with libibverbs: ```cpp ibv_mr* mr = ibv_reg_mr(pd, address, length, access); ``` This path relies on legacy peer-memory support for GPU allocations. Without `nvidia_peermem`, the call fails and returns `nullptr`. Although HybridEP already queries `CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED`, the result was not used by the buffer registration code. NVSHMEM's own DMA-BUF support does not cover these allocations because the HybridEP RDMA coordinator registers these buffers directly rather than through NVSHMEM's memory-registration path. The subsequent HybridEP code accesses fields such as `mr->lkey`, `mr->rkey`, and `mr->addr` without checking the result, producing the segmentation fault. ## DMA-BUF alignment requirement The first checked implementation attempted to export the exact CUDA pointer and size. That replaced the segmentation fault with a useful error and exposed a second issue: ```text HybridEP cannot DMA-BUF register GPU buffer 'combine.attn_output_flags': address=0xfa19e11800, length=3696, page_size=4096 ``` Ordinary `cudaMalloc` addresses are not guaranteed to be aligned to the host page size. `cuMemGetHandleForAddressRange` requires the exported address range to satisfy the DMA-BUF alignment requirements. The final implementation therefore exports the containing page-aligned range: ```text original buffer: [address, address + length) exported range: [align_down(address), align_up(address + length)) ``` The implementation also checks for arithmetic overflow before calculating the aligned range. ## Implementation A shared `register_gpu_memory` helper now performs the following sequence for every dispatch and combine buffer: 1. Call `ibv_reg_mr` first. This preserves existing behavior and avoids changing systems where legacy peer-memory registration already works. 2. If `ibv_reg_mr` fails, save its `errno`. The error is preserved before any CUDA or diagnostic calls can overwrite it. 3. Verify that CUDA reports DMA-BUF support. If DMA-BUF is unavailable, raise an error containing the buffer name, the legacy verbs error, and the fact that CUDA does not support DMA-BUF on the device. 4. Determine the host page size with `sysconf(_SC_PAGESIZE)`. 5. Align the start address down and the end address up to cover the entire CUDA allocation range. 6. Export the aligned range using: ```cpp cuMemGetHandleForAddressRange( &dmabuf_fd, aligned_address, aligned_length, CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0); ``` 7. Register the exported descriptor using: ```cpp ibv_reg_dmabuf_mr( pd, 0, aligned_length, aligned_address, dmabuf_fd, access); ``` 8. Close the exported file descriptor after registration. 9. Use the resulting memory region's `lkey` and `rkey`, while retaining the original CUDA buffer addresses for work requests and remote-address exchange. The last point is important for subranges of the exported allocation. `ibv_mr::addr` must not be assumed to represent the original, potentially unaligned CUDA buffer pointer. Dispatch and combine metadata now explicitly uses the corresponding CUDA allocation pointers. The fallback covers: ### Dispatch buffers - attention input tokens; - inter-node group tokens; - attention input flags; - inter-node group flags; - attention input probabilities; - inter-node group probabilities; - attention input scaling factors; - inter-node group scaling factors. ### Combine buffers - intra-node reduction tokens; - inter-node group tokens; - attention output flags; - inter-node group flags; - intra-node reduction probabilities; - inter-node group probabilities. When the fallback is selected, rank 0 of each node reports: ```text [HybridEP] legacy GPU memory registration is unavailable; using DMA-BUF registration ``` CUDA export failures and DMA-BUF verbs registration failures include the buffer name, CUDA error name and description or verbs `errno`, and the original `ibv_reg_mr` error where applicable. ## Validation The implementation was tested with three multinode MoE training configurations: - BF16; - FP8; - NVFP4. Each configuration used 64 B200 GPUs across eight nodes, eight ranks per node, and EP64. Each test ran 30 training iterations with synthetic data. ### DMA-BUF environment The environment had CUDA DMA-BUF support and did not load `nvidia_peermem`. Before this change: - all three configurations failed during HybridEP initialization; - the failure occurred at iteration 0; - the stack terminated in `RDMACoordinator::allocate_combine_buffers`; - no training metrics were produced. After this change: | Precision | Result | Iterations | Skipped/NaN | | --- | --- | ---: | ---: | | BF16 | Passed | 30/30 | 0 | | FP8 | Passed | 30/30 | 0 | | NVFP4 | Passed | 30/30 | 0 | The DMA-BUF fallback banner appeared on all eight nodes. Every rank exited successfully. ### Legacy peer-memory environment The same three configurations were also run on a second 8-node B200 environment where `ibv_reg_mr` can register GPU memory through the legacy peer-memory path. | Precision | Result | Iterations | Skipped/NaN | | --- | --- | ---: | ---: | | BF16 | Passed | 30/30 | 0 | | FP8 | Passed | 30/30 | 0 | | NVFP4 | Passed | 30/30 | 0 | The DMA-BUF fallback banner did not appear, confirming that the existing `ibv_reg_mr` path remained active. ### Performance sanity check Across the matched configurations, throughput in the DMA-BUF environment was within +/-2.4% of the legacy peer-memory environment: | Precision | DMA-BUF vs. legacy | | --- | ---: | | BF16 | +2.39% | | FP8 | -1.19% | | NVFP4 | +2.31% | The geometric mean of the matched ratios was +1.16%. These were single runs on two independent systems rather than a controlled performance benchmark, but they show no obvious throughput regression from the DMA-BUF registration path. ## Tested revisions The full multinode runtime matrix was executed with HybridEP revision: ```text 58397eac493b075b0baea5861ac6e9793d8d7a21 ``` The same patch applies cleanly to the current `hybrid-ep` tip used as the base of this change: ```text 10d4dd7377d5bce900fbb4b80cce863764892b95 ``` The current-tip branch passes `git diff --check`. The full 64-GPU runtime validation described above was performed on the earlier tested revision. ## Compatibility - Systems where `ibv_reg_mr` succeeds continue using the existing path. - DMA-BUF is attempted only after legacy registration fails. - Systems supporting neither registration mechanism now receive a descriptive exception instead of a null-pointer segmentation fault. - No Python API or configuration changes are required. - Intranode-only behavior is unchanged. ## Related issues - [deepseek-ai/DeepEP#367](https://github.com/deepseek-ai/DeepEP/issues/367) - [NVIDIA-NeMo/Megatron-Bridge#5051](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/5051) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-27 21:05:54 +00:00