Add support for GPT-NeoX (Pythia) (#50)

2023-04-28 00:32:10 -07:00
parent aa50b17ca7
commit a96d63c21d
9 changed files with 436 additions and 71 deletions
--- a/csrc/pos_encoding.cpp
+++ b/csrc/pos_encoding.cpp
@ -4,6 +4,7 @@ void rotary_embedding_neox(
  torch::Tensor& positions,
  torch::Tensor& query,
  torch::Tensor& key,
+  int head_size,
  torch::Tensor& cos_sin_cache);

 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {