Replace nvshmem#187
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR removes the NVSHMEM-based symmetric-memory pathway and replaces the fused cubed-sphere CUDA exchange mechanism with a UCX-backed process-group exchange of a packed edge-state buffer, updating build/test gating and documentation accordingly.
Changes:
- Reworked fused cubed-sphere exchange to use a per-process packed buffer plus UCX
send/recvinstead of PyTorch symmetric memory rendezvous/signal pads. - Updated the CUDA smoke test and CMake test registration to require UCX and validate the UCX process group.
- Removed NVSHMEM build plumbing and documented UCX usage for the fused cubed-sphere CUDA path.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_fused_cs_ucx_smoke.cu | Replaces symmetric-memory smoke coverage with UCX-based exchange checks and verification kernels. |
| tests/CMakeLists.txt | Renames/retargets the fused cubed-sphere smoke test and gates it on UCX_FOUND. |
| src/hydro/hydro_options.cpp | Updates comments to reflect the new packed-buffer exchange approach. |
| src/hydro/hydro_forward_fused.cpp | Implements UCX exchange logic, pool/buffer pointer management, and updates fused forward path accordingly. |
| src/hydro/fused_recon_riemann_dispatch.hpp | Updates public dispatch signatures to remove symmetric-memory sync/release phases and rename pointer fields. |
| src/hydro/fused_cubed_sphere_dispatch.cu | Removes symmetric-memory sync/release kernels and adapts pack/flux entrypoints to use the exchange buffer. |
| src/CMakeLists.txt | Drops NVSHMEM include/link directives from the CUDA library target. |
| docs/user_guide/distributed.rst | Documents that fused cubed-sphere CUDA exchange now uses the UCX backend and no longer needs symmetric-memory runtime. |
| CMakeLists.txt | Removes inclusion of the NVSHMEM CMake module. |
| cmake/nvshmem.cmake | Deletes the NVSHMEM discovery/validation module. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+267
to
+271
| if (peer_process != layout.options->process_rank()) { | ||
| remote_peers->push_back(peer_process); | ||
| } | ||
| bool dim_enabled = side <= SIDE_R ? exchange_dim2 : exchange_dim3; | ||
| if (!dim_enabled) continue; |
Comment on lines
+139
to
+153
| std::vector<CommWorkPtr> works; | ||
| works.reserve(2 * peers.size()); | ||
| for (int peer : peers) { | ||
| std::vector<torch::Tensor> send_tensors{pool.buffer}; | ||
| std::vector<torch::Tensor> recv_tensors{pool.peer_buffers[peer]}; | ||
| if (auto work = | ||
| comm->send(send_tensors, peer, kFusedCubedSphereExchangeTag)) { | ||
| works.push_back(std::move(work)); | ||
| } | ||
| if (auto work = | ||
| comm->recv(recv_tensors, peer, kFusedCubedSphereExchangeTag)) { | ||
| works.push_back(std::move(work)); | ||
| } | ||
| } | ||
| for (auto const& work : works) work->wait(); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.