gelu.cu 2.28 KB
Newer Older
Przemek Tredak's avatar
Przemek Tredak committed
1
/*************************************************************************
2
 * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Przemek Tredak's avatar
Przemek Tredak committed
3
4
5
 *
 * See LICENSE for license information.
 ************************************************************************/
6

7
#include "../util/math.h"
8
#include "./activation_template.h"
Przemek Tredak's avatar
Przemek Tredak committed
9

10
void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
11
  NVTE_API_CALL(nvte_gelu);
Przemek Tredak's avatar
Przemek Tredak committed
12
  using namespace transformer_engine;
13
  act_fn<fp32, Empty, gelu<fp32, fp32>>(input, output, stream);
14
15
}

16
void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output,
17
18
19
                cudaStream_t stream) {
  NVTE_API_CALL(nvte_dgelu);
  using namespace transformer_engine;
20
  dact_fn<fp32, Empty, dgelu<fp32, fp32>>(grad, input, output, stream);
Przemek Tredak's avatar
Przemek Tredak committed
21
}
22

23
void nvte_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
24
  NVTE_API_CALL(nvte_geglu);
25
  using namespace transformer_engine;
26
  gated_act_fn<fp32, Empty, gelu<fp32, fp32>>(input, output, stream);
27
28
}

29
void nvte_dgeglu(const NVTETensor grad, const NVTETensor input, NVTETensor output,
30
                 cudaStream_t stream) {
31
  NVTE_API_CALL(nvte_dgeglu);
32
  using namespace transformer_engine;
33
  dgated_act_fn<fp32, Empty, gelu<fp32, fp32>, dgelu<fp32, fp32>>(grad, input, output, stream);
34
}
35

36
void nvte_qgelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
37
38
  NVTE_API_CALL(nvte_qgelu);
  using namespace transformer_engine;
39
  act_fn<fp32, Empty, qgelu<fp32, fp32>>(input, output, stream);
40
41
}

42
43
void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output,
                 cudaStream_t stream) {
44
45
  NVTE_API_CALL(nvte_dqgelu);
  using namespace transformer_engine;
46
  dact_fn<fp32, Empty, dqgelu<fp32, fp32>>(grad, input, output, stream);
47
}
48

49
void nvte_qgeglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
50
51
  NVTE_API_CALL(nvte_qgeglu);
  using namespace transformer_engine;
52
  gated_act_fn<fp32, Empty, qgelu<fp32, fp32>>(input, output, stream);
53
54
}

55
56
void nvte_dqgeglu(const NVTETensor grad, const NVTETensor input, NVTETensor output,
                  cudaStream_t stream) {
57
58
  NVTE_API_CALL(nvte_dqgeglu);
  using namespace transformer_engine;
59
  dgated_act_fn<fp32, Empty, qgelu<fp32, fp32>, dqgelu<fp32, fp32>>(grad, input, output, stream);
60
}