Repository navigation
Expand file tree
/
Copy pathglu.metal
More file actions
63 lines (55 loc) · 2.08 KB
/
Copy pathglu.metal
File metadata and controls
63 lines (55 loc) · 2.08 KB
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
struct ds4_metal_args_glu {
int32_t ne00;
uint64_t nb01;
int32_t ne10;
uint64_t nb11;
int32_t ne0;
uint64_t nb1;
int32_t i00;
int32_t i10;
float alpha;
float limit;
};
// SwiGLU activation for the FFN inner state. DS4 clamps the shared expert with
// the same swiglu_limit used by routed experts.
kernel void kernel_swiglu_f32(
constant ds4_metal_args_glu & args,
device const char * src0,
device const char * src1,
device char * dst,
uint tgpig[[threadgroup_position_in_grid]],
uint tpitg[[thread_position_in_threadgroup]],
uint ntg[[threads_per_threadgroup]]) {
device const float * src0_row = (device const float *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
device const float * src1_row = (device const float *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
device float * dst_row = (device float *) ((device char *) dst + tgpig*args.nb1);
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
float x0 = src0_row[i0];
float x1 = src1_row[i0];
if (args.limit > 1.0e-6f) {
x0 = min(x0, args.limit);
x1 = clamp(x1, -args.limit, args.limit);
}
const float silu = x0 / (1.0f + exp(-x0));
dst_row[i0] = silu*x1*args.alpha;
}
}
kernel void kernel_swiglu_flat_f32(
constant ds4_metal_args_glu & args,
device const char * src0,
device const char * src1,
device char * dst,
uint i [[thread_position_in_grid]]) {
if (i >= (uint)args.ne0) return;
device const float * src0_f32 = (device const float *) src0 + args.i00;
device const float * src1_f32 = (device const float *) src1 + args.i10;
device float * dst_f32 = (device float *) dst;
float x0 = src0_f32[i];
float x1 = src1_f32[i];
if (args.limit > 1.0e-6f) {
x0 = min(x0, args.limit);
x1 = clamp(x1, -args.limit, args.limit);
}
const float silu = x0 / (1.0f + exp(-x0));
dst_f32[i] = silu*x1*args.alpha;
}