pytorch: guard cusolvermp/newton_schulz pybind + decl on NVTE_WITH_CUSOLVERMP

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Author: phu0ngng

transformer_engine/pytorch/csrc/extensions.h

@@ -631,12 +631,14 @@ void bulk_overlap_ag_with_external_gemm(CommOverlap &allgather_communicator, at:
  * Newton-Schulz (cuSolverMp)
  **************************************************************************************************/
 
+#ifdef NVTE_WITH_CUSOLVERMP
 int64_t cusolvermp_ctx_create(int64_t nccl_comm_ptr, int nranks, int rank);
 
 void cusolvermp_ctx_destroy(int64_t ctx_ptr);
 
 void newton_schulz(int64_t ctx_ptr, int64_t m, int64_t n, at::Tensor x, int64_t num_iterations,
                    std::vector<float> coefficients);
+#endif  // NVTE_WITH_CUSOLVERMP
 
 }  // namespace transformer_engine::pytorch
 

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -589,6 +589,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         &transformer_engine::pytorch::multi_tensor_compute_scale_inv_e8m0_cuda,
         "Fused compute E8M0 scale_inv from amax", py::call_guard<py::gil_scoped_release>());
 
+#ifdef NVTE_WITH_CUSOLVERMP
   // Newton-Schulz (cuSolverMp)
   m.def("cusolvermp_ctx_create", &transformer_engine::pytorch::cusolvermp_ctx_create,
         "Create cuSolverMp context for Newton-Schulz", py::arg("nccl_comm_ptr"), py::arg("nranks"),
@@ -599,6 +600,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Newton-Schulz matrix orthogonalization", py::arg("ctx_ptr"), py::arg("m"), py::arg("n"),
         py::arg("x"), py::arg("num_iterations"), py::arg("coefficients"),
         py::call_guard<py::gil_scoped_release>());
+#endif  // NVTE_WITH_CUSOLVERMP
 
   // Comm+GEMM Overlap
   m.def("bulk_overlap_ag_with_external_gemm",