[PyTorch] Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy

build_tools/pytorch.py

@@ -77,6 +77,16 @@ def setup_pytorch_extension(
 
     setup_mpi_flags(include_dirs, cxx_flags)
 
+    # Mirror the NCCL EP gate from setup.py / common CMake. When disabled, the
+    # ep.cpp source no-ops at the #ifdef boundary; without the define it would
+    # produce undefined references to nvte_ep_*.
+    if bool(int(os.getenv("NVTE_BUILD_WITH_NCCL_EP", "1"))):
+        cxx_flags.append("-DNVTE_WITH_NCCL_EP")
+        # PyTorch's symm-mem headers gate the NCCL_HAS_SYMMEM_* feature macros on
+        # USE_NCCL. The EP extension shares the symm-mem NCCL comm with torch, so
+        # it needs those macros visible.
+        cxx_flags.append("-DUSE_NCCL")
+
     library_dirs = []
     libraries = []
     if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", 0))):