kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #674] [Feature Request]: [TransferEngine] Migrate RDMA to use dma buf compatible library calls (ie ibv_reg_mr(...) calls to ibv_reg_dmabuf_mr(...)) #216

New issue

Closed

opened 2026-10-29 21:11:00 +00:00 by bsakaguchi · 8 comments

bsakaguchi commented

2026-10-29 21:11:00 +00:00

Owner

Originally created by @misterwilliam on GitHub (Jul 25, 2025).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/674

Describe your feature request

I would like to run Mooncake in an environment where nvidia-peermem is not installed.

To do that I need to migrate the RDMA memory register from using ibv_reg_mr(...) to using ibv_reg_dmabuf_mr(...). Also before calling I will also need to call cuMemGetHandleForAddressRange()2 to get the dma buf struct. This should have no performance impact, it just changes the way these libraries are called.

Using ibv_reg_mr(...) is considered deprecated 1 and using ibv_reg_dma_mr(...) is the new way which does not rely on proprietary software. Some more historical / technical background on this can be found here: https://medium.com/@datenlord/the-evolution-and-implementation-of-gpudirect-rdma-19751f7b9413

I am happy to contribute the pull request to make this happen. I just wanted to start the conversation.

Proposed code change:

int dmaBufSupported;
if (cuDeviceGetAttribute(&dmaBufSupported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, device) == CUDA_SUCCESS && dmaBufSupported) {
    ibv_reg_dmabuf_mr(...)
} else {
    ibv_reg_mr(...)
}

Before submitting a new issue...

Make sure you already searched for relevant issues and read the documentation

Originally created by @misterwilliam on GitHub (Jul 25, 2025). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/674 ### Describe your feature request I would like to run Mooncake in an environment where nvidia-peermem is not installed. To do that I need to migrate the RDMA memory register from using ibv_reg_mr(...) to using ibv_reg_dmabuf_mr(...). Also before calling I will also need to call cuMemGetHandleForAddressRange()[2] to get the dma buf struct. This should have no performance impact, it just changes the way these libraries are called. Using ibv_reg_mr(...) is considered deprecated [1] and using ibv_reg_dma_mr(...) is the new way which does not rely on proprietary software. Some more historical / technical background on this can be found here: https://medium.com/@datenlord/the-evolution-and-implementation-of-gpudirect-rdma-19751f7b9413 I am happy to contribute the pull request to make this happen. I just wanted to start the conversation. Proposed code change: ``` int dmaBufSupported; if (cuDeviceGetAttribute(&dmaBufSupported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, device) == CUDA_SUCCESS && dmaBufSupported) { ibv_reg_dmabuf_mr(...) } else { ibv_reg_mr(...) } ``` [1]: https://discourse.ubuntu.com/t/nvidia-gpudirect-over-infiniband-migration-paths/44425 [2]: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g51e719462c04ee90a6b0f8b2a75fe031 ### Before submitting a new issue... - [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

bsakaguchi closed this issue

2026-10-29 21:11:00 +00:00

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@staryxchen commented on GitHub (Jul 26, 2025):

Are you trying to run Mooncake in a non-Nvidia environment (GPU/RDMA NIC)? If so, I think we can control the use of different APIs to adapt to different hardware environments through compilation options.

@staryxchen commented on GitHub (Jul 26, 2025): Are you trying to run Mooncake in a non-Nvidia environment (GPU/RDMA NIC)? If so, I think we can control the use of different APIs to adapt to different hardware environments through compilation options.

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@misterwilliam commented on GitHub (Jul 26, 2025):

I am trying to use Mooncake with nvidia (b200), but I am trying to run with (Google) GKE Autopilot which uses COS (Container Optimized OS) which has constraints on being allowed to install kernel modules (ie install nvidia-peermem is not allowed).

My change should be thought of the updating nvidia RDMA implementation from deprecated API to supported API. I've updated my initial proposal with a proposed code change snippet to show it can be made to made backwards compatible in situations where dma-buf is not supported.

@misterwilliam commented on GitHub (Jul 26, 2025): I am trying to use Mooncake with nvidia (b200), but I am trying to run with (Google) GKE Autopilot which uses COS (Container Optimized OS) which has constraints on being allowed to install kernel modules (ie install nvidia-peermem is not allowed). My change should be thought of the updating nvidia RDMA implementation from deprecated API to supported API. I've updated my initial proposal with a proposed code change snippet to show it can be made to made backwards compatible in situations where dma-buf is not supported.

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@stmatengss commented on GitHub (Jul 27, 2025):

Regarding the difference between ibv_reg_mr and ibv_reg_dma_mr (see https://forums.developer.nvidia.com/t/gpudirect-rdma-difference-between-ibv-reg-mr-and-ibv-reg-dmabuf-mr/262313, https://man7.org/linux/man-pages/man3/ibv_reg_mr.3.html), it appears that ibv_reg_dma_mr is the preferred alternative to ibv_reg_mr when the NV peer-memory driver is unavailable. My concern is ensuring DRAM compatibility across all RDMA transport paths, as we don't differentiate between VRAM and DRAM.

@stmatengss commented on GitHub (Jul 27, 2025): Regarding the difference between `ibv_reg_mr` and `ibv_reg_dma_mr` (see [https://forums.developer.nvidia.com/t/gpudirect-rdma-difference-between-ibv-reg-mr-and-ibv-reg-dmabuf-mr/262313](https://forums.developer.nvidia.com/t/gpudirect-rdma-difference-between-ibv-reg-mr-and-ibv-reg-dmabuf-mr/262313), [https://man7.org/linux/man-pages/man3/ibv_reg_mr.3.html](https://man7.org/linux/man-pages/man3/ibv_reg_mr.3.html)), it appears that `ibv_reg_dma_mr` is the preferred alternative to `ibv_reg_mr` when the NV peer-memory driver is unavailable. My concern is ensuring DRAM compatibility across all RDMA transport paths, as we don't differentiate between VRAM and DRAM.

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@stmatengss commented on GitHub (Jul 27, 2025):

Are you trying to run Mooncake in a non-Nvidia environment (GPU/RDMA NIC)? If so, I think we can control the use of different APIs to adapt to different hardware environments through compilation options.

That's correct; it should be done in the compilation stage.

@stmatengss commented on GitHub (Jul 27, 2025): > Are you trying to run Mooncake in a non-Nvidia environment (GPU/RDMA NIC)? If so, I think we can control the use of different APIs to adapt to different hardware environments through compilation options. That's correct; it should be done in the compilation stage.

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@misterwilliam commented on GitHub (Jul 27, 2025):

So it sounds like we want to:

Use compilations options to guard the change.
Detect if the address is a VRAM or DRAM address. And if it is a VRAM use ibv_dmabuf_reg_mr(). Otherwise use ibv_reg_mr()

So something like this:

def get_device_type_for_addr(addr, location):
    if location[:2] == "cpu":
      return "cpu"
   elif location[:3] === "cuda":
      return "cuda"
   else:
      cudaPointerGetAttributes(attribs, addr);
      if attribs.type == cudaMemoryTypeDevice {
         return "cuda"
      } else {
         return "cpu"
      }

#ifndef NO_NVIDIA_PEERMEM

device_type = get_device_type_for_addr(..)
if device_type == "cuda":
    // use ibv_dmabuf_reg_mr(...)
else:
    // use ibv_reg_mr(...)

#endif

// use ibv_reg_mr(...)

@misterwilliam commented on GitHub (Jul 27, 2025): So it sounds like we want to: 1) Use compilations options to guard the change. 2) Detect if the address is a VRAM or DRAM address. And if it is a VRAM use `ibv_dmabuf_reg_mr()`. Otherwise use `ibv_reg_mr()` So something like this: ``` def get_device_type_for_addr(addr, location): if location[:2] == "cpu": return "cpu" elif location[:3] === "cuda": return "cuda" else: cudaPointerGetAttributes(attribs, addr); if attribs.type == cudaMemoryTypeDevice { return "cuda" } else { return "cpu" } #ifndef NO_NVIDIA_PEERMEM device_type = get_device_type_for_addr(..) if device_type == "cuda": // use ibv_dmabuf_reg_mr(...) else: // use ibv_reg_mr(...) #endif // use ibv_reg_mr(...) ```

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@stmatengss commented on GitHub (Jul 28, 2025):

So it sounds like we want to:

Use compilations options to guard the change.
Detect if the address is a VRAM or DRAM address. And if it is a VRAM use ibv_dmabuf_reg_mr(). Otherwise use ibv_reg_mr()

So something like this:

def get_device_type_for_addr(addr, location):
    if location[:2] == "cpu":
      return "cpu"
   elif location[:3] === "cuda":
      return "cuda"
   else:
      cudaPointerGetAttributes(attribs, addr);
      if attribs.type == cudaMemoryTypeDevice {
         return "cuda"
      } else {
         return "cpu"
      }

#ifndef NO_NVIDIA_PEERMEM

device_type = get_device_type_for_addr(..)
if device_type == "cuda":
    // use ibv_dmabuf_reg_mr(...)
else:
    // use ibv_reg_mr(...)

#endif

// use ibv_reg_mr(...)

Yes. Most users use GDR-based VRAM-to-VRAM transfer. Using compilation options can reduce code changes. Welcome to submit a PR

@stmatengss commented on GitHub (Jul 28, 2025): > So it sounds like we want to: > > 1. Use compilations options to guard the change. > 2. Detect if the address is a VRAM or DRAM address. And if it is a VRAM use `ibv_dmabuf_reg_mr()`. Otherwise use `ibv_reg_mr()` > > So something like this: > > ``` > def get_device_type_for_addr(addr, location): > if location[:2] == "cpu": > return "cpu" > elif location[:3] === "cuda": > return "cuda" > else: > cudaPointerGetAttributes(attribs, addr); > if attribs.type == cudaMemoryTypeDevice { > return "cuda" > } else { > return "cpu" > } > > #ifndef NO_NVIDIA_PEERMEM > > device_type = get_device_type_for_addr(..) > if device_type == "cuda": > // use ibv_dmabuf_reg_mr(...) > else: > // use ibv_reg_mr(...) > > #endif > > // use ibv_reg_mr(...) > ``` Yes. Most users use GDR-based VRAM-to-VRAM transfer. Using compilation options can reduce code changes. Welcome to submit a PR

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@misterwilliam commented on GitHub (Jul 29, 2025):

Sounds good. I expect I will get it ready in the next couple of days.

@misterwilliam commented on GitHub (Jul 29, 2025): Sounds good. I expect I will get it ready in the next couple of days.

bsakaguchi commented

2026-10-29 21:11:01 +00:00

Author

Owner

@misterwilliam commented on GitHub (Jul 31, 2025):

I made a pull request. Some things worth highlighting.

Since the location isn't passed into the relevant code I just auto detect if it is host memory or GPU memory and do the appropriate thing.
ibv_dmabuf_reg_mr() returns a ibv_mr struct with addr set to the offset link. There is existing code that assumes ibv_mr->addr is always set to the start of the memory region. I created a new struct MrMeta to track this info since we can't rely on ibv_mr to track it and updated all existing usage to use MrMeta.

@misterwilliam commented on GitHub (Jul 31, 2025): I made a pull request. Some things worth highlighting. 1. Since the location isn't passed into the relevant code I just auto detect if it is host memory or GPU memory and do the appropriate thing. 2. `ibv_dmabuf_reg_mr()` returns a `ibv_mr` struct with addr set to the offset [link](https://github.com/linux-rdma/rdma-core/blob/c0d5740d0bfbc0ab70f73ffad15247b8e4f712e8/libibverbs/verbs.c#L428). There is existing code that assumes `ibv_mr->addr` is always set to the start of the memory region. I created a new struct `MrMeta` to track this info since we can't rely on `ibv_mr` to track it and updated all existing usage to use `MrMeta`.

bsakaguchi referenced this issue

2026-10-29 21:27:08 +00:00

[PR #216] [MERGED] chore(scripts): set rpath for shared objects in build script #1401