char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
3af988fab
char/aliceai-t5-moe
1// Fused MoE combine epilogue :replaces the router-weight MUL + the ( n_expert_used-1 ) 2// cross-expert ADD chain with ONE weighted-sum-across-experts pass. 3// dst [row, tok] = sum_e experts[row, e, tok] * weights[0, e, tok] 4// experts: [n_embd, n_expert_used, n_tokens] f32 (contiguous after down-proj GEMM) 5// weights: [1, n_expert_used, n_tokens] f32 6// dst: [n_embd, n_tokens] f32 7// One read of experts + one write of dst (eliminates the intermediate weighted 8// buffer and the k-1 elementwise add round-trips). Vectorized float4 over rows. 9// strides e1/e2/w1/w2/d1 are in ELEMENTS (floats). 10 11// Same weighted sum, with the per-expert bias add folded in. 12// 13// The MoE down projection's bias is applied by an in-place add_id whose only 14// consumer is this combine, so it costs a full read plus a full write of a 15// tensor that is read once more immediately afterwards. Reading the raw matmul 16// output here and adding the bias row while it is already in registers removes 17// that pass. Kept as a separate kernel so the unfused path is untouched. 18__kernel void kernel_moe_combine_bias_f32( 19 __global const char * e_buf, ulong off_e, 20 __global const char * w_buf, ulong off_w, 21 __global const char * b_buf, ulong off_b, // per-expert bias rows 22 __global const char * i_buf, ulong off_i, // expert ids 23 __global char * d_buf, ulong off_d, 24 int n_embd4, // n_embd / 4 25 int k, // n_expert_used 26 int n_tokens, 27 uint e1, uint e2, // experts strides (elements): per-expert, per-token 28 uint w1, uint w2, // weights strides (elements) 29 uint d1, // dst per-token stride (elements) 30 ulong nb_b1, // bias row stride (bytes) 31 ulong nb_i1) // ids row stride (bytes) - ids is a view, not packed 32{ 33 const uint r4 = get_global_id(0); 34 const uint tok = get_global_id(1); 35 if (r4 >= (uint)n_embd4 || tok >= (uint)n_tokens) return; 36 37 __global const float * E = (__global const float *)(e_buf + off_e) + tok*e2 + r4*4u; 38 __global const float * W = (__global const float *)(w_buf + off_w) + tok*w2; 39 __global const char * B = b_buf + off_b; 40 __global const char * I = i_buf + off_i + (ulong)tok*nb_i1; 41 42 float4 acc = (float4)(0.0f); 43 for (int e = 0; e < k; ++e) { 44 const int i11 = *((__global const int *)(I + (ulong)e*sizeof(int))); 45 __global const float * Brow = (__global const float *)(B + (ulong)i11*nb_b1) + r4*4u; 46 const float4 v = vload4(0, E + (uint)e*e1) + vload4(0, Brow); 47 acc = mad(v, (float4)(W[(uint)e*w1]), acc); 48 } 49 50 __global float * D = (__global float *)(d_buf + off_d) + tok*d1 + r4*4u; 51 vstore4(acc, 0, D); 52} 53 54__kernel void kernel_moe_combine_f32( 55 __global const char * e_buf, ulong off_e, 56 __global const char * w_buf, ulong off_w, 57 __global char * d_buf, ulong off_d, 58 int n_embd4, // n_embd / 4 59 int k, // n_expert_used 60 int n_tokens, 61 uint e1, uint e2, // experts strides (elements): per-expert, per-token 62 uint w1, uint w2, // weights strides (elements) 63 uint d1) // dst per-token stride (elements) 64{ 65 const uint r4 = get_global_id(0); 66 const uint tok = get_global_id(1); 67 if (r4 >= (uint)n_embd4 || tok >= (uint)n_tokens) return; 68 69 __global const float * E = (__global const float *)(e_buf + off_e) + tok*e2 + r4*4u; 70 __global const float * W = (__global const float *)(w_buf + off_w) + tok*w2; 71 72 float4 acc = (float4)(0.0f); 73 for (int e = 0; e < k; ++e) { 74 acc = mad(vload4(0, E + (uint)e*e1), (float4)(W[(uint)e*w1]), acc); 75 } 76 77 __global float * D = (__global float *)(d_buf + off_d) + tok*d1 + r4*4u; 78 vstore4(acc, 0, D); 79}