linting

Signed-off-by: Alp Dener <adener@nvidia.com>

Author: denera

3rdparty/cudnn-frontend

@@ -24,7 +24,7 @@
     MXFP8Quantizer,
 )
 import transformer_engine.pytorch.cpp_extensions as tex
-from transformer_engine.pytorch.cpp_extensions.gemm import get_cublas_workspace_size_bytes
+
 from transformer_engine.pytorch.module.base import fill_userbuffers_buffer_for_all_gather
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)

tests/pytorch/distributed/run_gemm_with_overlap.py

@@ -84,8 +84,11 @@ CommOverlapCore::CommOverlapCore(int64_t nccl_comm_ptr, int tp_rank, int tp_size
       "Comm+GEMM overlap with cuBLASMp backend requires TE to be built with NVTE_WITH_CUBLASMP=1.");
   _with_cublasmp = true;
 
-  nvte_comm_gemm_ctx_create(reinterpret_cast<ncclComm_t>(nccl_comm_ptr), tp_size, tp_rank);
+  _cublasmp_ctx = nvte_comm_gemm_ctx_create(reinterpret_cast<ncclComm_t>(nccl_comm_ptr), tp_size,
+                                            tp_rank);
 
+  _tp_id = tp_rank;
+  _tp_size = tp_size;
   _num_comm_sm = num_comm_sm;
   _is_p2p = is_p2p;
   _atomic_gemm = atomic_gemm;

transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp

@@ -331,7 +331,6 @@ def add_ub(
         comm_priority: int = 0,
         gemm_priority: int = 0,
         pipeline_rs_overlap_first_gemm: bool = False,
-        with_cublasmp: bool = False,
     ) -> None:
         if atomic_gemm:
             warnings.warn(
@@ -506,7 +505,7 @@ def fill_userbuffers_buffer_for_all_gather(
     """
     # cuBlasMp already handles its own buffer filling and quantization factors
     if comm.with_cublasmp():
-        return
+        return local_tensor, local_tensor
 
     # Tensor dimensions
     local_shape = local_tensor.size()