char-fork/llama.cpp

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

Hongqiang Wangopencl: fold the gpt-oss MoE per-expert bias adds into the epilogue (op/kernel fusion) (#26431)3af988fab

char/aliceai-t5-moe
3.8 KiB79 linesraw
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}