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

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