Add kernel for GeGLU with approximate GELU (#3337)

2024-03-12 22:06:17 -07:00
parent 49a3c8662b
commit 602358f8a8
5 changed files with 49 additions and 7 deletions
--- a/csrc/pybind.cpp
+++ b/csrc/pybind.cpp
@@ -25,7 +25,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  ops.def(
    "gelu_and_mul",
    &gelu_and_mul,
-    "Activation function used in GeGLU.");
+    "Activation function used in GeGLU with `none` approximation.");
+  ops.def(
+    "gelu_tanh_and_mul",
+    &gelu_tanh_and_mul,
+    "Activation function used in GeGLU with `tanh` approximation.");
  ops.def(
    "gelu_new",
    &gelu_new,