Add ability to do RDMA without nvidia-peermem by misterwilliam · Pull Request #704 · kvcache-ai/Mooncake

misterwilliam · 2025-07-31T21:05:51Z

Migrate from ibv_reg_mr() to ibv_reg_dmabuf_mr(). This allows us to not need nvidia-peermem for RDMA because ibv_reg_dmabuf_mr() does not require nvidia-peermem on nvidia GPUs while ibv_reg_mr() does require nvidia-peermem.

Need to introduce a new struct (MrMeta) to track starting address of MR because ibv_reg_dmabuf_mr() is iova based and doesn't set mr->address to starting address, so we need to track it ourselves.

Closes #674

stmatengss · 2025-08-01T04:50:05Z


 #include "transport/rdma_transport/rdma_context.h"

+#include <cuda.h>


By default, the CUDA library is not required. It should be included using #ifdef...#endif.

stmatengss · 2025-08-01T04:51:42Z

Additionally, you should note the USE_CXL option. It is recommended to set it to ON when using ibv_reg_dmabuf_mr.

staryxchen · 2025-08-01T02:43:52Z

 option(WITH_METRICS "enable metrics and metrics reporting thread" ON)
 option(USE_3FS "option for using 3FS storage backend" OFF)
-
+option(NO_NVIDIA_PEERMEM "build to support RDMA without nvidia-peermem" OFF)


what about WITH_NVIDIA_PEERMEM?

misterwilliam · 2025-08-01T19:46:42Z

Additionally, you should note the USE_CXL option. It is recommended to set it to ON when using ibv_reg_dmabuf_mr.

Can you elaborate on this? I'm not sure what you are asking me to do. Maybe I am lacking the proper understanding of how CXL relates to dmabuf support. As far as I understand CXL support implies that the underlying host memory access is super fast and so you don't need to bother RDMA. If that is the correct understanding then it would seem to me that my change doesn't have interactions with CXL support.

stmatengss · 2025-08-02T06:39:53Z

Additionally, you should note the USE_CXL option. It is recommended to set it to ON when using ibv_reg_dmabuf_mr.

Can you elaborate on this? I'm not sure what you are asking me to do. Maybe I am lacking the proper understanding of how CXL relates to dmabuf support. As far as I understand CXL support implies that the underlying host memory access is super fast and so you don't need to bother RDMA. If that is the correct understanding then it would seem to me that my change doesn't have interactions with CXL support.

Sorry for misleading. USE_CXL should be USE_CUDA.

staryxchen · 2025-08-03T08:27:10Z

-    if (!mr) {
+
+    MrMeta mrMeta;
+#ifndef WITH_NVIDIA_PEERMEM


Consider the case where WITH_NVIDIA_PEERMEM is not defined and USE_CUDA is also not defined.

Maybe just adding a comment is sufficient? I added comment in doc string for the WITH_NVIDIA_PEERMEM option that mentions that USE_CUDA=ON is required.

The new code allows someone to run RDMA on an NVIDIA GPU without installing the kernel module nvidia-peermem. This code path would break if you tried running it on a non-nvidia machine without CUDA. Therefore the only drawback of the current situation is if someone wanted run the same binary on nvidia GPUs without nvidia-peermem and also non-nvidia GPUs (say AMD GPUs). To support that case we would have to implement the AMD method of constructing the dmabuf or add some code that detects what GPU type we are and only run through the dmabuf code if we are on a nvidia and fallback to non-dmabuf code otherwise. Adding AMD support is beyond the scope of the effort (also I don't have access to AMD GPUs to test the code). And adding the code that checks if we are running on an nvidia device and falling back to non-dmabuf if we are not on an nvidia device isn't worth the added complexity?

My concern is that if USE_CUDA is not defined, we can always call ibv_reg_mr without worrying about whether nvidia-peermem is supported. However, in the current implementation, if nvidia-peermem is not supported, USE_CUDA must be defined (otherwise the CUDA API cannot be used). For users who only use DRAM, this seems to result in incompatible behavior.

I changed the compilation guard logic from ifnfdef WITH_NVIDIA_PEERMEM to if !defined(WITH_NVIDIA_PEERMEM) && defined(USE_CUDA). This means that if someone compiles the binary with USE_CUDA=OFF and WITH_NVIDIA_PEERMEM=OFF the binary with just use ibv_reg_mr() without needing CUDA. So if someone runs this binary just using DRAM everything will work.

I think that addresses your concern. Feel free to let me know if it doesn't.

staryxchen · 2025-08-03T08:29:29Z


    for (auto &entry : memory_region_list_) {
-        int ret = ibv_dereg_mr(entry);
+        int ret = ibv_dereg_mr(entry.mr);


When deregistering MR, do we consider whether the memory is VRAM or DRAM?

As far as I know it doesn't matter. I've also tested with pytorch and my current implementation works.

# NVIDA cuda >>> gt = torch.tensor([[3., -2.], [2., -2.]], device=torch.device("cuda", 0)) >>> engine.register_memory(gt.data_ptr(), gt.nbytes) 0 >>> engine.unregister_memory(gt.data_ptr()) 0 # CPU host memory >>> t = torch.tensor([[2., -2.], [2., -2.]]) >>> engine.register_memory(t.data_ptr(), t.nbytes) 0 >>> engine.unregister_memory(t.data_ptr()) 0

staryxchen · 2025-08-03T08:30:05Z

            return ERR_CONTEXT;
        }

+#ifndef WITH_NVIDIA_PEERMEM


same as Line R209

staryxchen · 2025-08-03T08:31:54Z

    volatile int outstanding;
 };

+struct MrMeta {


I think it would be easier to understand if the full name were used here. Like MemoryRegionMeta?

misterwilliam · 2025-08-04T03:05:15Z

Additionally, you should note the USE_CXL option. It is recommended to set it to ON when using ibv_reg_dmabuf_mr.

Can you elaborate on this? I'm not sure what you are asking me to do. Maybe I am lacking the proper understanding of how CXL relates to dmabuf support. As far as I understand CXL support implies that the underlying host memory access is super fast and so you don't need to bother RDMA. If that is the correct understanding then it would seem to me that my change doesn't have interactions with CXL support.

Sorry for misleading. USE_CXL should be USE_CUDA.

Ah. That makes sense. I added a comment that if WITH_NVIDIA_PEERMEM=OFF then USE_CUDA=ON is required.

stmatengss · 2025-08-04T08:18:19Z

Everything appears to be in order. Minor note: the format check wasn't passed. Please adhere to the .clang-format coding style and consider running 'clang-format -i FILENAME'. Thanks!

Migrate from ibv_reg_mr() to ibv_reg_dmabuf_mr(). This allows us to not need nvidia-peermem for RDMA because ibv_reg_dmabuf_mr() does not require nvidia-peermem on nvidia GPUs while ibv_reg_mr() does require nvidia-peermem. Need to introduce a new struct (MrMeta) to track starting address of MR because ibv_reg_dmabuf_mr() is iova based and doesn't set mr->address to starting address, so we need to track it ourselves.

misterwilliam · 2025-08-04T18:17:52Z

Everything appears to be in order. Minor note: the format check wasn't passed. Please adhere to the .clang-format coding style and consider running 'clang-format -i FILENAME'. Thanks!

Done

stmatengss · 2025-08-07T11:46:45Z

@alogfans PTAL. thanks

stmatengss · 2025-09-03T12:01:16Z

sgl-project/sglang#9950, @misterwilliam Could you solve this problem?

misterwilliam · 2025-09-03T15:56:29Z

I added some comments. I don't have access to a machine with ROCm so it would be really hard for me to implement the change correctly.

Migrate from ibv_reg_mr() to ibv_reg_dmabuf_mr(). This allows us to not need nvidia-peermem for RDMA because ibv_reg_dmabuf_mr() does not require nvidia-peermem on nvidia GPUs while ibv_reg_mr() does require nvidia-peermem. Need to introduce a new struct (MrMeta) to track starting address of MR because ibv_reg_dmabuf_mr() is iova based and doesn't set mr->address to starting address, so we need to track it ourselves.

@misterwilliam

Fixes kvcache-ai#751. Adds a parallel `#elif defined(USE_HIP)` branch in RdmaContext::registerMemoryRegionInternal that mirrors the existing CUDA dmabuf path (added by kvcache-ai#704) using ROCm's `hsa_amd_portable_export_dmabuf()` instead of `cuMemGetHandleForAddressRange(...DMA_BUF_FD...)`. This lets Mooncake register AMD GPU memory for RDMA without requiring an nvidia-peermem-equivalent kernel module — the path UCX's ROCm backend (uct/rocm/base/rocm_base.c) already uses successfully. Same host-vs-device split as the CUDA branch: `hipPointerGetAttributes` detects host memory and falls back to `ibv_reg_mr`; device/managed memory goes through the dmabuf path. `hipMemGetAddressRange` is used to get the true allocation base because `addr` may sit at an offset within a larger hipMalloc block (caching allocators pack tensors). CMake: added `hsa-runtime64` to the HIP link line in mooncake-transfer-engine/src/CMakeLists.txt. Validation: - Standalone dmabuf probe verified PASS on: * AMD MI355X (gfx950) + Pensando ionic + ROCm 7.2.2 * AMD MI300X (gfx942) + Broadcom Thor2 (bnxt_re) + ROCm 7.0.2 Probe source + container recipe: https://github.com/andyluo7/dynamo/blob/amd-poc-consumer-polish/amd-mi355x-poc/advanced/debug-probes/dmabuf_register_probe.cpp - Standalone compile check confirms all HIP/HSA/ibverbs symbols in the new branch resolve and link cleanly with hsa-runtime64 + libibverbs. End-to-end SGLang+Mooncake disagg validation (T3) on MI355X+ionic will follow in a comment once a full Mooncake build with submodules completes. CC @misterwilliam @stmatengss @alogfans (active on kvcache-ai#751) Closes kvcache-ai#751 Signed-off-by: Andy Luo <anluo@amd.com>

@misterwilliam

…ixes #751) (#2225) * [TransferEngine][ROCm] Add HIP dmabuf MR registration for AMD GPUs Fixes #751. Adds a parallel `#elif defined(USE_HIP)` branch in RdmaContext::registerMemoryRegionInternal that mirrors the existing CUDA dmabuf path (added by #704) using ROCm's `hsa_amd_portable_export_dmabuf()` instead of `cuMemGetHandleForAddressRange(...DMA_BUF_FD...)`. This lets Mooncake register AMD GPU memory for RDMA without requiring an nvidia-peermem-equivalent kernel module — the path UCX's ROCm backend (uct/rocm/base/rocm_base.c) already uses successfully. Same host-vs-device split as the CUDA branch: `hipPointerGetAttributes` detects host memory and falls back to `ibv_reg_mr`; device/managed memory goes through the dmabuf path. `hipMemGetAddressRange` is used to get the true allocation base because `addr` may sit at an offset within a larger hipMalloc block (caching allocators pack tensors). CMake: added `hsa-runtime64` to the HIP link line in mooncake-transfer-engine/src/CMakeLists.txt. Validation: - Standalone dmabuf probe verified PASS on: * AMD MI355X (gfx950) + Pensando ionic + ROCm 7.2.2 * AMD MI300X (gfx942) + Broadcom Thor2 (bnxt_re) + ROCm 7.0.2 Probe source + container recipe: https://github.com/andyluo7/dynamo/blob/amd-poc-consumer-polish/amd-mi355x-poc/advanced/debug-probes/dmabuf_register_probe.cpp - Standalone compile check confirms all HIP/HSA/ibverbs symbols in the new branch resolve and link cleanly with hsa-runtime64 + libibverbs. End-to-end SGLang+Mooncake disagg validation (T3) on MI355X+ionic will follow in a comment once a full Mooncake build with submodules completes. CC @misterwilliam @stmatengss @alogfans (active on #751) Closes #751 --------- Signed-off-by: Andy Luo <anluo@amd.com> Signed-off-by: Andy Luo <andy.luo@amd.com> Co-authored-by: Claude Sonnet 4 <noreply@anthropic.com>

@misterwilliam

…ixes kvcache-ai#751) (kvcache-ai#2225) * [TransferEngine][ROCm] Add HIP dmabuf MR registration for AMD GPUs Fixes kvcache-ai#751. Adds a parallel `#elif defined(USE_HIP)` branch in RdmaContext::registerMemoryRegionInternal that mirrors the existing CUDA dmabuf path (added by kvcache-ai#704) using ROCm's `hsa_amd_portable_export_dmabuf()` instead of `cuMemGetHandleForAddressRange(...DMA_BUF_FD...)`. This lets Mooncake register AMD GPU memory for RDMA without requiring an nvidia-peermem-equivalent kernel module — the path UCX's ROCm backend (uct/rocm/base/rocm_base.c) already uses successfully. Same host-vs-device split as the CUDA branch: `hipPointerGetAttributes` detects host memory and falls back to `ibv_reg_mr`; device/managed memory goes through the dmabuf path. `hipMemGetAddressRange` is used to get the true allocation base because `addr` may sit at an offset within a larger hipMalloc block (caching allocators pack tensors). CMake: added `hsa-runtime64` to the HIP link line in mooncake-transfer-engine/src/CMakeLists.txt. Validation: - Standalone dmabuf probe verified PASS on: * AMD MI355X (gfx950) + Pensando ionic + ROCm 7.2.2 * AMD MI300X (gfx942) + Broadcom Thor2 (bnxt_re) + ROCm 7.0.2 Probe source + container recipe: https://github.com/andyluo7/dynamo/blob/amd-poc-consumer-polish/amd-mi355x-poc/advanced/debug-probes/dmabuf_register_probe.cpp - Standalone compile check confirms all HIP/HSA/ibverbs symbols in the new branch resolve and link cleanly with hsa-runtime64 + libibverbs. End-to-end SGLang+Mooncake disagg validation (T3) on MI355X+ionic will follow in a comment once a full Mooncake build with submodules completes. CC @misterwilliam @stmatengss @alogfans (active on kvcache-ai#751) Closes kvcache-ai#751 --------- Signed-off-by: Andy Luo <anluo@amd.com> Signed-off-by: Andy Luo <andy.luo@amd.com> Co-authored-by: Claude Sonnet 4 <noreply@anthropic.com>

stmatengss reviewed Aug 1, 2025

View reviewed changes

staryxchen reviewed Aug 1, 2025

View reviewed changes

misterwilliam force-pushed the main branch from 9a70a6b to b433a20 Compare August 1, 2025 19:33

staryxchen reviewed Aug 3, 2025

View reviewed changes

misterwilliam force-pushed the main branch from b433a20 to 4c8eddc Compare August 4, 2025 02:49

misterwilliam force-pushed the main branch from 4c8eddc to 1aa9d3c Compare August 4, 2025 03:06

misterwilliam force-pushed the main branch from 1aa9d3c to d68b4e8 Compare August 4, 2025 18:05

stmatengss requested a review from alogfans August 5, 2025 07:00

stmatengss approved these changes Aug 9, 2025

View reviewed changes

stmatengss merged commit 1350b74 into kvcache-ai:main Aug 9, 2025
11 checks passed

stmatengss mentioned this pull request Aug 18, 2025

[Bug]: Mooncake with SGLang on ROCm failed to register memory #751

Closed

1 task

andyluo7 mentioned this pull request May 25, 2026

[TransferEngine][ROCm] Add HIP dmabuf MR registration for AMD GPUs (fixes #751) #2225

Merged


		#include "transport/rdma_transport/rdma_context.h"

		#include <cuda.h>

Uh oh!

Conversation

misterwilliam commented Jul 31, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stmatengss commented Aug 1, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

misterwilliam commented Aug 1, 2025

Uh oh!

stmatengss commented Aug 2, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

misterwilliam commented Aug 4, 2025

Uh oh!

stmatengss commented Aug 4, 2025

Uh oh!

misterwilliam commented Aug 4, 2025

Uh oh!

stmatengss commented Aug 7, 2025

Uh oh!

Uh oh!

stmatengss commented Sep 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

misterwilliam commented Sep 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

stmatengss commented Sep 3, 2025 •

edited

Loading