Build uploaded using `kernels`.
Browse files- build/torch27-cxx11-cu128-aarch64-linux/activation/__init__.py +75 -0
- build/torch27-cxx11-cu128-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc +0 -0
- build/torch27-cxx11-cu128-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc +0 -0
- build/torch27-cxx11-cu128-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc +0 -0
- build/{torch28-cxx11-cu129-aarch64-linux/activation/_activation_0c3eb4e_dirty.abi3.so → torch27-cxx11-cu128-aarch64-linux/activation/_activation_320b408.abi3.so} +2 -2
- build/torch27-cxx11-cu128-aarch64-linux/activation/_ops.py +9 -0
- build/torch27-cxx11-cu128-aarch64-linux/activation/layers.py +179 -0
- build/torch28-cxx11-cu129-aarch64-linux/activation/__init__.py +18 -0
- build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc +0 -0
- build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc +0 -0
- build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc +0 -0
- build/torch28-cxx11-cu129-aarch64-linux/activation/_activation_320b408.abi3.so +3 -0
- build/torch28-cxx11-cu129-aarch64-linux/activation/_ops.py +3 -3
- build/torch28-cxx11-cu129-aarch64-linux/activation/layers.py +51 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/__init__.py +75 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/_activation_320b408.abi3.so +3 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/_ops.py +9 -0
- build/torch29-cxx11-cu126-aarch64-linux/activation/layers.py +179 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/__init__.py +75 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/_activation_320b408.abi3.so +3 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/_ops.py +9 -0
- build/torch29-cxx11-cu128-aarch64-linux/activation/layers.py +179 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/__init__.py +75 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc +0 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/_activation_320b408.abi3.so +3 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/_ops.py +9 -0
- build/torch29-cxx11-cu130-aarch64-linux/activation/layers.py +179 -0
build/torch27-cxx11-cu128-aarch64-linux/activation/__init__.py
ADDED
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
from ._ops import ops
|
4 |
+
|
5 |
+
from . import layers
|
6 |
+
|
7 |
+
|
8 |
+
def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
9 |
+
ops.silu_and_mul(out, x)
|
10 |
+
return out
|
11 |
+
|
12 |
+
|
13 |
+
def mul_and_silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
14 |
+
ops.mul_and_silu(out, x)
|
15 |
+
return out
|
16 |
+
|
17 |
+
|
18 |
+
def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
19 |
+
ops.gelu_and_mul(out, x)
|
20 |
+
return out
|
21 |
+
|
22 |
+
|
23 |
+
def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
24 |
+
ops.gelu_tanh_and_mul(out, x)
|
25 |
+
return out
|
26 |
+
|
27 |
+
|
28 |
+
def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0) -> None:
|
29 |
+
ops.fatrelu_and_mul(out, x, threshold)
|
30 |
+
return out
|
31 |
+
|
32 |
+
|
33 |
+
def gelu(out: torch.Tensor, x: torch.Tensor) -> None:
|
34 |
+
ops.gelu(out, x)
|
35 |
+
return out
|
36 |
+
|
37 |
+
def silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
38 |
+
ops.silu(out, x)
|
39 |
+
return out
|
40 |
+
|
41 |
+
|
42 |
+
def gelu_tanh(out: torch.Tensor, x: torch.Tensor) -> None:
|
43 |
+
ops.gelu_tanh(out, x)
|
44 |
+
return out
|
45 |
+
|
46 |
+
|
47 |
+
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
48 |
+
ops.gelu_fast(out, x)
|
49 |
+
return out
|
50 |
+
|
51 |
+
|
52 |
+
def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
|
53 |
+
ops.gelu_new(out, x)
|
54 |
+
return out
|
55 |
+
|
56 |
+
|
57 |
+
def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
|
58 |
+
ops.gelu_quick(out, x)
|
59 |
+
return out
|
60 |
+
|
61 |
+
|
62 |
+
__all__ = [
|
63 |
+
"silu_and_mul",
|
64 |
+
"mul_and_silu",
|
65 |
+
"gelu_and_mul",
|
66 |
+
"gelu_tanh_and_mul",
|
67 |
+
"fatrelu_and_mul",
|
68 |
+
"gelu_fast",
|
69 |
+
"gelu_new",
|
70 |
+
"gelu_quick",
|
71 |
+
"gelu_tanh",
|
72 |
+
"silu",
|
73 |
+
"gelu",
|
74 |
+
"layers",
|
75 |
+
]
|
build/torch27-cxx11-cu128-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc
ADDED
Binary file (3.25 kB). View file
|
|
build/torch27-cxx11-cu128-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc
ADDED
Binary file (527 Bytes). View file
|
|
build/torch27-cxx11-cu128-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc
ADDED
Binary file (8.92 kB). View file
|
|
build/{torch28-cxx11-cu129-aarch64-linux/activation/_activation_0c3eb4e_dirty.abi3.so → torch27-cxx11-cu128-aarch64-linux/activation/_activation_320b408.abi3.so}
RENAMED
@@ -1,3 +1,3 @@
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:
|
3 |
-
size
|
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:34bdeb9ab72686850aef0a16b225b1b956162edb2cf46cba65c5e5b92ae267ae
|
3 |
+
size 4207000
|
build/torch27-cxx11-cu128-aarch64-linux/activation/_ops.py
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from . import _activation_320b408
|
3 |
+
ops = torch.ops._activation_320b408
|
4 |
+
|
5 |
+
def add_op_namespace_prefix(op_name: str):
|
6 |
+
"""
|
7 |
+
Prefix op by namespace.
|
8 |
+
"""
|
9 |
+
return f"_activation_320b408::{op_name}"
|
build/torch27-cxx11-cu128-aarch64-linux/activation/layers.py
ADDED
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
from ._ops import ops
|
5 |
+
|
6 |
+
|
7 |
+
class SiluAndMul(nn.Module):
|
8 |
+
"""An activation function for SwiGLU.
|
9 |
+
|
10 |
+
The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
11 |
+
|
12 |
+
Shapes:
|
13 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
14 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
15 |
+
"""
|
16 |
+
|
17 |
+
can_torch_compile: bool = True
|
18 |
+
|
19 |
+
def forward(self, x: torch.Tensor):
|
20 |
+
d = x.shape[-1] // 2
|
21 |
+
output_shape = x.shape[:-1] + (d,)
|
22 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
23 |
+
ops.silu_and_mul(out, x)
|
24 |
+
return out
|
25 |
+
|
26 |
+
class Silu(nn.Module):
|
27 |
+
"""An activation function for SiLU.
|
28 |
+
|
29 |
+
The function computes x -> silu(x).
|
30 |
+
|
31 |
+
Shapes:
|
32 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
33 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
34 |
+
"""
|
35 |
+
|
36 |
+
can_torch_compile: bool = True
|
37 |
+
|
38 |
+
def forward(self, x: torch.Tensor):
|
39 |
+
out = torch.empty_like(x)
|
40 |
+
ops.silu(out, x)
|
41 |
+
return out
|
42 |
+
|
43 |
+
class Gelu(nn.Module):
|
44 |
+
"""An activation function for GELU.
|
45 |
+
|
46 |
+
The function computes x -> gelu(x).
|
47 |
+
|
48 |
+
Shapes:
|
49 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
50 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
51 |
+
"""
|
52 |
+
|
53 |
+
can_torch_compile: bool = True
|
54 |
+
|
55 |
+
def forward(self, x: torch.Tensor):
|
56 |
+
out = torch.empty_like(x)
|
57 |
+
ops.gelu(out, x)
|
58 |
+
return out
|
59 |
+
|
60 |
+
class GeluTanh(nn.Module):
|
61 |
+
"""An activation function for GELU with `tanh` approximation.
|
62 |
+
|
63 |
+
The function computes x -> gelu_tanh(x).
|
64 |
+
|
65 |
+
Shapes:
|
66 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
67 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
68 |
+
"""
|
69 |
+
|
70 |
+
can_torch_compile: bool = True
|
71 |
+
|
72 |
+
def forward(self, x: torch.Tensor):
|
73 |
+
out = torch.empty_like(x)
|
74 |
+
ops.gelu_tanh(out, x)
|
75 |
+
return out
|
76 |
+
|
77 |
+
|
78 |
+
class MulAndSilu(nn.Module):
|
79 |
+
"""An activation function for SwiGLU.
|
80 |
+
|
81 |
+
The function computes x -> x[:d] * silu(x[d:]) where d = x.shape[-1] // 2.
|
82 |
+
|
83 |
+
Shapes:
|
84 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
85 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
86 |
+
"""
|
87 |
+
|
88 |
+
can_torch_compile: bool = True
|
89 |
+
|
90 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
91 |
+
d = x.shape[-1] // 2
|
92 |
+
output_shape = x.shape[:-1] + (d,)
|
93 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
94 |
+
ops.mul_and_silu(out, x)
|
95 |
+
return out
|
96 |
+
|
97 |
+
|
98 |
+
class GeluAndMul(nn.Module):
|
99 |
+
"""An activation function for GeGLU.
|
100 |
+
|
101 |
+
The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
102 |
+
|
103 |
+
Shapes:
|
104 |
+
x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
|
105 |
+
return: (batch_size, seq_len, d) or (num_tokens, d)
|
106 |
+
"""
|
107 |
+
|
108 |
+
can_torch_compile: bool = True
|
109 |
+
|
110 |
+
def forward(self, x: torch.Tensor):
|
111 |
+
d = x.shape[-1] // 2
|
112 |
+
output_shape = x.shape[:-1] + (d,)
|
113 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
114 |
+
ops.gelu_and_mul(out, x)
|
115 |
+
return out
|
116 |
+
|
117 |
+
|
118 |
+
class GeluTanhAndMul(nn.Module):
|
119 |
+
can_torch_compile: bool = True
|
120 |
+
|
121 |
+
def forward(self, x: torch.Tensor):
|
122 |
+
d = x.shape[-1] // 2
|
123 |
+
output_shape = x.shape[:-1] + (d,)
|
124 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
125 |
+
ops.gelu_tanh_and_mul(out, x)
|
126 |
+
return out
|
127 |
+
|
128 |
+
|
129 |
+
class FatreluAndMul(nn.Module):
|
130 |
+
"""An activation function for FATReLU.
|
131 |
+
|
132 |
+
The function computes x -> FATReLU(x[:d]) * x[d:] where
|
133 |
+
d = x.shape[-1] // 2.
|
134 |
+
This is used in openbmb/MiniCPM-S-1B-sft.
|
135 |
+
|
136 |
+
Shapes:
|
137 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
138 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
139 |
+
"""
|
140 |
+
|
141 |
+
can_torch_compile: bool = True
|
142 |
+
|
143 |
+
def __init__(self, threshold: float = 0.0):
|
144 |
+
super().__init__()
|
145 |
+
self.threshold = threshold
|
146 |
+
|
147 |
+
def forward(self, x: torch.Tensor):
|
148 |
+
d = x.shape[-1] // 2
|
149 |
+
output_shape = x.shape[:-1] + (d,)
|
150 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
151 |
+
ops.fatrelu_and_mul(out, x, self.threshold)
|
152 |
+
return out
|
153 |
+
|
154 |
+
|
155 |
+
class FastGELU(nn.Module):
|
156 |
+
can_torch_compile: bool = True
|
157 |
+
|
158 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
159 |
+
out = torch.empty_like(x)
|
160 |
+
ops.gelu_fast(out, x)
|
161 |
+
return out
|
162 |
+
|
163 |
+
|
164 |
+
class NewGELU(nn.Module):
|
165 |
+
can_torch_compile: bool = True
|
166 |
+
|
167 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
168 |
+
out = torch.empty_like(x)
|
169 |
+
ops.gelu_new(out, x)
|
170 |
+
return out
|
171 |
+
|
172 |
+
|
173 |
+
class QuickGELU(nn.Module):
|
174 |
+
can_torch_compile: bool = True
|
175 |
+
|
176 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
177 |
+
out = torch.empty_like(x)
|
178 |
+
ops.gelu_quick(out, x)
|
179 |
+
return out
|
build/torch28-cxx11-cu129-aarch64-linux/activation/__init__.py
CHANGED
@@ -30,6 +30,20 @@ def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0)
|
|
30 |
return out
|
31 |
|
32 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
33 |
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
34 |
ops.gelu_fast(out, x)
|
35 |
return out
|
@@ -47,11 +61,15 @@ def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
47 |
|
48 |
__all__ = [
|
49 |
"silu_and_mul",
|
|
|
50 |
"gelu_and_mul",
|
51 |
"gelu_tanh_and_mul",
|
52 |
"fatrelu_and_mul",
|
53 |
"gelu_fast",
|
54 |
"gelu_new",
|
55 |
"gelu_quick",
|
|
|
|
|
|
|
56 |
"layers",
|
57 |
]
|
|
|
30 |
return out
|
31 |
|
32 |
|
33 |
+
def gelu(out: torch.Tensor, x: torch.Tensor) -> None:
|
34 |
+
ops.gelu(out, x)
|
35 |
+
return out
|
36 |
+
|
37 |
+
def silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
38 |
+
ops.silu(out, x)
|
39 |
+
return out
|
40 |
+
|
41 |
+
|
42 |
+
def gelu_tanh(out: torch.Tensor, x: torch.Tensor) -> None:
|
43 |
+
ops.gelu_tanh(out, x)
|
44 |
+
return out
|
45 |
+
|
46 |
+
|
47 |
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
48 |
ops.gelu_fast(out, x)
|
49 |
return out
|
|
|
61 |
|
62 |
__all__ = [
|
63 |
"silu_and_mul",
|
64 |
+
"mul_and_silu",
|
65 |
"gelu_and_mul",
|
66 |
"gelu_tanh_and_mul",
|
67 |
"fatrelu_and_mul",
|
68 |
"gelu_fast",
|
69 |
"gelu_new",
|
70 |
"gelu_quick",
|
71 |
+
"gelu_tanh",
|
72 |
+
"silu",
|
73 |
+
"gelu",
|
74 |
"layers",
|
75 |
]
|
build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc
CHANGED
Binary files a/build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc and b/build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc differ
|
|
build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc
CHANGED
Binary files a/build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc and b/build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc differ
|
|
build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc
CHANGED
Binary files a/build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc and b/build/torch28-cxx11-cu129-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc differ
|
|
build/torch28-cxx11-cu129-aarch64-linux/activation/_activation_320b408.abi3.so
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:3af83bae80c8641200010ba586e5a2cac271fa4fcd344e3532ea7d5094fd7c17
|
3 |
+
size 4275744
|
build/torch28-cxx11-cu129-aarch64-linux/activation/_ops.py
CHANGED
@@ -1,9 +1,9 @@
|
|
1 |
import torch
|
2 |
-
from . import
|
3 |
-
ops = torch.ops.
|
4 |
|
5 |
def add_op_namespace_prefix(op_name: str):
|
6 |
"""
|
7 |
Prefix op by namespace.
|
8 |
"""
|
9 |
-
return f"
|
|
|
1 |
import torch
|
2 |
+
from . import _activation_320b408
|
3 |
+
ops = torch.ops._activation_320b408
|
4 |
|
5 |
def add_op_namespace_prefix(op_name: str):
|
6 |
"""
|
7 |
Prefix op by namespace.
|
8 |
"""
|
9 |
+
return f"_activation_320b408::{op_name}"
|
build/torch28-cxx11-cu129-aarch64-linux/activation/layers.py
CHANGED
@@ -23,6 +23,57 @@ class SiluAndMul(nn.Module):
|
|
23 |
ops.silu_and_mul(out, x)
|
24 |
return out
|
25 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
26 |
|
27 |
class MulAndSilu(nn.Module):
|
28 |
"""An activation function for SwiGLU.
|
|
|
23 |
ops.silu_and_mul(out, x)
|
24 |
return out
|
25 |
|
26 |
+
class Silu(nn.Module):
|
27 |
+
"""An activation function for SiLU.
|
28 |
+
|
29 |
+
The function computes x -> silu(x).
|
30 |
+
|
31 |
+
Shapes:
|
32 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
33 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
34 |
+
"""
|
35 |
+
|
36 |
+
can_torch_compile: bool = True
|
37 |
+
|
38 |
+
def forward(self, x: torch.Tensor):
|
39 |
+
out = torch.empty_like(x)
|
40 |
+
ops.silu(out, x)
|
41 |
+
return out
|
42 |
+
|
43 |
+
class Gelu(nn.Module):
|
44 |
+
"""An activation function for GELU.
|
45 |
+
|
46 |
+
The function computes x -> gelu(x).
|
47 |
+
|
48 |
+
Shapes:
|
49 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
50 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
51 |
+
"""
|
52 |
+
|
53 |
+
can_torch_compile: bool = True
|
54 |
+
|
55 |
+
def forward(self, x: torch.Tensor):
|
56 |
+
out = torch.empty_like(x)
|
57 |
+
ops.gelu(out, x)
|
58 |
+
return out
|
59 |
+
|
60 |
+
class GeluTanh(nn.Module):
|
61 |
+
"""An activation function for GELU with `tanh` approximation.
|
62 |
+
|
63 |
+
The function computes x -> gelu_tanh(x).
|
64 |
+
|
65 |
+
Shapes:
|
66 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
67 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
68 |
+
"""
|
69 |
+
|
70 |
+
can_torch_compile: bool = True
|
71 |
+
|
72 |
+
def forward(self, x: torch.Tensor):
|
73 |
+
out = torch.empty_like(x)
|
74 |
+
ops.gelu_tanh(out, x)
|
75 |
+
return out
|
76 |
+
|
77 |
|
78 |
class MulAndSilu(nn.Module):
|
79 |
"""An activation function for SwiGLU.
|
build/torch29-cxx11-cu126-aarch64-linux/activation/__init__.py
ADDED
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
from ._ops import ops
|
4 |
+
|
5 |
+
from . import layers
|
6 |
+
|
7 |
+
|
8 |
+
def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
9 |
+
ops.silu_and_mul(out, x)
|
10 |
+
return out
|
11 |
+
|
12 |
+
|
13 |
+
def mul_and_silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
14 |
+
ops.mul_and_silu(out, x)
|
15 |
+
return out
|
16 |
+
|
17 |
+
|
18 |
+
def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
19 |
+
ops.gelu_and_mul(out, x)
|
20 |
+
return out
|
21 |
+
|
22 |
+
|
23 |
+
def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
24 |
+
ops.gelu_tanh_and_mul(out, x)
|
25 |
+
return out
|
26 |
+
|
27 |
+
|
28 |
+
def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0) -> None:
|
29 |
+
ops.fatrelu_and_mul(out, x, threshold)
|
30 |
+
return out
|
31 |
+
|
32 |
+
|
33 |
+
def gelu(out: torch.Tensor, x: torch.Tensor) -> None:
|
34 |
+
ops.gelu(out, x)
|
35 |
+
return out
|
36 |
+
|
37 |
+
def silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
38 |
+
ops.silu(out, x)
|
39 |
+
return out
|
40 |
+
|
41 |
+
|
42 |
+
def gelu_tanh(out: torch.Tensor, x: torch.Tensor) -> None:
|
43 |
+
ops.gelu_tanh(out, x)
|
44 |
+
return out
|
45 |
+
|
46 |
+
|
47 |
+
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
48 |
+
ops.gelu_fast(out, x)
|
49 |
+
return out
|
50 |
+
|
51 |
+
|
52 |
+
def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
|
53 |
+
ops.gelu_new(out, x)
|
54 |
+
return out
|
55 |
+
|
56 |
+
|
57 |
+
def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
|
58 |
+
ops.gelu_quick(out, x)
|
59 |
+
return out
|
60 |
+
|
61 |
+
|
62 |
+
__all__ = [
|
63 |
+
"silu_and_mul",
|
64 |
+
"mul_and_silu",
|
65 |
+
"gelu_and_mul",
|
66 |
+
"gelu_tanh_and_mul",
|
67 |
+
"fatrelu_and_mul",
|
68 |
+
"gelu_fast",
|
69 |
+
"gelu_new",
|
70 |
+
"gelu_quick",
|
71 |
+
"gelu_tanh",
|
72 |
+
"silu",
|
73 |
+
"gelu",
|
74 |
+
"layers",
|
75 |
+
]
|
build/torch29-cxx11-cu126-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc
ADDED
Binary file (3.25 kB). View file
|
|
build/torch29-cxx11-cu126-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc
ADDED
Binary file (527 Bytes). View file
|
|
build/torch29-cxx11-cu126-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc
ADDED
Binary file (8.92 kB). View file
|
|
build/torch29-cxx11-cu126-aarch64-linux/activation/_activation_320b408.abi3.so
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:f9c24e0eb75a09a9fc19e7096276d560226f198617291681c1a18e94002a629e
|
3 |
+
size 2963480
|
build/torch29-cxx11-cu126-aarch64-linux/activation/_ops.py
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from . import _activation_320b408
|
3 |
+
ops = torch.ops._activation_320b408
|
4 |
+
|
5 |
+
def add_op_namespace_prefix(op_name: str):
|
6 |
+
"""
|
7 |
+
Prefix op by namespace.
|
8 |
+
"""
|
9 |
+
return f"_activation_320b408::{op_name}"
|
build/torch29-cxx11-cu126-aarch64-linux/activation/layers.py
ADDED
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
from ._ops import ops
|
5 |
+
|
6 |
+
|
7 |
+
class SiluAndMul(nn.Module):
|
8 |
+
"""An activation function for SwiGLU.
|
9 |
+
|
10 |
+
The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
11 |
+
|
12 |
+
Shapes:
|
13 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
14 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
15 |
+
"""
|
16 |
+
|
17 |
+
can_torch_compile: bool = True
|
18 |
+
|
19 |
+
def forward(self, x: torch.Tensor):
|
20 |
+
d = x.shape[-1] // 2
|
21 |
+
output_shape = x.shape[:-1] + (d,)
|
22 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
23 |
+
ops.silu_and_mul(out, x)
|
24 |
+
return out
|
25 |
+
|
26 |
+
class Silu(nn.Module):
|
27 |
+
"""An activation function for SiLU.
|
28 |
+
|
29 |
+
The function computes x -> silu(x).
|
30 |
+
|
31 |
+
Shapes:
|
32 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
33 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
34 |
+
"""
|
35 |
+
|
36 |
+
can_torch_compile: bool = True
|
37 |
+
|
38 |
+
def forward(self, x: torch.Tensor):
|
39 |
+
out = torch.empty_like(x)
|
40 |
+
ops.silu(out, x)
|
41 |
+
return out
|
42 |
+
|
43 |
+
class Gelu(nn.Module):
|
44 |
+
"""An activation function for GELU.
|
45 |
+
|
46 |
+
The function computes x -> gelu(x).
|
47 |
+
|
48 |
+
Shapes:
|
49 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
50 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
51 |
+
"""
|
52 |
+
|
53 |
+
can_torch_compile: bool = True
|
54 |
+
|
55 |
+
def forward(self, x: torch.Tensor):
|
56 |
+
out = torch.empty_like(x)
|
57 |
+
ops.gelu(out, x)
|
58 |
+
return out
|
59 |
+
|
60 |
+
class GeluTanh(nn.Module):
|
61 |
+
"""An activation function for GELU with `tanh` approximation.
|
62 |
+
|
63 |
+
The function computes x -> gelu_tanh(x).
|
64 |
+
|
65 |
+
Shapes:
|
66 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
67 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
68 |
+
"""
|
69 |
+
|
70 |
+
can_torch_compile: bool = True
|
71 |
+
|
72 |
+
def forward(self, x: torch.Tensor):
|
73 |
+
out = torch.empty_like(x)
|
74 |
+
ops.gelu_tanh(out, x)
|
75 |
+
return out
|
76 |
+
|
77 |
+
|
78 |
+
class MulAndSilu(nn.Module):
|
79 |
+
"""An activation function for SwiGLU.
|
80 |
+
|
81 |
+
The function computes x -> x[:d] * silu(x[d:]) where d = x.shape[-1] // 2.
|
82 |
+
|
83 |
+
Shapes:
|
84 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
85 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
86 |
+
"""
|
87 |
+
|
88 |
+
can_torch_compile: bool = True
|
89 |
+
|
90 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
91 |
+
d = x.shape[-1] // 2
|
92 |
+
output_shape = x.shape[:-1] + (d,)
|
93 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
94 |
+
ops.mul_and_silu(out, x)
|
95 |
+
return out
|
96 |
+
|
97 |
+
|
98 |
+
class GeluAndMul(nn.Module):
|
99 |
+
"""An activation function for GeGLU.
|
100 |
+
|
101 |
+
The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
102 |
+
|
103 |
+
Shapes:
|
104 |
+
x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
|
105 |
+
return: (batch_size, seq_len, d) or (num_tokens, d)
|
106 |
+
"""
|
107 |
+
|
108 |
+
can_torch_compile: bool = True
|
109 |
+
|
110 |
+
def forward(self, x: torch.Tensor):
|
111 |
+
d = x.shape[-1] // 2
|
112 |
+
output_shape = x.shape[:-1] + (d,)
|
113 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
114 |
+
ops.gelu_and_mul(out, x)
|
115 |
+
return out
|
116 |
+
|
117 |
+
|
118 |
+
class GeluTanhAndMul(nn.Module):
|
119 |
+
can_torch_compile: bool = True
|
120 |
+
|
121 |
+
def forward(self, x: torch.Tensor):
|
122 |
+
d = x.shape[-1] // 2
|
123 |
+
output_shape = x.shape[:-1] + (d,)
|
124 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
125 |
+
ops.gelu_tanh_and_mul(out, x)
|
126 |
+
return out
|
127 |
+
|
128 |
+
|
129 |
+
class FatreluAndMul(nn.Module):
|
130 |
+
"""An activation function for FATReLU.
|
131 |
+
|
132 |
+
The function computes x -> FATReLU(x[:d]) * x[d:] where
|
133 |
+
d = x.shape[-1] // 2.
|
134 |
+
This is used in openbmb/MiniCPM-S-1B-sft.
|
135 |
+
|
136 |
+
Shapes:
|
137 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
138 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
139 |
+
"""
|
140 |
+
|
141 |
+
can_torch_compile: bool = True
|
142 |
+
|
143 |
+
def __init__(self, threshold: float = 0.0):
|
144 |
+
super().__init__()
|
145 |
+
self.threshold = threshold
|
146 |
+
|
147 |
+
def forward(self, x: torch.Tensor):
|
148 |
+
d = x.shape[-1] // 2
|
149 |
+
output_shape = x.shape[:-1] + (d,)
|
150 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
151 |
+
ops.fatrelu_and_mul(out, x, self.threshold)
|
152 |
+
return out
|
153 |
+
|
154 |
+
|
155 |
+
class FastGELU(nn.Module):
|
156 |
+
can_torch_compile: bool = True
|
157 |
+
|
158 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
159 |
+
out = torch.empty_like(x)
|
160 |
+
ops.gelu_fast(out, x)
|
161 |
+
return out
|
162 |
+
|
163 |
+
|
164 |
+
class NewGELU(nn.Module):
|
165 |
+
can_torch_compile: bool = True
|
166 |
+
|
167 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
168 |
+
out = torch.empty_like(x)
|
169 |
+
ops.gelu_new(out, x)
|
170 |
+
return out
|
171 |
+
|
172 |
+
|
173 |
+
class QuickGELU(nn.Module):
|
174 |
+
can_torch_compile: bool = True
|
175 |
+
|
176 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
177 |
+
out = torch.empty_like(x)
|
178 |
+
ops.gelu_quick(out, x)
|
179 |
+
return out
|
build/torch29-cxx11-cu128-aarch64-linux/activation/__init__.py
ADDED
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
from ._ops import ops
|
4 |
+
|
5 |
+
from . import layers
|
6 |
+
|
7 |
+
|
8 |
+
def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
9 |
+
ops.silu_and_mul(out, x)
|
10 |
+
return out
|
11 |
+
|
12 |
+
|
13 |
+
def mul_and_silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
14 |
+
ops.mul_and_silu(out, x)
|
15 |
+
return out
|
16 |
+
|
17 |
+
|
18 |
+
def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
19 |
+
ops.gelu_and_mul(out, x)
|
20 |
+
return out
|
21 |
+
|
22 |
+
|
23 |
+
def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
24 |
+
ops.gelu_tanh_and_mul(out, x)
|
25 |
+
return out
|
26 |
+
|
27 |
+
|
28 |
+
def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0) -> None:
|
29 |
+
ops.fatrelu_and_mul(out, x, threshold)
|
30 |
+
return out
|
31 |
+
|
32 |
+
|
33 |
+
def gelu(out: torch.Tensor, x: torch.Tensor) -> None:
|
34 |
+
ops.gelu(out, x)
|
35 |
+
return out
|
36 |
+
|
37 |
+
def silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
38 |
+
ops.silu(out, x)
|
39 |
+
return out
|
40 |
+
|
41 |
+
|
42 |
+
def gelu_tanh(out: torch.Tensor, x: torch.Tensor) -> None:
|
43 |
+
ops.gelu_tanh(out, x)
|
44 |
+
return out
|
45 |
+
|
46 |
+
|
47 |
+
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
48 |
+
ops.gelu_fast(out, x)
|
49 |
+
return out
|
50 |
+
|
51 |
+
|
52 |
+
def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
|
53 |
+
ops.gelu_new(out, x)
|
54 |
+
return out
|
55 |
+
|
56 |
+
|
57 |
+
def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
|
58 |
+
ops.gelu_quick(out, x)
|
59 |
+
return out
|
60 |
+
|
61 |
+
|
62 |
+
__all__ = [
|
63 |
+
"silu_and_mul",
|
64 |
+
"mul_and_silu",
|
65 |
+
"gelu_and_mul",
|
66 |
+
"gelu_tanh_and_mul",
|
67 |
+
"fatrelu_and_mul",
|
68 |
+
"gelu_fast",
|
69 |
+
"gelu_new",
|
70 |
+
"gelu_quick",
|
71 |
+
"gelu_tanh",
|
72 |
+
"silu",
|
73 |
+
"gelu",
|
74 |
+
"layers",
|
75 |
+
]
|
build/torch29-cxx11-cu128-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc
ADDED
Binary file (3.25 kB). View file
|
|
build/torch29-cxx11-cu128-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc
ADDED
Binary file (527 Bytes). View file
|
|
build/torch29-cxx11-cu128-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc
ADDED
Binary file (8.92 kB). View file
|
|
build/torch29-cxx11-cu128-aarch64-linux/activation/_activation_320b408.abi3.so
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:08ee3dfa4d481eaf44ac3c11a0843598c05950f779dba66abd468fecb7839b32
|
3 |
+
size 4208760
|
build/torch29-cxx11-cu128-aarch64-linux/activation/_ops.py
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from . import _activation_320b408
|
3 |
+
ops = torch.ops._activation_320b408
|
4 |
+
|
5 |
+
def add_op_namespace_prefix(op_name: str):
|
6 |
+
"""
|
7 |
+
Prefix op by namespace.
|
8 |
+
"""
|
9 |
+
return f"_activation_320b408::{op_name}"
|
build/torch29-cxx11-cu128-aarch64-linux/activation/layers.py
ADDED
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
from ._ops import ops
|
5 |
+
|
6 |
+
|
7 |
+
class SiluAndMul(nn.Module):
|
8 |
+
"""An activation function for SwiGLU.
|
9 |
+
|
10 |
+
The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
11 |
+
|
12 |
+
Shapes:
|
13 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
14 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
15 |
+
"""
|
16 |
+
|
17 |
+
can_torch_compile: bool = True
|
18 |
+
|
19 |
+
def forward(self, x: torch.Tensor):
|
20 |
+
d = x.shape[-1] // 2
|
21 |
+
output_shape = x.shape[:-1] + (d,)
|
22 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
23 |
+
ops.silu_and_mul(out, x)
|
24 |
+
return out
|
25 |
+
|
26 |
+
class Silu(nn.Module):
|
27 |
+
"""An activation function for SiLU.
|
28 |
+
|
29 |
+
The function computes x -> silu(x).
|
30 |
+
|
31 |
+
Shapes:
|
32 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
33 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
34 |
+
"""
|
35 |
+
|
36 |
+
can_torch_compile: bool = True
|
37 |
+
|
38 |
+
def forward(self, x: torch.Tensor):
|
39 |
+
out = torch.empty_like(x)
|
40 |
+
ops.silu(out, x)
|
41 |
+
return out
|
42 |
+
|
43 |
+
class Gelu(nn.Module):
|
44 |
+
"""An activation function for GELU.
|
45 |
+
|
46 |
+
The function computes x -> gelu(x).
|
47 |
+
|
48 |
+
Shapes:
|
49 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
50 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
51 |
+
"""
|
52 |
+
|
53 |
+
can_torch_compile: bool = True
|
54 |
+
|
55 |
+
def forward(self, x: torch.Tensor):
|
56 |
+
out = torch.empty_like(x)
|
57 |
+
ops.gelu(out, x)
|
58 |
+
return out
|
59 |
+
|
60 |
+
class GeluTanh(nn.Module):
|
61 |
+
"""An activation function for GELU with `tanh` approximation.
|
62 |
+
|
63 |
+
The function computes x -> gelu_tanh(x).
|
64 |
+
|
65 |
+
Shapes:
|
66 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
67 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
68 |
+
"""
|
69 |
+
|
70 |
+
can_torch_compile: bool = True
|
71 |
+
|
72 |
+
def forward(self, x: torch.Tensor):
|
73 |
+
out = torch.empty_like(x)
|
74 |
+
ops.gelu_tanh(out, x)
|
75 |
+
return out
|
76 |
+
|
77 |
+
|
78 |
+
class MulAndSilu(nn.Module):
|
79 |
+
"""An activation function for SwiGLU.
|
80 |
+
|
81 |
+
The function computes x -> x[:d] * silu(x[d:]) where d = x.shape[-1] // 2.
|
82 |
+
|
83 |
+
Shapes:
|
84 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
85 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
86 |
+
"""
|
87 |
+
|
88 |
+
can_torch_compile: bool = True
|
89 |
+
|
90 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
91 |
+
d = x.shape[-1] // 2
|
92 |
+
output_shape = x.shape[:-1] + (d,)
|
93 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
94 |
+
ops.mul_and_silu(out, x)
|
95 |
+
return out
|
96 |
+
|
97 |
+
|
98 |
+
class GeluAndMul(nn.Module):
|
99 |
+
"""An activation function for GeGLU.
|
100 |
+
|
101 |
+
The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
102 |
+
|
103 |
+
Shapes:
|
104 |
+
x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
|
105 |
+
return: (batch_size, seq_len, d) or (num_tokens, d)
|
106 |
+
"""
|
107 |
+
|
108 |
+
can_torch_compile: bool = True
|
109 |
+
|
110 |
+
def forward(self, x: torch.Tensor):
|
111 |
+
d = x.shape[-1] // 2
|
112 |
+
output_shape = x.shape[:-1] + (d,)
|
113 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
114 |
+
ops.gelu_and_mul(out, x)
|
115 |
+
return out
|
116 |
+
|
117 |
+
|
118 |
+
class GeluTanhAndMul(nn.Module):
|
119 |
+
can_torch_compile: bool = True
|
120 |
+
|
121 |
+
def forward(self, x: torch.Tensor):
|
122 |
+
d = x.shape[-1] // 2
|
123 |
+
output_shape = x.shape[:-1] + (d,)
|
124 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
125 |
+
ops.gelu_tanh_and_mul(out, x)
|
126 |
+
return out
|
127 |
+
|
128 |
+
|
129 |
+
class FatreluAndMul(nn.Module):
|
130 |
+
"""An activation function for FATReLU.
|
131 |
+
|
132 |
+
The function computes x -> FATReLU(x[:d]) * x[d:] where
|
133 |
+
d = x.shape[-1] // 2.
|
134 |
+
This is used in openbmb/MiniCPM-S-1B-sft.
|
135 |
+
|
136 |
+
Shapes:
|
137 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
138 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
139 |
+
"""
|
140 |
+
|
141 |
+
can_torch_compile: bool = True
|
142 |
+
|
143 |
+
def __init__(self, threshold: float = 0.0):
|
144 |
+
super().__init__()
|
145 |
+
self.threshold = threshold
|
146 |
+
|
147 |
+
def forward(self, x: torch.Tensor):
|
148 |
+
d = x.shape[-1] // 2
|
149 |
+
output_shape = x.shape[:-1] + (d,)
|
150 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
151 |
+
ops.fatrelu_and_mul(out, x, self.threshold)
|
152 |
+
return out
|
153 |
+
|
154 |
+
|
155 |
+
class FastGELU(nn.Module):
|
156 |
+
can_torch_compile: bool = True
|
157 |
+
|
158 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
159 |
+
out = torch.empty_like(x)
|
160 |
+
ops.gelu_fast(out, x)
|
161 |
+
return out
|
162 |
+
|
163 |
+
|
164 |
+
class NewGELU(nn.Module):
|
165 |
+
can_torch_compile: bool = True
|
166 |
+
|
167 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
168 |
+
out = torch.empty_like(x)
|
169 |
+
ops.gelu_new(out, x)
|
170 |
+
return out
|
171 |
+
|
172 |
+
|
173 |
+
class QuickGELU(nn.Module):
|
174 |
+
can_torch_compile: bool = True
|
175 |
+
|
176 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
177 |
+
out = torch.empty_like(x)
|
178 |
+
ops.gelu_quick(out, x)
|
179 |
+
return out
|
build/torch29-cxx11-cu130-aarch64-linux/activation/__init__.py
ADDED
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
from ._ops import ops
|
4 |
+
|
5 |
+
from . import layers
|
6 |
+
|
7 |
+
|
8 |
+
def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
9 |
+
ops.silu_and_mul(out, x)
|
10 |
+
return out
|
11 |
+
|
12 |
+
|
13 |
+
def mul_and_silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
14 |
+
ops.mul_and_silu(out, x)
|
15 |
+
return out
|
16 |
+
|
17 |
+
|
18 |
+
def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
19 |
+
ops.gelu_and_mul(out, x)
|
20 |
+
return out
|
21 |
+
|
22 |
+
|
23 |
+
def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
24 |
+
ops.gelu_tanh_and_mul(out, x)
|
25 |
+
return out
|
26 |
+
|
27 |
+
|
28 |
+
def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0) -> None:
|
29 |
+
ops.fatrelu_and_mul(out, x, threshold)
|
30 |
+
return out
|
31 |
+
|
32 |
+
|
33 |
+
def gelu(out: torch.Tensor, x: torch.Tensor) -> None:
|
34 |
+
ops.gelu(out, x)
|
35 |
+
return out
|
36 |
+
|
37 |
+
def silu(out: torch.Tensor, x: torch.Tensor) -> None:
|
38 |
+
ops.silu(out, x)
|
39 |
+
return out
|
40 |
+
|
41 |
+
|
42 |
+
def gelu_tanh(out: torch.Tensor, x: torch.Tensor) -> None:
|
43 |
+
ops.gelu_tanh(out, x)
|
44 |
+
return out
|
45 |
+
|
46 |
+
|
47 |
+
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
48 |
+
ops.gelu_fast(out, x)
|
49 |
+
return out
|
50 |
+
|
51 |
+
|
52 |
+
def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
|
53 |
+
ops.gelu_new(out, x)
|
54 |
+
return out
|
55 |
+
|
56 |
+
|
57 |
+
def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
|
58 |
+
ops.gelu_quick(out, x)
|
59 |
+
return out
|
60 |
+
|
61 |
+
|
62 |
+
__all__ = [
|
63 |
+
"silu_and_mul",
|
64 |
+
"mul_and_silu",
|
65 |
+
"gelu_and_mul",
|
66 |
+
"gelu_tanh_and_mul",
|
67 |
+
"fatrelu_and_mul",
|
68 |
+
"gelu_fast",
|
69 |
+
"gelu_new",
|
70 |
+
"gelu_quick",
|
71 |
+
"gelu_tanh",
|
72 |
+
"silu",
|
73 |
+
"gelu",
|
74 |
+
"layers",
|
75 |
+
]
|
build/torch29-cxx11-cu130-aarch64-linux/activation/__pycache__/__init__.cpython-313.pyc
ADDED
Binary file (3.25 kB). View file
|
|
build/torch29-cxx11-cu130-aarch64-linux/activation/__pycache__/_ops.cpython-313.pyc
ADDED
Binary file (527 Bytes). View file
|
|
build/torch29-cxx11-cu130-aarch64-linux/activation/__pycache__/layers.cpython-313.pyc
ADDED
Binary file (8.92 kB). View file
|
|
build/torch29-cxx11-cu130-aarch64-linux/activation/_activation_320b408.abi3.so
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:73748b54059552f5983322f7dedc36ed349b38ad6fb9318301bb4965b1fe49aa
|
3 |
+
size 4094968
|
build/torch29-cxx11-cu130-aarch64-linux/activation/_ops.py
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from . import _activation_320b408
|
3 |
+
ops = torch.ops._activation_320b408
|
4 |
+
|
5 |
+
def add_op_namespace_prefix(op_name: str):
|
6 |
+
"""
|
7 |
+
Prefix op by namespace.
|
8 |
+
"""
|
9 |
+
return f"_activation_320b408::{op_name}"
|
build/torch29-cxx11-cu130-aarch64-linux/activation/layers.py
ADDED
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
from ._ops import ops
|
5 |
+
|
6 |
+
|
7 |
+
class SiluAndMul(nn.Module):
|
8 |
+
"""An activation function for SwiGLU.
|
9 |
+
|
10 |
+
The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
11 |
+
|
12 |
+
Shapes:
|
13 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
14 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
15 |
+
"""
|
16 |
+
|
17 |
+
can_torch_compile: bool = True
|
18 |
+
|
19 |
+
def forward(self, x: torch.Tensor):
|
20 |
+
d = x.shape[-1] // 2
|
21 |
+
output_shape = x.shape[:-1] + (d,)
|
22 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
23 |
+
ops.silu_and_mul(out, x)
|
24 |
+
return out
|
25 |
+
|
26 |
+
class Silu(nn.Module):
|
27 |
+
"""An activation function for SiLU.
|
28 |
+
|
29 |
+
The function computes x -> silu(x).
|
30 |
+
|
31 |
+
Shapes:
|
32 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
33 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
34 |
+
"""
|
35 |
+
|
36 |
+
can_torch_compile: bool = True
|
37 |
+
|
38 |
+
def forward(self, x: torch.Tensor):
|
39 |
+
out = torch.empty_like(x)
|
40 |
+
ops.silu(out, x)
|
41 |
+
return out
|
42 |
+
|
43 |
+
class Gelu(nn.Module):
|
44 |
+
"""An activation function for GELU.
|
45 |
+
|
46 |
+
The function computes x -> gelu(x).
|
47 |
+
|
48 |
+
Shapes:
|
49 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
50 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
51 |
+
"""
|
52 |
+
|
53 |
+
can_torch_compile: bool = True
|
54 |
+
|
55 |
+
def forward(self, x: torch.Tensor):
|
56 |
+
out = torch.empty_like(x)
|
57 |
+
ops.gelu(out, x)
|
58 |
+
return out
|
59 |
+
|
60 |
+
class GeluTanh(nn.Module):
|
61 |
+
"""An activation function for GELU with `tanh` approximation.
|
62 |
+
|
63 |
+
The function computes x -> gelu_tanh(x).
|
64 |
+
|
65 |
+
Shapes:
|
66 |
+
x: (num_tokens, d) or (batch_size, seq_len, d)
|
67 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
68 |
+
"""
|
69 |
+
|
70 |
+
can_torch_compile: bool = True
|
71 |
+
|
72 |
+
def forward(self, x: torch.Tensor):
|
73 |
+
out = torch.empty_like(x)
|
74 |
+
ops.gelu_tanh(out, x)
|
75 |
+
return out
|
76 |
+
|
77 |
+
|
78 |
+
class MulAndSilu(nn.Module):
|
79 |
+
"""An activation function for SwiGLU.
|
80 |
+
|
81 |
+
The function computes x -> x[:d] * silu(x[d:]) where d = x.shape[-1] // 2.
|
82 |
+
|
83 |
+
Shapes:
|
84 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
85 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
86 |
+
"""
|
87 |
+
|
88 |
+
can_torch_compile: bool = True
|
89 |
+
|
90 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
91 |
+
d = x.shape[-1] // 2
|
92 |
+
output_shape = x.shape[:-1] + (d,)
|
93 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
94 |
+
ops.mul_and_silu(out, x)
|
95 |
+
return out
|
96 |
+
|
97 |
+
|
98 |
+
class GeluAndMul(nn.Module):
|
99 |
+
"""An activation function for GeGLU.
|
100 |
+
|
101 |
+
The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
|
102 |
+
|
103 |
+
Shapes:
|
104 |
+
x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
|
105 |
+
return: (batch_size, seq_len, d) or (num_tokens, d)
|
106 |
+
"""
|
107 |
+
|
108 |
+
can_torch_compile: bool = True
|
109 |
+
|
110 |
+
def forward(self, x: torch.Tensor):
|
111 |
+
d = x.shape[-1] // 2
|
112 |
+
output_shape = x.shape[:-1] + (d,)
|
113 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
114 |
+
ops.gelu_and_mul(out, x)
|
115 |
+
return out
|
116 |
+
|
117 |
+
|
118 |
+
class GeluTanhAndMul(nn.Module):
|
119 |
+
can_torch_compile: bool = True
|
120 |
+
|
121 |
+
def forward(self, x: torch.Tensor):
|
122 |
+
d = x.shape[-1] // 2
|
123 |
+
output_shape = x.shape[:-1] + (d,)
|
124 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
125 |
+
ops.gelu_tanh_and_mul(out, x)
|
126 |
+
return out
|
127 |
+
|
128 |
+
|
129 |
+
class FatreluAndMul(nn.Module):
|
130 |
+
"""An activation function for FATReLU.
|
131 |
+
|
132 |
+
The function computes x -> FATReLU(x[:d]) * x[d:] where
|
133 |
+
d = x.shape[-1] // 2.
|
134 |
+
This is used in openbmb/MiniCPM-S-1B-sft.
|
135 |
+
|
136 |
+
Shapes:
|
137 |
+
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
138 |
+
return: (num_tokens, d) or (batch_size, seq_len, d)
|
139 |
+
"""
|
140 |
+
|
141 |
+
can_torch_compile: bool = True
|
142 |
+
|
143 |
+
def __init__(self, threshold: float = 0.0):
|
144 |
+
super().__init__()
|
145 |
+
self.threshold = threshold
|
146 |
+
|
147 |
+
def forward(self, x: torch.Tensor):
|
148 |
+
d = x.shape[-1] // 2
|
149 |
+
output_shape = x.shape[:-1] + (d,)
|
150 |
+
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
151 |
+
ops.fatrelu_and_mul(out, x, self.threshold)
|
152 |
+
return out
|
153 |
+
|
154 |
+
|
155 |
+
class FastGELU(nn.Module):
|
156 |
+
can_torch_compile: bool = True
|
157 |
+
|
158 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
159 |
+
out = torch.empty_like(x)
|
160 |
+
ops.gelu_fast(out, x)
|
161 |
+
return out
|
162 |
+
|
163 |
+
|
164 |
+
class NewGELU(nn.Module):
|
165 |
+
can_torch_compile: bool = True
|
166 |
+
|
167 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
168 |
+
out = torch.empty_like(x)
|
169 |
+
ops.gelu_new(out, x)
|
170 |
+
return out
|
171 |
+
|
172 |
+
|
173 |
+
class QuickGELU(nn.Module):
|
174 |
+
can_torch_compile: bool = True
|
175 |
+
|
176 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
177 |
+
out = torch.empty_like(x)
|
178 |
+
ops.gelu_quick(out, x)
|
179 |
+
return out
|