micro optimizations

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

Author: zhongbozhu

transformer_engine/pytorch/csrc/extensions.h

@@ -190,22 +190,24 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p
  **************************************************************************************************/
 
 std::vector<at::Tensor> megacpp_grouped_mlp_forward(
-    const at::Tensor &input, const at::Tensor &split_sizes, py::handle fc1_weight,
+    const at::Tensor &input, at::ScalarType act_dtype, const at::Tensor &split_sizes,
+    py::handle fc1_weight,
     py::handle fc1_bias, py::handle fc2_weight, py::handle fc2_bias,
     const std::optional<at::Tensor> &act_scales, const std::string &activation,
     int64_t glu_interleave_size, double activation_limit, double activation_alpha,
-    double activation_glu_linear_offset);
+    double activation_glu_linear_offset, py::handle gemm_scratch);
 
 py::tuple megacpp_grouped_mlp_backward(
-    const at::Tensor &grad_output, const at::Tensor &split_sizes, const at::Tensor &x_offsets,
+    const at::Tensor &grad_output, at::ScalarType act_dtype, const at::Tensor &split_sizes,
+    const at::Tensor &x_offsets,
     const at::Tensor &fc1_offsets, const at::Tensor &fc2_offsets, const at::Tensor &fc2_dy_offsets,
     const at::Tensor &base_offsets, const at::Tensor &x, const at::Tensor &fc1_activation_input,
     const at::Tensor &fc2_x, const std::optional<at::Tensor> &act_scales, py::handle fc1_weight,
     py::handle fc2_weight, py::handle fc1_wgrad_output, bool fc1_compute_wgrad,
     bool fc1_accumulate_wgrad, py::handle fc2_wgrad_output, bool fc2_compute_wgrad,
     bool fc2_accumulate_wgrad, const std::string &activation, int64_t glu_interleave_size,
     double activation_limit, double activation_alpha, double activation_glu_linear_offset,
-    bool act_scales_requires_grad, bool input_requires_grad);
+    bool act_scales_requires_grad, bool input_requires_grad, py::handle gemm_scratch);
 
 /***************************************************************************************************
  * Transpose

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -358,22 +358,25 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         &transformer_engine::pytorch::te_general_grouped_gemm_for_discrete_out,
         "Grouped GEMM for discrete output list");
   m.def("megacpp_grouped_mlp_forward", &transformer_engine::pytorch::megacpp_grouped_mlp_forward,
-        "Mega C++ grouped MLP forward", py::arg("input"), py::arg("split_sizes"),
-        py::arg("fc1_weight"), py::arg("fc1_bias"), py::arg("fc2_weight"), py::arg("fc2_bias"),
+        "Mega C++ grouped MLP forward", py::arg("input"), py::arg("act_dtype"),
+        py::arg("split_sizes"), py::arg("fc1_weight"), py::arg("fc1_bias"),
+        py::arg("fc2_weight"), py::arg("fc2_bias"),
         py::arg("act_scales"), py::arg("activation"), py::arg("glu_interleave_size"),
         py::arg("activation_limit") = 0.0, py::arg("activation_alpha") = 0.0,
-        py::arg("activation_glu_linear_offset") = 0.0);
+        py::arg("activation_glu_linear_offset") = 0.0,
+        py::arg("gemm_scratch") = py::none());
   m.def("megacpp_grouped_mlp_backward", &transformer_engine::pytorch::megacpp_grouped_mlp_backward,
-        "Mega C++ grouped MLP backward", py::arg("grad_output"), py::arg("split_sizes"),
-        py::arg("x_offsets"), py::arg("fc1_offsets"), py::arg("fc2_offsets"),
-        py::arg("fc2_dy_offsets"), py::arg("base_offsets"), py::arg("x"),
+        "Mega C++ grouped MLP backward", py::arg("grad_output"), py::arg("act_dtype"),
+        py::arg("split_sizes"), py::arg("x_offsets"), py::arg("fc1_offsets"),
+        py::arg("fc2_offsets"), py::arg("fc2_dy_offsets"), py::arg("base_offsets"), py::arg("x"),
         py::arg("fc1_activation_input"), py::arg("fc2_x"), py::arg("act_scales"),
         py::arg("fc1_weight"), py::arg("fc2_weight"), py::arg("fc1_wgrad_output"),
         py::arg("fc1_compute_wgrad"), py::arg("fc1_accumulate_wgrad"), py::arg("fc2_wgrad_output"),
         py::arg("fc2_compute_wgrad"), py::arg("fc2_accumulate_wgrad"), py::arg("activation"),
         py::arg("glu_interleave_size"), py::arg("activation_limit") = 0.0,
         py::arg("activation_alpha") = 0.0, py::arg("activation_glu_linear_offset") = 0.0,
-        py::arg("act_scales_requires_grad") = true, py::arg("input_requires_grad") = true);
+        py::arg("act_scales_requires_grad") = true, py::arg("input_requires_grad") = true,
+        py::arg("gemm_scratch") = py::none());
   m.def("fp8_transpose", &transformer_engine::pytorch::fp8_transpose, "Transpose with FP8 I/O",
         py::arg("input"), py::arg("dtype"), py::kw_only(), py::arg("out"),
         py::call_guard<py::gil_scoped_release>());