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/ops.h
+++ b/csrc/ops.h
@@ -57,6 +57,10 @@ void silu_and_mul(
  torch::Tensor& out,
  torch::Tensor& input);

+void gelu_and_mul(
+  torch::Tensor& out,
+  torch::Tensor& input);
+
 void gelu_new(
  torch::Tensor& out,
  torch::Tensor& input);