cleaned public facing C++ api for CommOverlapCore

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

Author: denera

transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp

@@ -386,50 +386,8 @@ CublasMpDims compute_rs_dims(const TensorWrapper &A, bool transa, const TensorWr
   return {m, n, k};
 }
 
-CublasMpDims compute_ar_dims(const TensorWrapper &A, bool transa, const TensorWrapper &B,
-                             bool transb, int tp_size) {
-  // AR shares the same m/n/k semantics as RS at descriptor level.
-  return compute_rs_dims(A, transa, B, transb, tp_size);
-}
-
 }  // namespace
 
-void CommOverlapCore::cublasmp_ag_gemm(const TensorWrapper &A, bool transa, const TensorWrapper &B,
-                                       bool transb, TensorWrapper &D, TensorWrapper &bias,
-                                       TensorWrapper &pre_gelu_out, bool grad, bool accumulate,
-                                       cudaStream_t stream_main) {
-  auto [m, n, k] = compute_ag_dims(A, transa, B, transb, _tp_size);
-  // col-major GEMM compute overlapped with all-gather on input B
-  // (M/P, K) x [(K, N/P) -(AG)-> (K, N)] = (M/P, N)
-  nvte_all_gather_gemm(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
-                       pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
-                       stream_main, _algo_type);
-}
-
-void CommOverlapCore::cublasmp_gemm_rs(const TensorWrapper &A, bool transa, const TensorWrapper &B,
-                                       bool transb, TensorWrapper &D, TensorWrapper &bias,
-                                       TensorWrapper &pre_gelu_out, bool grad, bool accumulate,
-                                       cudaStream_t stream_main) {
-  auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size);
-  // col-major GEMM compute overlapped with reduce-scatter on the output
-  // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P)
-  nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
-                           pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
-                           stream_main, _algo_type);
-}
-
-void CommOverlapCore::cublasmp_gemm_ar(const TensorWrapper &A, bool transa, const TensorWrapper &B,
-                                       bool transb, TensorWrapper &D, TensorWrapper &bias,
-                                       TensorWrapper &pre_gelu_out, bool grad, bool accumulate,
-                                       cudaStream_t stream_main) {
-  auto [m, n, k] = compute_ar_dims(A, transa, B, transb, _tp_size);
-  // col-major GEMM compute overlapped with all-reduce on the output
-  // (M, K/P) x (K/P, N) = (M, N) -(AR)-> (M, N)
-  nvte_gemm_all_reduce(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
-                       pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
-                       stream_main, _algo_type);
-}
-
 /***************************************************************************************************
  * Comm+GEMM Overlap Base (Pipelined / Collective)
  **************************************************************************************************/
@@ -549,8 +507,13 @@ void CommOverlapBase::atomic_gemm_overlap_rs(const TensorWrapper &A, bool transa
                                              bool use_split_accumulator, TensorWrapper &rs_output,
                                              cudaStream_t stream_main) {
   if (_with_cublasmp) {
-    return cublasmp_gemm_rs(A, transa, B, transb, D, bias, pre_gelu_out, grad, accumulate,
-                            stream_main);
+    auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size);
+    // col-major GEMM compute overlapped with reduce-scatter on the output
+    // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P)
+    nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
+                            pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
+                            stream_main, _algo_type);
+    return;
   }
 
   int ori_sms = _ub_comm->sms;
@@ -651,8 +614,13 @@ void CommOverlapBase::split_overlap_rs(const TensorWrapper &A, bool transa, cons
                                        bool grad, bool accumulate, bool use_split_accumulator,
                                        TensorWrapper &rs_output, cudaStream_t stream_main) {
   if (_with_cublasmp) {
-    return cublasmp_gemm_rs(A, transa, B, transb, D, bias, pre_gelu_out, grad, accumulate,
-                            stream_main);
+    auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size);
+    // col-major GEMM compute overlapped with reduce-scatter on the output
+    // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P)
+    nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
+                            pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
+                            stream_main, _algo_type);
+    return;
   }
 
   // Get GEMM dimensions
@@ -968,8 +936,13 @@ void CommOverlapP2PBase::atomic_gemm_overlap_ag(
     TensorWrapper &bias, TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad,
     bool accumulate, bool use_split_accumulator, TensorWrapper &B_copy, cudaStream_t stream_main) {
   if (_with_cublasmp) {
-    return cublasmp_ag_gemm(A, transa, B, transb, D, bias, pre_gelu_out, grad, accumulate,
-                            stream_main);
+    auto [m, n, k] = compute_ag_dims(A, transa, B, transb, _tp_size);
+    // col-major GEMM compute overlapped with all-gather on input B
+    // (M/P, K) x [(K, N/P) -(AG)-> (K, N)] = (M/P, N)
+    nvte_all_gather_gemm(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
+                        pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
+                        stream_main, _algo_type);
+    return;
   }
 
   int ori_sms = _ub_comm->sms;
@@ -1075,8 +1048,13 @@ void CommOverlapP2PBase::split_overlap_ag(const TensorWrapper &A, bool transa,
                                           bool use_split_accumulator, TensorWrapper &B_copy,
                                           cudaStream_t stream_main) {
   if (_with_cublasmp) {
-    return cublasmp_ag_gemm(A, transa, B, transb, D, bias, pre_gelu_out, grad, accumulate,
-                            stream_main);
+    auto [m, n, k] = compute_ag_dims(A, transa, B, transb, _tp_size);
+    // col-major GEMM compute overlapped with all-gather on input B
+    // (M/P, K) x [(K, N/P) -(AG)-> (K, N)] = (M/P, N)
+    nvte_all_gather_gemm(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
+                        pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
+                        stream_main, _algo_type);
+    return;
   }
 
   int ori_sms = _ub_comm->sms;
@@ -1247,8 +1225,13 @@ void CommOverlapP2PBase::atomic_gemm_overlap_rs(
     bool accumulate, bool use_split_accumulator, TensorWrapper &rs_output,
     cudaStream_t stream_main) {
   if (_with_cublasmp) {
-    return cublasmp_gemm_rs(A, transa, B, transb, D, bias, pre_gelu_out, grad, accumulate,
-                            stream_main);
+    auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size);
+    // col-major GEMM compute overlapped with reduce-scatter on the output
+    // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P)
+    nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
+                            pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
+                            stream_main, _algo_type);
+    return;
   }
 
   int ori_sms = _ub_comm->sms;
@@ -1316,8 +1299,13 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa,
                                           bool use_split_accumulator, TensorWrapper &rs_output,
                                           cudaStream_t stream_main) {
   if (_with_cublasmp) {
-    return cublasmp_gemm_rs(A, transa, B, transb, D, bias, pre_gelu_out, grad, accumulate,
-                            stream_main);
+    auto [m, n, k] = compute_rs_dims(A, transa, B, transb, _tp_size);
+    // col-major GEMM compute overlapped with reduce-scatter on the output
+    // (M, K/P) x (K/P, N) = (M, N) -(RS)-> (M, N/P)
+    nvte_gemm_reduce_scatter(_cublasmp_ctx, m, n, k, A.data(), B.data(), D.data(), bias.data(),
+                            pre_gelu_out.data(), transa, transb, grad, accumulate, _num_comm_sm,
+                            stream_main, _algo_type);
+    return;
   }
 
   int ori_sms = _ub_comm->sms;

transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h

@@ -128,18 +128,6 @@ class CommOverlapCore {
 
   bool with_cublasmp() { return _with_cublasmp; }
 
-  void cublasmp_ag_gemm(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb,
-                        TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out,
-                        bool grad, bool accumulate, cudaStream_t stream_main);
-
-  void cublasmp_gemm_rs(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb,
-                        TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out,
-                        bool grad, bool accumulate, cudaStream_t stream_main);
-
-  void cublasmp_gemm_ar(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb,
-                        TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out,
-                        bool grad, bool accumulate, cudaStream_t stream_main);
-
   virtual void bulk_overlap(const TensorWrapper &A, bool transa, const TensorWrapper &B,
                             bool transb, TensorWrapper &D, TensorWrapper &bias,
                             TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad,

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -170,16 +170,17 @@ Error_Type GemmInitV2FFI(Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type
                                std::vector<size_t>{static_cast<size_t>(bias.element_count())});
       }
       TensorWrapper pre_gelu_out_(get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING));
+      TensorWrapper dummy;
       // Match GemmV2FFI's operand swap: rhs becomes A, lhs becomes B.
       cudaStream_t prepare_stream = cudaStreamPerThread;
       if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) {
-        executor->cublasmp_ag_gemm(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, d_,
-                                   bias_, pre_gelu_out_, false /*grad*/, false /*accumulate*/,
-                                   prepare_stream);
+        executor->split_overlap_ag(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, d_,
+                                   bias_, pre_gelu_out_, dummy, false /*grad*/, false /*accumulate*/,
+                                   false /*use_split_accumulator*/, dummy, prepare_stream);
       } else if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) {
-        executor->cublasmp_gemm_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, d_,
-                                   bias_, pre_gelu_out_, false /*grad*/, false /*accumulate*/,
-                                   prepare_stream);
+        executor->split_overlap_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, d_,
+                                   bias_, pre_gelu_out_, dummy, false /*grad*/, false /*accumulate*/,
+                                   false /*use_split_accumulator*/, dummy, prepare_stream);
       }
       NVTE_CHECK_CUDA(cudaStreamSynchronize(prepare_stream));
     }

transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp

@@ -274,18 +274,32 @@ void cublasmp_capture_warmup(te::CommOverlapCore *core, int tp_size, te::CommOve
   NVTE_CHECK_CUDA(cudaMemset(a_ptr, 0, a_bytes));
   NVTE_CHECK_CUDA(cudaMemset(b_ptr, 0, b_bytes));
 
-  te::TensorWrapper A_tw, B_tw, D_tw, bias_tw, pre_gelu_tw;
+  te::TensorWrapper A_tw, B_tw, D_tw, bias_tw, pre_gelu_tw, dummy;
   A_tw.set_rowwise_data(a_ptr, te::DType::kBFloat16, a_shape);
   B_tw.set_rowwise_data(b_ptr, te::DType::kBFloat16, b_shape);
   D_tw.set_rowwise_data(d_ptr, te::DType::kBFloat16, d_shape);
 
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   if (comm_type == te::CommOverlapType::AG) {
-    core->cublasmp_ag_gemm(A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw,
-                           pre_gelu_tw, /*grad=*/false, /*accumulate=*/false, stream);
+    if (core->is_atomic_gemm()) {
+      core->atomic_gemm_overlap_ag(
+          A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy,
+          /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream);
+    } else {
+      core->split_overlap_ag(
+          A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy,
+          /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream);
+    }
   } else {
-    core->cublasmp_gemm_rs(A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw,
-                           pre_gelu_tw, /*grad=*/false, /*accumulate=*/false, stream);
+    if (core->is_atomic_gemm()) {
+      core->atomic_gemm_overlap_rs(
+          A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy,
+          /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream);
+    } else {
+      core->split_overlap_rs(
+          A_tw, /*transa=*/true, B_tw, /*transb=*/false, D_tw, bias_tw, pre_gelu_tw, dummy,
+          /*grad=*/false, /*accumulate=*/false, /*use_split_accumulator=*/false, dummy, stream);
+    }
   }
   NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
   cudaFree(a_ptr);