[SymmMem] Add multimem support for NCCL and NVSHMEM (pytorch#172185)

kwen2501 · pytorchmergebot · commit 1c83214342d2 · 2026-01-12T23:28:17.000Z
Pull Request resolved: pytorch#172185 Approved by: https://github.com/Skylion007, https://github.com/dzmitry-huba ghstack dependencies: pytorch#172163
diff --git a/test/distributed/test_nccl.py b/test/distributed/test_nccl.py
@@ -411,6 +411,31 @@ def test_mempool_compute_ops(self):
         expected = torch.mm(x, w) * self.world_size
         self.assertEqual(y, expected)
 
+    @skip_but_pass_in_sandcastle_if(TEST_WITH_ROCM, "Skip NCCL tests for ROCm")
+    @skip_but_pass_in_sandcastle_if(IS_WINDOWS, "NCCL doesn't support Windows")
+    @skip_if_lt_x_gpu(2)
+    @requires_nccl_version(
+        (2, 29), "NCCL Symmetric Memory multicast support from nccl 2.29"
+    )
+    def test_multicast_ptr(self) -> None:
+        """
+        Get the multicast pointer
+        """
+        from torch._C._autograd import DeviceType
+        from torch._C._distributed_c10d import _SymmetricMemory
+
+        symm_mem.set_backend("NCCL")
+        torch.cuda.set_device(self.rank)
+        c10d.all_reduce(torch.ones(1, device=self.device))
+        group_name = c10d.group.WORLD.group_name
+
+        tensor = symm_mem.empty(1, device=self.device)
+        handle = symm_mem.rendezvous(tensor, group_name)
+        if _SymmetricMemory.has_multicast_support(DeviceType.CUDA, self.device.index):
+            self.assertNotEqual(handle.multicast_ptr, 0)
+        else:
+            self.assertEqual(handle.multicast_ptr, 0)
+
 
 instantiate_device_type_tests(TestNCCL, globals(), only_for="cuda")
 
diff --git a/test/distributed/test_nvshmem.py b/test/distributed/test_nvshmem.py
@@ -3,6 +3,7 @@
 # To run:
 # python test/distributed/test_nvshmem.py
 
+
 import os
 
 import torch
@@ -31,12 +32,26 @@ def requires_nvshmem():
     )
 
 
+def has_nvls_support():
+    if not symm_mem.is_nvshmem_available():
+        return False
+
+    if os.environ.get("NVSHMEM_DISABLE_NVLS", "0") == "1":
+        return False
+
+    # Set NVSHMEM as SymmMem backend before running the check
+    symm_mem.set_backend("NVSHMEM")
+    from torch._C._autograd import DeviceType
+    from torch._C._distributed_c10d import _SymmetricMemory
+
+    return _SymmetricMemory.has_multicast_support(DeviceType.CUDA, 0)
+
+
 def requires_nvls():
-    """Skip test if NVLS (NVLink Switch) is not available."""
-    nvls_disabled = os.environ.get("NVSHMEM_DISABLE_NVLS", "0") == "1"
+    """Skip test if NVLS (NVLink SHARP) is not available."""
     return skip_but_pass_in_sandcastle_if(
-        nvls_disabled,
-        "Test requires NVLS which is disabled via NVSHMEM_DISABLE_NVLS=1",
+        not has_nvls_support(),
+        "Test requires NVLink SHARP support",
     )
 
 
@@ -225,6 +240,20 @@ def test_get_remote_tensors(self) -> None:
         for peer, tensor in enumerate(remote_tensors):
             self.assertEqual(tensor, peer)
 
+    def test_multicast_ptr(self) -> None:
+        """
+        Get the multicast pointer
+        """
+        self._init_device()
+        group_name = dist.group.WORLD.group_name
+
+        tensor = symm_mem.empty(1, device=self.device)
+        handle = symm_mem.rendezvous(tensor, group_name)
+        if has_nvls_support():
+            self.assertNotEqual(handle.multicast_ptr, 0)
+        else:
+            self.assertEqual(handle.multicast_ptr, 0)
+
     @skipIfRocm
     def test_nvshmem_put(self) -> None:
         self._init_device()
diff --git a/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu b/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.cu
@@ -142,7 +142,7 @@ bool CUDASymmetricMemory::has_multicast_support() {
 }
 
 void* CUDASymmetricMemory::get_multicast_ptr() {
-  return pai_->mc_addr_;
+  return static_cast<char*>(pai_->mc_addr_) + offset_;
 }
 
 size_t CUDASymmetricMemory::get_offset() {
diff --git a/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu b/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu
@@ -138,6 +138,12 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
       arr_size,
       cudaMemcpyDeviceToHost));
 #endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
+    C10D_NCCL_CHECK(
+        ncclGetLsaMultimemDevicePointer(buffer_win_, offset_, &mc_addr_),
+        "Failed to get multicast pointer");
+#endif
   }
 
   // Exact copy is not needed / supported
@@ -159,6 +165,8 @@ class NCCLPeerAllocInfo : public c10::intrusive_ptr_target {
   std::string group_name_;
   ncclWindow_t buffer_win_;
   ncclWindow_t signal_handle_;
+  // Multicast address
+  void* mc_addr_ = nullptr;
 
   friend class NCCLSymmetricMemory;
 };
@@ -195,13 +203,14 @@ size_t NCCLSymmetricMemory::get_buffer_size() {
 }
 
 bool NCCLSymmetricMemory::has_multicast_support() {
-  // TODO
-  return false;
+  return pai_->mc_addr_ != nullptr;
 }
 
 void* NCCLSymmetricMemory::get_multicast_ptr() {
-  // TODO
-  return nullptr;
+  if (!has_multicast_support()) {
+    return nullptr;
+  }
+  return static_cast<char*>(pai_->mc_addr_) + offset_;
 }
 
 void NCCLSymmetricMemory::barrier(int channel, size_t timeout_ms) {
@@ -312,8 +321,7 @@ class NCCLSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
   }
 
   bool has_multicast_support(int device_idx) override {
-    // TODO
-    return false;
+    return device_has_multicast_support(device_idx);
   }
 
   c10::DeviceType supported_device_type() override {
diff --git a/torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu b/torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu
@@ -138,6 +138,10 @@ class NVSHMEMPeerAllocInfo : public c10::intrusive_ptr_target {
         signal_pads_.data(),
         arr_size,
         cudaMemcpyHostToDevice));
+
+    // Initialize multicast address
+    // On unsupported platforms, this API returns a nullptr.
+    mc_addr_ = nvshmemx_mc_ptr(NVSHMEM_TEAM_WORLD, base_ptr_);
   }
 
  private:
@@ -151,6 +155,8 @@ class NVSHMEMPeerAllocInfo : public c10::intrusive_ptr_target {
   void** signal_pads_dev_;
   // Whether the world is within CUDA P2P only, not network
   bool world_within_cuda_p2p_;
+  // Multicast address
+  void* mc_addr_;
 
   friend class NVSHMEMSymmetricMemory;
 };
@@ -206,13 +212,15 @@ class NVSHMEMSymmetricMemory : public SymmetricMemory {
   }
 
   bool has_multicast_support() override {
-    // TODO
-    return false;
+    // On unsupported platforms, this API returns a nullptr.
+    return pai_->mc_addr_ != nullptr;
   }
 
   void* get_multicast_ptr() override {
-    // TODO
-    return nullptr;
+    if (!has_multicast_support()) {
+      return nullptr;
+    }
+    return static_cast<char*>(pai_->mc_addr_) + offset_;
   }
 
   size_t get_offset() override {
@@ -432,9 +440,8 @@ class NVSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
   };
 
   bool has_multicast_support(int device_idx) override {
-    // TODO
-    return false;
-  };
+    return device_has_multicast_support(device_idx);
+  }
 
   c10::DeviceType supported_device_type() override {
     return c10::DeviceType::CUDA;

Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ bool CUDASymmetricMemory::has_multicast_support() {`
`142`	`142`	`}`
`143`	`143`
`144`	`144`	`void* CUDASymmetricMemory::get_multicast_ptr() {`
`145`		`- return pai_->mc_addr_;`
	`145`	`+ return static_cast<char*>(pai_->mc_addr_) + offset_;`
`146`	`146`	`}`
`147`	`147`
`148`	`148`	`size_t CUDASymmetricMemory::get_offset() {`