pos_encoding.cpp 367 Bytes
Newer Older
1
2
#include <torch/extension.h>

3
void rotary_embedding(
4
5
6
  torch::Tensor& positions,
  torch::Tensor& query,
  torch::Tensor& key,
7
  int head_size,
8
9
  torch::Tensor& cos_sin_cache,
  bool is_neox);
10
11
12

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
13
14
15
    "rotary_embedding",
    &rotary_embedding,
    "Apply GPT-NeoX or GPT-J style rotary embedding to query and key");
16
}