[DCU] remove redundant gemm

transformer_engine/common/gemm/cublaslt_gemm.cu

@@ -1166,82 +1166,13 @@ void nvte_cublas_batchgemm(const NVTETensor A, const NVTETensor B, NVTETensor D,
             stream);
 }
 
-// add for batchgemm
-void nvte_cublas_batchgemm_v2(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias,
-                      NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,
-                      NVTETensor workspace, bool accumulate, bool use_split_accumulator,
-                      int math_sm_count, int batch_count, cudaStream_t stream) {
-  NVTE_API_CALL(nvte_cublas_batchgemm_v2);
-  using namespace transformer_engine;
-  const Tensor *inputA = convertNVTETensorCheck(A);
-  const Tensor *inputB = convertNVTETensorCheck(B);
-  Tensor *outputD = convertNVTETensor(D);
-  const Tensor *biasTensor = convertNVTETensor(bias);
-  Tensor *outputGelu = convertNVTETensor(pre_gelu_out);
-  Tensor *wspace = convertNVTETensor(workspace);
-  if ((biasTensor->data.dptr != nullptr) || (outputGelu->data.dptr != nullptr)) {
-    NVTE_ERROR("MOE batchgemm not surpport bias or gelu.");
-  }
-
-  int m, n, k;
-  if (!transa && transb) {
-  // for NT
-  m = transa ? inputA->data.shape[0]/batch_count : inputA->data.shape[1];
-  k = transa ? inputA->data.shape[1] : inputA->data.shape[0]/batch_count;
-  n = transb ? inputB->data.shape[1] : inputB->data.shape[0]/batch_count;
-  }  else if(transa && !transb){
-  // for TN
-  m = transa ? inputA->data.shape[0]/batch_count: inputA->data.shape[1];
-  k = transa ? inputA->data.shape[1] : inputA->data.shape[0]/batch_count;
-  n = transb ? inputB->data.shape[1] : inputB->data.shape[0]/batch_count;
-  } else if(!transa && !transb){
-  // for NN
-  m = transa ? inputA->data.shape[0]/batch_count : inputA->data.shape[1];
-  k = transa ? inputA->data.shape[1] : inputA->data.shape[0]/batch_count;
-  n = transb ? inputB->data.shape[1] : inputB->data.shape[0]/batch_count; }
-  int lda, ldb, ldd;
-  if (transa && !transb) {  // TN
-    lda = k;
-    ldb = k;
-    ldd = m;
-  } else if (!transa && !transb) {  // NN
-    lda = m; 
-    ldb = k;
-    ldd = m;
-  } else if (!transa && transb) {  // NT
-    lda = m;
-    ldb = n;
-    ldd = m;
-  } else {  // TT
-    NVTE_ERROR("TT layout not allowed.");
-  }
-  hipblas_batchgemm(inputA,
-            inputB,
-            outputD,
-            biasTensor,
-            outputGelu,
-            m, n, k,
-            lda, ldb, ldd,
-            (transa) ? HIPBLAS_OP_T : HIPBLAS_OP_N,
-            (transb) ? HIPBLAS_OP_T : HIPBLAS_OP_N,
-            grad, wspace->data.dptr,
-            wspace->data.shape[0],
-            accumulate, use_split_accumulator,
-            math_sm_count,
-            0,
-            0,
-            false,
-            nullptr,
-            batch_count,
-            stream);
-}
 
 // add for batchgemm
-void nvte_cublas_batchgemm_v3(const NVTETensor A, const NVTETensor B, const NVTETensor A_scales, const NVTETensor B_scales, NVTETensor D, const NVTETensor bias,
+void nvte_cublas_batchgemm_tensorwise_int8(const NVTETensor A, const NVTETensor B, const NVTETensor A_scales, const NVTETensor B_scales, NVTETensor D, const NVTETensor bias,
                       NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,
                       NVTETensor workspace, bool accumulate, bool use_split_accumulator,
                       int math_sm_count, int batch_count, cudaStream_t stream) {
-  NVTE_API_CALL(nvte_cublas_batchgemm_v3);
+  NVTE_API_CALL(nvte_cublas_batchgemm_tensorwise_int8);
   using namespace transformer_engine;
   const Tensor *inputA = convertNVTETensorCheck(A);
   const Tensor *inputB = convertNVTETensorCheck(B);
@@ -1297,16 +1228,7 @@ void nvte_cublas_batchgemm_v3(const NVTETensor A, const NVTETensor B, const NVTE
 
   handle = hipblaslt_handles[0];
 
-  hipblaslt_batchgemm_tensorwise_int8(inputA, inputB, inputA_scales, inputB_scales, outputD, biasTensor, outputGelu, 
-                 m, n, k, lda, ldb, ldd,
-                 (transa) ? HIPBLAS_OP_T : HIPBLAS_OP_N,
-                 (transb) ? HIPBLAS_OP_T : HIPBLAS_OP_N,
-                 grad,
-                 wspace->data.dptr,
-                 wspace->data.shape[0], accumulate, use_split_accumulator,
-                 math_sm_count, 0, 0, 
-                 false, nullptr, batch_count, stream,
-                 handle);
+  NVTE_ERROR("Remove nvte_cublas_batchgemm_tensorwise_int8 for now.");
 
 }
 #endif

transformer_engine/common/include/transformer_engine/gemm.h

@@ -152,12 +152,7 @@ void nvte_cublas_batchgemm(const NVTETensor A, const NVTETensor B, NVTETensor D,
                       NVTETensor workspace, bool accumulate, bool use_split_accumulator,
                       int math_sm_count, int batch_count, cudaStream_t stream);
 
-void nvte_cublas_batchgemm_v2(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias,
-                      NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,
-                      NVTETensor workspace, bool accumulate, bool use_split_accumulator,
-                      int math_sm_count, int batch_count, cudaStream_t stream);
-
-void nvte_cublas_batchgemm_v3(const NVTETensor A, const NVTETensor B, const NVTETensor A_scales, const NVTETensor B_scales, NVTETensor D, const NVTETensor bias,
+void nvte_cublas_batchgemm_tensorwise_int8(const NVTETensor A, const NVTETensor B, const NVTETensor A_scales, const NVTETensor B_scales, NVTETensor D, const NVTETensor bias,
                       NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,
                       NVTETensor workspace, bool accumulate, bool use_split_accumulator,
                       int math_sm_count, int batch_count, cudaStream_t stream);

transformer_engine/pytorch/csrc/extensions/gemm.cpp

@@ -588,7 +588,7 @@ std::vector<py::object> generic_batchgemm(py::handle A, bool transa, py::handle
     } else {
       // Launch GEMM
       NVTE_SCOPED_GIL_RELEASE({
-        nvte_cublas_batchgemm_v2(A_tensor.data(), B_tensor.data(), D_tensor.data(), bias_tensor.data(),
+        nvte_cublas_batchgemm(A_tensor.data(), B_tensor.data(), D_tensor.data(), bias_tensor.data(),
                          te_pre_gelu_out.data(), transa, transb, grad, te_workspace.data(),
                          accumulate, use_split_accumulator, num_math_sms, batch_count, main_stream);
       });
@@ -724,7 +724,7 @@ std::vector<py::object> tensorwise_int8_batchgemm(py::handle A, bool transa, py:
     } else {
       // Launch GEMM
       NVTE_SCOPED_GIL_RELEASE({
-        nvte_cublas_batchgemm_v3(A_tensor.data(), B_tensor.data(), A_scales_tensor.data(), B_scales_tensor.data(), D_tensor.data(), bias_tensor.data(),
+        nvte_cublas_batchgemm_tensorwise_int8(A_tensor.data(), B_tensor.data(), A_scales_tensor.data(), B_scales_tensor.data(), D_tensor.data(), bias_tensor.data(),
                          te_pre_gelu_out.data(), transa, transb, grad, te_workspace.data(),
                          accumulate, use_split_accumulator, num_math_sms, batch_count, main_stream);
       });