#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) 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 }
}