Optimize data movement (#20)

2023-04-02 00:30:17 -07:00
parent 1f01a18d39
commit 897cb2ae28
17 changed files with 275 additions and 135 deletions
--- a/csrc/pos_encoding.cpp
+++ b/csrc/pos_encoding.cpp
@ -1,8 +1,6 @@
 #include <torch/extension.h>

 void rotary_embedding_neox(
-  torch::Tensor& out_query,
-  torch::Tensor& out_key,
  torch::Tensor& positions,
  torch::Tensor& query,
  torch::Tensor& key,