"test/vscode:/vscode.git/clone" did not exist on "bb3974807e3d9ab847de71a58a8f2985810a8166"
pos_encoding.cpp 340 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
#include <torch/extension.h>

void rotary_embedding_neox(
  torch::Tensor& positions,
  torch::Tensor& query,
  torch::Tensor& key,
  torch::Tensor& cos_sin_cache);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
    "rotary_embedding_neox",
    &rotary_embedding_neox,
    "Apply GPT-NeoX style rotary embedding to query and key");
}