[feat] support GDR in mooncake backend#131
Conversation
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
|
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
|
What will happen if user use TCP transport and GDR on? |
There was a problem hiding this comment.
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_utilswith 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-pythonto 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 |
| 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() |
| staging = self._gdr_staging | ||
| buffer_size = staging.size | ||
|
|
||
| # Caller guarantees all tensors are CUDA; make contiguous outside the lock. |
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. |
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
| # Ensure all pending GPU work in this process is done before the staging stream reads. | ||
| torch.cuda.synchronize() |
There was a problem hiding this comment.
The current_stream event approach is only safe when all tensors in the batch were produced on the caller's current stream at call time — a guarantee we cannot enforce or verify, since PyTorch provides no API to query which stream produced a given tensor.
torch.cuda.synchronize() is the only correct option here; in RL workloads this sync lands at a natural computation boundary where the device is already nearly idle, so the practical overhead is negligible.
| if (gpu_tensor_indices and use_gdr_path) or non_tensor_indices: | ||
| if custom_backend_meta is None or len(custom_backend_meta) != len(keys): | ||
| raise ValueError( | ||
| "custom_backend_meta is required when GDR is enabled (for n_chunks) " | ||
| "or when any dtype is None (for packed_size)." | ||
| ) |
There was a problem hiding this comment.
At get() time there is no way to determine which keys were stored as oversized chunks without reading n_chunks from custom_backend_meta — the information simply does not exist elsewhere.
Making custom_backend_meta optional would silently treat a missing entry as "no chunking", producing corrupted results if the key was actually split; a strict validation that fails loudly is safer than a relaxed contract that corrupts data.
Signed-off-by: xupinjie <xupinjie321@outlook.com>
CLA Signature Guide@xupinjie , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (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 |
Add GDR support for mooncake backend