Optimize GeGLU layer in Gemma (#2975)

2024-02-21 20:17:52 -08:00
parent 93dc5a2870
commit fd5dcc5c81
6 changed files with 107 additions and 76 deletions
--- a/csrc/pybind.cpp
+++ b/csrc/pybind.cpp
@@ -22,6 +22,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    "silu_and_mul",
    &silu_and_mul,
    "Activation function used in SwiGLU.");
+  ops.def(
+    "gelu_and_mul",
+    &gelu_and_mul,
+    "Activation function used in GeGLU.");
  ops.def(
    "gelu_new",
    &gelu_new,