[feat] support GDR in mooncake backend by xupinjie · Pull Request #131 · Ascend/TransferQueue

xupinjie · 2026-06-30T08:27:24Z

Add GDR support for mooncake backend

ascend-robot · 2026-06-30T08:27:34Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`c572550` [feat] support GDR in mooncake ...](`c572550`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

xupinjie · 2026-06-30T08:30:49Z

ascend-robot · 2026-06-30T09:16:39Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`901d4b8` [feat] support GDR in mooncake ...](`901d4b8`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

ascend-robot · 2026-06-30T09:18:58Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`8297b8f` [feat] support GDR in mooncake ...](`8297b8f`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

ascend-robot · 2026-06-30T09:34:58Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`8c7637a` [feat] support GDR in mooncake ...](`8c7637a`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

ascend-robot · 2026-06-30T11:44:39Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`e861d15` [feat] support GDR in mooncake ...](`e861d15`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

0oshowero0 · 2026-07-03T06:54:57Z

+            from cuda import cuda as cuda_driver
+            from cuda import cudart


Will this still work on cuda-python 13.x?

Valid concern. The current imports use the legacy cuda.cuda/cuda.cudart API which was removed in 13.x. We'll fix this by trying cuda.bindings (13.x/new API) first and falling back to cuda.cuda/cuda.cudart (12.x).

0oshowero0 · 2026-07-03T06:58:26Z

What will happen if user use TCP transport and GDR on?

Copilot

Pull request overview

Adds GPUDirect RDMA (GDR) support to the MooncakeStore backend by introducing a persistent, pre-registered CUDA staging buffer and integrating a GDR-aware PUT/GET/CLEAR data path for tensors (including chunking for oversized tensors).

Changes:

Introduce transfer_queue.utils.mooncake_utils with staging-buffer management (GdrStaging) and helper utilities for aligned packing, grouping, and chunked sub-keys.
Add a GDR-enabled tensor transfer path to MooncakeStoreClient (PUT/GET), including oversized tensor chunking via :c{i} sub-keys and CLEAR expansion to remove sub-keys.
Extend unit/e2e/perf tooling and config to exercise and benchmark the new GDR mode, and add cuda-python to the Mooncake optional dependency set.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
transfer_queue/utils/mooncake_utils.py	New GDR utilities: aligned layout, chunked sub-keys, grouping, and persistent CUDA staging buffer implementation.
transfer_queue/storage/clients/mooncake_client.py	Integrates GDR staging into MooncakeStoreClient PUT/GET/CLEAR, including chunked oversized tensor handling and staging lifecycle management.
tests/test_mooncake_utils.py	New unit tests for the mooncake GDR utilities and MooncakeStoreClient.clear() chunk-subkey expansion behavior.
tests/e2e/test_e2e_lifecycle_consistency.py	Adds a GDR-enabled Mooncake backend config and a CUDA-only E2E roundtrip scenario for GPU tensors.
scripts/performance_test/perftest.py	Adds GDR awareness to perf test harness (device setup + GET timing semantics).
scripts/performance_test/perftest_config.yaml	Documents/configures MooncakeStore `use_gdr` flag in perf test config.
pyproject.toml	Adds `cuda-python` to the `mooncake` optional dependency group.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+        # GdrStaging instance created eagerly but cudaMalloc is deferred to first use.
+        # Skip GDR if CUDA context is not initialized in this process (e.g. CPU-only workers)
+        gdr_eligible = self.use_gdr and buffer_bytes > 0 and torch.cuda.is_initialized()
+        self._gdr_staging: GdrStaging | None = GdrStaging(buffer_bytes) if gdr_eligible else None



+        gdr_meta: dict[str, dict | None] = {}
+        if use_gdr_path and tensor_keys:
+            gdr_meta = dict(zip(tensor_keys, self._put_tensors_gdr(tensor_keys, tensor_values), strict=False))


        custom_backend_meta: list[dict | None] = [
-            {"packed_size": next(sizes_iter)} if not isinstance(value, torch.Tensor) else None for value in values
+            gdr_meta.get(key) if isinstance(value, torch.Tensor) else {"packed_size": next(sizes_iter)}
+            for key, value in zip(keys, values, strict=False)


+                    for i, sub_key in enumerate(sub_keys):
+                        chunk_size = min(buffer_size, total - i * buffer_size)
+                        self._batch_get_into_with_retry([sub_key], [staging.ptr], [chunk_size])
+                        staging.memcpy_d2d_async(final_tensor.data_ptr() + i * buffer_size, staging.ptr, chunk_size)
+                    staging.synchronize()


+                else:
+                    # Normal group: aligned total fits; batch_get_into then unpack.
+                    offsets, _ = _aligned_offsets(g_nbytes)
+                    sub_ptrs = [staging._ptr + off for off in offsets]


+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGdrStagingCuda:


+        staging = self._gdr_staging
+        buffer_size = staging.size
+
+        # Caller guarantees all tensors are CUDA; make contiguous outside the lock.


0oshowero0 · 2026-07-03T12:47:07Z

    # Network device name.
    # Set to "" to let Mooncake auto-select available devices.
    device_name: ""
+    # GPU Direct RDMA. When true, CUDA tensors are transferred directly from GPU memory


Please also modify the main config

Please also modify the main config

Done.

xupinjie · 2026-07-07T08:04:22Z

What will happen if user use TCP transport and GDR on?

Valid. GDR requires RDMA transport. We'll add a validation in init that raises a clear ValueError (or at minimum a warning + auto-disable) when use_gdr=True and protocol != "rdma", so users get a clear error message immediately rather than a cryptic failure at transfer time.

ascend-robot · 2026-07-07T08:31:19Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`94cf818` [feat] support GDR in mooncake ...](`94cf818`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

Signed-off-by: xupinjie <xupinjie321@outlook.com>

ascend-robot · 2026-07-07T08:37:49Z

CLA Signature Guide

@xupinjie , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[`20cd384` [feat] support GDR in mooncake ...](`20cd384`)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

0oshowero0 · 2026-07-08T04:08:10Z

+                "local_buffer_size": 134217728,  # 128MB
+                "metadata_server": "localhost:50050",
+                "master_server_address": "localhost:50051",
+                "protocol": "tcp",


This test still using TCP+GDR path.

Currently we lack of GPU CI resources, so maybe we need to regularly run the GDR ci manually..

0oshowero0 · 2026-07-08T04:11:04Z

+        # GdrStaging instance created eagerly but cudaMalloc is deferred to first use.
+        # Skip GDR if CUDA context is not initialized in this process (e.g. CPU-only workers)
+        gdr_eligible = self.use_gdr and buffer_bytes > 0 and torch.cuda.is_initialized()
+        self._gdr_staging: GdrStaging | None = GdrStaging(buffer_bytes) if gdr_eligible else None


I think we should provide a simple GDR doc that describes the required dependencies, design choices (e.g. when there is no cuda context we fall back to rdma)

0oshowero0 · 2026-07-08T05:53:13Z

+        self._stream: torch.cuda.Stream | None = None
+        self._initialized = False
+
+    def lazy_init(self, store) -> None:


Is there any race condition risks for lazy_init? If user process has different threads that call put or get at the same time, there might be problems

0oshowero0 · 2026-07-08T05:56:33Z

+            self._initialized = False
+
+    @contextlib.contextmanager
+    def acquire(self):


This is another design choice that we may need to point out in doc: the current single staging buffer design is for RL scenario where we don't have many concurrent calls that requires a shared client. So a single buffer is enough, and is also good for saving HBM

0oshowero0 · 2026-07-08T05:57:44Z

    },
+    # MooncakeStore with GDR staging buffer enabled.
+    # GPU tensors are transferred via the persistent staging buffer path;
+    # CPU tensors and non-tensors fall through to the original CPU RDMA path.


This should be align with current implementation

And we should point out in the doc

ascend-robot added the ascend-cla/no label Jun 30, 2026

xupinjie force-pushed the pinjie/GDR_v2 branch from c572550 to 901d4b8 Compare June 30, 2026 09:16

xupinjie force-pushed the pinjie/GDR_v2 branch from 901d4b8 to 8297b8f Compare June 30, 2026 09:18

xupinjie mentioned this pull request Jun 30, 2026

[feat] Add GPU Direct RDMA (GDR) support for MooncakeStore backend. #79

Closed

xupinjie force-pushed the pinjie/GDR_v2 branch from 8297b8f to 8c7637a Compare June 30, 2026 09:34

xupinjie force-pushed the pinjie/GDR_v2 branch from 8c7637a to e861d15 Compare June 30, 2026 11:44

0oshowero0 reviewed Jul 3, 2026

View reviewed changes

0oshowero0 requested a review from Copilot July 3, 2026 07:00

Copilot started reviewing on behalf of 0oshowero0 July 3, 2026 07:00 View session

Copilot AI reviewed Jul 3, 2026

View reviewed changes

0oshowero0 reviewed Jul 3, 2026

View reviewed changes

xupinjie force-pushed the pinjie/GDR_v2 branch from e861d15 to 94cf818 Compare July 7, 2026 08:31

[feat] support GDR in mooncake backend

20cd384

Signed-off-by: xupinjie <xupinjie321@outlook.com>

xupinjie force-pushed the pinjie/GDR_v2 branch from 94cf818 to 20cd384 Compare July 7, 2026 08:37

0oshowero0 reviewed Jul 8, 2026

View reviewed changes

		@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
		class TestGdrStagingCuda:

Uh oh!

Conversation

xupinjie commented Jun 30, 2026

Uh oh!

ascend-robot commented Jun 30, 2026

CLA Signature Guide

Uh oh!

xupinjie commented Jun 30, 2026

Uh oh!

ascend-robot commented Jun 30, 2026

CLA Signature Guide

Uh oh!

ascend-robot commented Jun 30, 2026

CLA Signature Guide

Uh oh!

ascend-robot commented Jun 30, 2026

CLA Signature Guide

Uh oh!

ascend-robot commented Jun 30, 2026

CLA Signature Guide

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

0oshowero0 commented Jul 3, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xupinjie commented Jul 7, 2026

Uh oh!

ascend-robot commented Jul 7, 2026

CLA Signature Guide

Uh oh!

ascend-robot commented Jul 7, 2026

CLA Signature Guide

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants