Hongqiang Wang
3af988fabc
* opencl: fold the gpt-oss MoE bias adds into swiglu_oai Default on, opt out with GGML_OPENCL_FUSE_MOE_BIAS_GLU=0. * opencl: fold the MoE down-projection bias into the combine Default on, opt out with GGML_OPENCL_FUSE_MOE_BIAS_COMBINE=0.
77 lines
3.0 KiB
Common Lisp
77 lines
3.0 KiB
Common Lisp
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
|
|
|
//------------------------------------------------------------------------------
|
|
// add_id(gate) + add_id(up) + swiglu_oai, fused
|
|
//
|
|
// gpt-oss-class MoE FFNs run three full passes over the same
|
|
// [n_ff, n_expert_used, n_tokens] f32 tensor: a per-expert bias add on the gate
|
|
// matmul output, the same on the up matmul output, then swiglu_oai over the
|
|
// two. Both bias adds are in-place, so each costs a full read plus a full write
|
|
// of a tensor that is only read once more. Folding them into the swiglu pass
|
|
// leaves two reads and one write instead of six passes.
|
|
//
|
|
// Grouping matches kernel_add_id: group 0 = expert slot (i1), group 1 = token
|
|
// (i2). For a contiguous destination that addressing is identical to the flat
|
|
// row walk kernel_swiglu_oai uses, since row i1 + i2*ne1 sits at
|
|
// i1*nb1 + i2*ne1*nb1.
|
|
//------------------------------------------------------------------------------
|
|
kernel void kernel_add_id_add_id_swiglu_oai(
|
|
global char * src_g,
|
|
ulong offset_g,
|
|
global char * src_gb,
|
|
ulong offset_gb,
|
|
global char * src_u,
|
|
ulong offset_u,
|
|
global char * src_ub,
|
|
ulong offset_ub,
|
|
global char * src_ids,
|
|
ulong offset_ids,
|
|
global char * dst,
|
|
ulong offsetd,
|
|
ulong nb01_g,
|
|
ulong nb02_g,
|
|
ulong nb01_u,
|
|
ulong nb02_u,
|
|
ulong nb11_g,
|
|
ulong nb11_u,
|
|
ulong nb21,
|
|
ulong nbd1,
|
|
ulong nbd2,
|
|
int ne0,
|
|
float limit,
|
|
float alpha
|
|
) {
|
|
src_g = (global char *)(src_g + offset_g);
|
|
src_gb = (global char *)(src_gb + offset_gb);
|
|
src_u = (global char *)(src_u + offset_u);
|
|
src_ub = (global char *)(src_ub + offset_ub);
|
|
src_ids = (global char *)(src_ids + offset_ids);
|
|
dst = (global char *)(dst + offsetd);
|
|
|
|
const int i1 = get_group_id(0);
|
|
const int i2 = get_group_id(1);
|
|
|
|
// The ids tensor is a view into a [n_expert, n_tokens] buffer, so its row
|
|
// stride is nb21 and the k selected ids are NOT contiguous per token.
|
|
const int i11 = *((global const int *) (src_ids + i1*sizeof(int) + i2*nb21));
|
|
|
|
global const float * g_row = (global const float *)(src_g + i1*nb01_g + i2*nb02_g);
|
|
global const float * u_row = (global const float *)(src_u + i1*nb01_u + i2*nb02_u);
|
|
global const float * gb_row = (global const float *)(src_gb + i11*nb11_g);
|
|
global const float * ub_row = (global const float *)(src_ub + i11*nb11_u);
|
|
global float * d_row = (global float *)(dst + i1*nbd1 + i2*nbd2);
|
|
|
|
for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
|
|
float x0 = g_row[i0] + gb_row[i0];
|
|
float x1 = u_row[i0] + ub_row[i0];
|
|
|
|
x0 = min(x0, limit);
|
|
x1 = max(min(x1, limit), -limit);
|
|
|
|
float out_glu = x0 / (1.0f + exp(-x0 * alpha));
|
|
out_glu = out_glu * (1.0f + x1);
|
|
|
|
d_row[i0] = out_glu;
|
|
}
|
|
}
|