diff --git a/drivers/cxl/core/port.c b/drivers/cxl/core/port.c
index 7152013b0f23a25e857e904b55d4948ae8ecaa46..2e910003297698fedeb5afdbe363b20baa030405 100644
--- a/drivers/cxl/core/port.c
+++ b/drivers/cxl/core/port.c
@@ -821,6 +821,13 @@ static struct cxl_port *__devm_cxl_add_port(struct device *host,
 	if (rc)
 		goto err;
 
+	/* update memdev endpoint */
+	if (is_cxl_memdev(uport_dev)) {
+		struct cxl_memdev *cxlmd = to_cxl_memdev(uport_dev);
+
+		cxlmd->endpoint = port;
+	}
+
 	rc = device_add(dev);
 	if (rc)
 		goto err;
@@ -1331,7 +1338,6 @@ int cxl_endpoint_autoremove(struct cxl_memdev *cxlmd, struct cxl_port *endpoint)
 
 	get_device(host);
 	get_device(&endpoint->dev);
-	cxlmd->endpoint = endpoint;
 	cxlmd->depth = endpoint->depth;
 	return devm_add_action_or_reset(dev, delete_endpoint, cxlmd);
 }
diff --git a/drivers/cxl/mem.c b/drivers/cxl/mem.c
index e087febf9af047c81dfb11b06d91d92d442d586c..e9e4a4a796b2da70087a1e31d8126ee09fa8a488 100644
--- a/drivers/cxl/mem.c
+++ b/drivers/cxl/mem.c
@@ -174,14 +174,6 @@ static int cxl_mem_probe(struct device *dev)
 	if (rc)
 		return rc;
 
-	if (resource_size(&cxlds->pmem_res) && IS_ENABLED(CONFIG_CXL_PMEM)) {
-		rc = devm_cxl_add_nvdimm(cxlmd);
-		if (rc == -ENODEV)
-			dev_info(dev, "PMEM disabled by platform\n");
-		else
-			return rc;
-	}
-
 	/*
 	 * The kernel may be operating out of CXL memory on this device,
 	 * there is no spec defined way to determine whether this device
diff --git a/drivers/cxl/port.c b/drivers/cxl/port.c
index 6240e05b95424b9c9627fcde4a0e50171bde1761..e10b5ad5fd6f7b73cd4f82027086d9411f8e5805 100644
--- a/drivers/cxl/port.c
+++ b/drivers/cxl/port.c
@@ -120,6 +120,12 @@ static int cxl_endpoint_port_probe(struct cxl_port *port)
 	if (rc)
 		return rc;
 
+	if (resource_size(&cxlds->pmem_res) && IS_ENABLED(CONFIG_CXL_PMEM)) {
+		rc = devm_cxl_add_nvdimm(cxlmd);
+		if (rc == -ENODEV)
+			dev_info(cxlds->dev, "PMEM disabled by platform\n");
+	}
+
 	/*
 	 * This can't fail in practice as CXL root exit unregisters all
 	 * descendant ports and that in turn synchronizes with cxl_port_probe()