char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aman Guptaggml: add SWIGLU_CLAMP (#27930)0190529ec

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.2 KiB178 linesraw
1enable f16;
2
3#ifdef TYPE_F32
4#define DataType f32
5#endif
6#ifdef TYPE_F16
7#define DataType f16
8#endif
9
10#ifdef OP_REGLU
11fn op(a: DataType, b: DataType) -> DataType {
12    return max(a, 0) * b;
13}
14#endif
15
16#ifdef OP_GEGLU
17const SQRT_2_OVER_PI: DataType =  0.79788456080286535587989211986876;
18const GELU_COEF_A: DataType = 0.044715;
19
20fn op(a: DataType, b: DataType) -> DataType {
21    let val = SQRT_2_OVER_PI * a * (1.0 + GELU_COEF_A * a * a);
22    return 0.5 * a * (2.0 - 2.0/ (exp(2* val) + 1)) * b;
23}
24#endif
25
26#ifdef OP_SWIGLU
27fn op(a: DataType, b: DataType) -> DataType {
28    return a / (1.0 + exp(-a)) * b;
29}
30#endif
31#ifdef OP_SWIGLU_OAI
32fn op(a: f32, b: f32) -> f32 {
33    let xi = min(a, params.limit);
34    let gi = max(min(b, params.limit), -params.limit);
35    var out_glu = xi / (1.0 + exp(-xi * params.alpha));
36    out_glu = out_glu * (1.0 + gi);
37    return out_glu;
38}
39#endif
40#ifdef OP_SWIGLU_CLAMP
41fn op(a: DataType, b: DataType) -> DataType {
42    let limit = DataType(params.limit);
43    let gate = min(a, limit);
44    let up = clamp(b, -limit, limit);
45    return gate / (1.0 + exp(-gate)) * up;
46}
47#endif
48#ifdef OP_GEGLU_ERF
49const p_erf: DataType = 0.3275911;
50const a1_erf: DataType = 0.254829592;
51const a2_erf: DataType = -0.284496736;
52const a3_erf: DataType = 1.421413741;
53const a4_erf: DataType = -1.453152027;
54const a5_erf: DataType = 1.061405429;
55const SQRT_2_INV: DataType = 0.7071067811865476;
56
57fn op(a: DataType, b: DataType) -> DataType {
58    let a_div_sqr2 = a * SQRT_2_INV;
59    let sign_x = sign(a_div_sqr2);
60    let x = abs(a_div_sqr2);
61    let t = 1.0 / (1.0 + p_erf * x);
62    let y = 1.0 - (((((a5_erf * t + a4_erf) * t + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x));
63    let erf_approx = sign_x * y;
64    return 0.5 * a * (1.0 + erf_approx) * b;
65}
66#endif
67#ifdef OP_GEGLU_QUICK
68const GELU_QUICK_COEF: DataType = -1.702;
69
70fn op(a: DataType, b: DataType) -> DataType {
71    return a * (1.0 / (1.0 + exp(GELU_QUICK_COEF * a))) * b;
72}
73#endif
74
75struct Params {
76    offset_src0: u32,
77    offset_src1: u32,
78    offset_dst: u32,
79
80    // Strides (in elements)
81    stride_src01: u32,
82    stride_src02: u32,
83    stride_src03: u32,
84
85    stride_src11: u32,
86    stride_src12: u32,
87    stride_src13: u32,
88
89    stride_dst1: u32,
90    stride_dst2: u32,
91    stride_dst3: u32,
92
93    // shape of dst
94    ne: u32,
95    ne0: u32,
96    ne1: u32,
97    ne2: u32,
98
99    swapped: u32,
100    alpha: f32,
101    limit: f32,
102}
103
104@group(0) @binding(0)
105var<storage, read_write> src0: array<DataType>;
106
107#ifdef SRC_OVERLAP
108@group(0) @binding(1)
109var<storage, read_write> dst: array<DataType>;
110
111@group(0) @binding(2)
112var<uniform> params: Params;
113
114fn a_value(base: u32) -> DataType {
115    return src0[base];
116}
117
118fn b_value(base: u32) -> DataType {
119    return src0[base];
120}
121
122#elif defined(NO_SPLIT)
123@group(0) @binding(1)
124var<storage, read_write> dst: array<DataType>;
125
126@group(0) @binding(2)
127var<uniform> params: Params;
128
129fn a_value(base: u32) -> DataType {
130    let offset: u32 = select(0, params.ne0, params.swapped != 0);
131    return src0[base + offset];
132}
133
134fn b_value(base: u32) -> DataType {
135    let offset: u32 = select(params.ne0, 0, params.swapped != 0);
136    return src0[base + offset];
137}
138
139#else
140@group(0) @binding(1)
141var<storage, read_write> src1: array<DataType>;
142
143@group(0) @binding(2)
144var<storage, read_write> dst: array<DataType>;
145
146@group(0) @binding(3)
147var<uniform> params: Params;
148
149fn a_value(base: u32) -> DataType {
150    return src0[base];
151}
152
153fn b_value(base: u32) -> DataType {
154    return src1[base];
155}
156
157#endif
158
159@compute @workgroup_size(WG_SIZE)
160fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
161    if (gid.x >= params.ne) {
162        return;
163    }
164
165    var i = gid.x;
166    let i3 = i / (params.ne2 * params.ne1 * params.ne0);
167    i = i % (params.ne2 * params.ne1 * params.ne0);
168    let i2 = i / (params.ne1 * params.ne0);
169    i = i % (params.ne1 * params.ne0);
170    let i1 = i / params.ne0;
171    let i0 = i % params.ne0;
172
173    let i_a = params.offset_src0 + i3 * params.stride_src03 + i2 * params.stride_src02 + i1 * params.stride_src01 + i0;
174    let i_b = params.offset_src1 + i3 * params.stride_src13 + i2 * params.stride_src12 + i1 * params.stride_src11 + i0;
175    let i_dst = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1 + i0;
176
177    dst[i_dst] = op(a_value(i_a), b_value(i_b));
178}