wehub-resource-sync
41 lines
1.3 KiB
Metal
41 lines
1.3 KiB
Metal
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;
|
|
}
|
|
}
|