FlagGems/include/flag_gems/operators.h at ee5fc74e0ba9f4e63b7fe84f529703753d5e4cb0 · AdvancedCompiler/FlagGems · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
#pragma once
#include <optional>
#include "torch/torch.h"

namespace flag_gems {
at::Tensor zeros(at::IntArrayRef size,
                 c10::optional<at::ScalarType> dtype = ::std::nullopt,
                 c10::optional<at::Layout> layout = ::std::nullopt,
                 c10::optional<at::Device> device = ::std::nullopt,
                 c10::optional<bool> pin_memory = ::std::nullopt);
at::Tensor add_tensor(const at::Tensor &a_, const at::Tensor &b_);
at::Tensor mm_tensor(const at::Tensor &mat1, const at::Tensor &mat2);
at::Tensor sum_dim(const at::Tensor &self,
                   at::OptionalIntArrayRef dim,
                   bool keepdim = false,
                   ::std::optional<at::ScalarType> dtype = ::std::nullopt);

at::Tensor rms_norm(const at::Tensor &input, const at::Tensor &weight, double epsilon = 1e-5);
void fused_add_rms_norm(at::Tensor &input,
                        at::Tensor &residual,
                        const at::Tensor &weight,
                        double epsilon = 1e-5);
at::Tensor nonzero(const at::Tensor &inp);
// Rotary embedding
void rotary_embedding_inplace(at::Tensor &q,
                              at::Tensor &k,
                              const at::Tensor &cos,
                              const at::Tensor &sin,
                              const std::optional<at::Tensor> &position_ids = std::nullopt,
                              bool rotary_interleaved = false);
std::tuple<at::Tensor, at::Tensor> rotary_embedding(
    const at::Tensor &q,
    const at::Tensor &k,
    const at::Tensor &cos,
    const at::Tensor &sin,
    const std::optional<at::Tensor> &position_ids = std::nullopt,
    bool rotary_interleaved = false);

at::Tensor cat(const at::TensorList &tensors, int64_t dim = 0);
at::Tensor bmm(const at::Tensor &A, const at::Tensor &B);
}  // namespace flag_gems