[c10d][nvshmem] add nvshmem build rules and dependency for libtorch_cuda (#159562)

jushg · facebook-github-bot · commit e7a150d5fa23 · 2025-08-11T12:42:31.000-07:00
Summary:

Add guarded build option for nvshmem-related c10d code with `-c fbcode.caffe2_use_nvshmem`

Guarded clause include nvshmem device + host code (static-linked) + these 2 files:
- `torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu`
-    `torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu`

Differential Revision: D78813568
diff --git a/build_variables.bzl b/build_variables.bzl
@@ -754,6 +754,11 @@ libtorch_cuda_distributed_extra_sources = [
     "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp",
 ]
 
+libtorch_nvshmem_sources = [
+    "torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu",
+    "torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu",
+]
+
 libtorch_cuda_distributed_sources = libtorch_cuda_distributed_base_sources + libtorch_cuda_distributed_extra_sources
 
 libtorch_cuda_sources = libtorch_cuda_core_sources + libtorch_cuda_distributed_sources + [
