char-fork/llama.cpp

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

Hongqiang Wangopencl: quant lm_head / decode GEMV and medium-batch GEMM optimizations (speculative decoding/MTP) (#26477)95ef7fc16

char/aliceai-t5-moe
10.8 KiB369 linesraw
1#pragma OPENCL EXTENSION cl_khr_fp16 : enable
2
3#ifdef cl_intel_subgroups
4#pragma OPENCL EXTENSION cl_intel_subgroups : enable
5#else
6#pragma OPENCL EXTENSION cl_khr_subgroups : enable
7#endif
8
9#ifdef cl_intel_required_subgroup_size
10#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
11#define INTEL_GPU 1
12#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
13#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
14#elif defined(cl_qcom_reqd_sub_group_size)
15#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
16#define ADRENO_GPU 1
17#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
18#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
19#endif
20
21//------------------------------------------------------------------------------
22// rms_norm
23//------------------------------------------------------------------------------
24// This kernel depends on subgroup size.
25#ifdef INTEL_GPU
26REQD_SUBGROUP_SIZE_32
27#elif defined (ADRENO_GPU)
28REQD_SUBGROUP_SIZE_64
29#endif
30kernel void kernel_rms_norm(
31        global void * src0,
32        ulong offset0,
33        global float * dst,
34        ulong offsetd,
35        int ne00,
36        int ne01,
37        int ne02,
38        int ne03,
39        ulong nb01,
40        ulong nb02,
41        ulong nb03,
42        float eps,
43        local float * sum // Note, the size depends on number of subgroups
44) {
45    src0 = (global void*)((global char*)src0 + offset0);
46    dst = (global float*)((global char*)dst + offsetd);
47
48    int i03 = get_group_id(2);
49    int i02 = get_group_id(1);
50    int i01 = get_group_id(0);
51
52    global float4 * x = (global float4 *) ((global char *) src0 + i03*nb03 + i02*nb02 + i01*nb01);
53    global float * x_scalar = (global float *) x;
54    float4 sumf = 0;
55    float all_sum = 0;
56
57    // parallel sum
58    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
59        sumf += x[i00] * x[i00];
60    }
61    all_sum = sumf.s0 + sumf.s1 + sumf.s2 + sumf.s3;
62    all_sum = sub_group_reduce_add(all_sum);
63    if (get_sub_group_local_id() == 0) {
64        sum[get_sub_group_id()] = all_sum;
65    }
66
67    barrier(CLK_LOCAL_MEM_FENCE);
68    // broadcast
69    for (uint i = get_local_size(0) / get_max_sub_group_size() / 2; i > 0; i /= 2) {
70       if (get_local_id(0) < i) {
71           sum[get_local_id(0)] += sum[get_local_id(0) + i];
72       }
73    }
74    if (get_local_id(0) == 0) {
75        for (int i = 4 * (ne00 / 4); i < ne00; i++) {
76            sum[0] += x_scalar[i];
77        }
78        sum[0] /= ne00;
79    }
80
81    barrier(CLK_LOCAL_MEM_FENCE);
82
83    const float mean  = sum[0];
84    const float scale = 1.0f/sqrt(mean + eps);
85
86    global float4 * y = (global float4 *) (dst + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00);
87    global float * y_scalar = (global float *) y;
88    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
89        y[i00] = x[i00] * scale;
90    }
91    if (get_local_id(0) == 0) {
92        for (int i00 = 4 * (ne00 / 4); i00 < ne00; i00++) {
93            y_scalar[i00] = x_scalar[i00] * scale;
94        }
95    }
96}
97
98//------------------------------------------------------------------------------
99// rms_norm_mul
100//------------------------------------------------------------------------------
101#ifdef INTEL_GPU
102REQD_SUBGROUP_SIZE_32
103#elif defined (ADRENO_GPU)
104REQD_SUBGROUP_SIZE_64
105#endif
106kernel void kernel_rms_norm_mul(
107        global char * src0,
108        ulong offset0,
109        global char * src1,
110        ulong offset1,
111        global char * dst,
112        ulong offsetd,
113        int ne00,
114        int ne01,
115        int ne02,
116        int ne03,
117        ulong nb01,
118        ulong nb02,
119        ulong nb03,
120        int ne10,
121        int ne11,
122        int ne12,
123        int ne13,
124        ulong nb11,
125        ulong nb12,
126        ulong nb13,
127        ulong nb1,
128        ulong nb2,
129        ulong nb3,
130        float eps,
131        local float * sum
132) {
133    src0 = src0 + offset0;
134    src1 = src1 + offset1;
135    dst  = dst  + offsetd;
136
137    // The size of sum is sizeof(float)*subgroup_size.
138    // Each subgroup writes its partial sum to this array.
139    // So the number of subgroups per workgroup for this kernel cannot exceed the subgroup size.
140    // This is generally true -
141    // for subgroup size 64, workgroup size should be less than 4096 (the max is usually 1024).
142    if (get_sub_group_id() == 0) {
143        sum[get_sub_group_local_id()] = 0.0f;
144    }
145
146    int i03 = get_group_id(2);
147    int i02 = get_group_id(1);
148    int i01 = get_group_id(0);
149
150    global float4 * x = (global float4 *) (src0 + i03*nb03 + i02*nb02 + i01*nb01);
151    global float4 * f = (global float4 *) (src1 + (i03%ne13)*nb13 + (i02%ne12)*nb12 + (i01%ne11)*nb11);
152
153    float sumf = 0;
154
155    // parallel sum
156    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
157        sumf += dot(x[i00], x[i00]);
158    }
159    sumf = sub_group_reduce_add(sumf);
160
161    barrier(CLK_LOCAL_MEM_FENCE);
162
163    if (get_sub_group_local_id() == 0) {
164        sum[get_sub_group_id()] = sumf;
165    }
166
167    barrier(CLK_LOCAL_MEM_FENCE);
168
169    //for (uint i = get_local_size(0) / get_max_sub_group_size() / 2; i > 0; i /= 2) {
170    //   if (get_local_id(0) < i) {
171    //       sum[get_local_id(0)] += sum[get_local_id(0) + i];
172    //   }
173    //}
174    //if (get_local_id(0) == 0) {
175    //    sum[0] /= ne00;
176    //}
177
178    //barrier(CLK_LOCAL_MEM_FENCE);
179
180    sumf = sum[get_sub_group_local_id()];
181    sumf = sub_group_reduce_add(sumf);
182
183    float mean  = sumf / ne00;
184    float scale = 1.0f/sqrt(mean + eps);
185
186    global float4 * y = (global float4 *) (dst + i03*nb3 + i02*nb2 + i01*nb1);
187    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
188        y[i00] = (x[i00] * scale) * f[i00%(ne10/4)];
189    }
190}
191
192//------------------------------------------------------------------------------
193// rms_norm + mul (norm weight) + add (residual), fused. Mirrors
194// kernel_rms_norm_mul with an extra residual operand src2: computes
195//   y = (rmsnorm(x) * w) + g
196// in one dispatch, removing one kernel launch + one global round-trip per
197// residual block (the dominant per-layer adjacency on Gemma matformers).
198//------------------------------------------------------------------------------
199kernel void kernel_rms_norm_mul_add(
200        global char * src0,
201        ulong offset0,
202        global char * src1,
203        ulong offset1,
204        global char * src2,
205        ulong offset2,
206        global char * dst,
207        ulong offsetd,
208        int ne00,
209        int ne01,
210        int ne02,
211        int ne03,
212        ulong nb01,
213        ulong nb02,
214        ulong nb03,
215        int ne10,
216        int ne11,
217        int ne12,
218        int ne13,
219        ulong nb11,
220        ulong nb12,
221        ulong nb13,
222        int ne20,
223        int ne21,
224        int ne22,
225        int ne23,
226        ulong nb21,
227        ulong nb22,
228        ulong nb23,
229        ulong nb1,
230        ulong nb2,
231        ulong nb3,
232        float eps,
233        local float * sum
234) {
235    src0 = src0 + offset0;
236    src1 = src1 + offset1;
237    src2 = src2 + offset2;
238    dst  = dst  + offsetd;
239
240    if (get_sub_group_id() == 0) {
241        sum[get_sub_group_local_id()] = 0.0f;
242    }
243
244    int i03 = get_group_id(2);
245    int i02 = get_group_id(1);
246    int i01 = get_group_id(0);
247
248    global float4 * x = (global float4 *) (src0 + i03*nb03 + i02*nb02 + i01*nb01);
249    global float4 * f = (global float4 *) (src1 + (i03%ne13)*nb13 + (i02%ne12)*nb12 + (i01%ne11)*nb11);
250    global float4 * g = (global float4 *) (src2 + (i03%ne23)*nb23 + (i02%ne22)*nb22 + (i01%ne21)*nb21);
251
252    float sumf = 0;
253
254    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
255        sumf += dot(x[i00], x[i00]);
256    }
257    sumf = sub_group_reduce_add(sumf);
258
259    barrier(CLK_LOCAL_MEM_FENCE);
260
261    if (get_sub_group_local_id() == 0) {
262        sum[get_sub_group_id()] = sumf;
263    }
264
265    barrier(CLK_LOCAL_MEM_FENCE);
266
267    sumf = sum[get_sub_group_local_id()];
268    sumf = sub_group_reduce_add(sumf);
269
270    float mean  = sumf / ne00;
271    float scale = 1.0f/sqrt(mean + eps);
272
273    global float4 * y = (global float4 *) (dst + i03*nb3 + i02*nb2 + i01*nb1);
274    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
275        y[i00] = (x[i00] * scale) * f[i00%(ne10/4)] + g[i00%(ne20/4)];
276    }
277}
278
279//------------------------------------------------------------------------------
280// rms_norm + mul(norm weight) + add(residual) + mul(scalar scale), fused.
281// Computes y = ((rmsnorm(x) * w) + g) * s, where s is a broadcast SCALAR (e.g.
282// Gemma-4 layer_output_scale). Folds the trailing per-layer l_out scale-mul into
283// the residual-norm kernel: one extra dispatch + global round-trip saved per
284// layer. src3 points at the single scale value.
285//------------------------------------------------------------------------------
286kernel void kernel_rms_norm_mul_add_scale(
287        global char * src0,
288        ulong offset0,
289        global char * src1,
290        ulong offset1,
291        global char * src2,
292        ulong offset2,
293        global char * src3,
294        ulong offset3,
295        global char * dst,
296        ulong offsetd,
297        int ne00,
298        int ne01,
299        int ne02,
300        int ne03,
301        ulong nb01,
302        ulong nb02,
303        ulong nb03,
304        int ne10,
305        int ne11,
306        int ne12,
307        int ne13,
308        ulong nb11,
309        ulong nb12,
310        ulong nb13,
311        int ne20,
312        int ne21,
313        int ne22,
314        int ne23,
315        ulong nb21,
316        ulong nb22,
317        ulong nb23,
318        ulong nb1,
319        ulong nb2,
320        ulong nb3,
321        float eps,
322        local float * sum
323) {
324    src0 = src0 + offset0;
325    src1 = src1 + offset1;
326    src2 = src2 + offset2;
327    src3 = src3 + offset3;
328    dst  = dst  + offsetd;
329
330    const float sc = *((global float *) src3);
331
332    if (get_sub_group_id() == 0) {
333        sum[get_sub_group_local_id()] = 0.0f;
334    }
335
336    int i03 = get_group_id(2);
337    int i02 = get_group_id(1);
338    int i01 = get_group_id(0);
339
340    global float4 * x = (global float4 *) (src0 + i03*nb03 + i02*nb02 + i01*nb01);
341    global float4 * f = (global float4 *) (src1 + (i03%ne13)*nb13 + (i02%ne12)*nb12 + (i01%ne11)*nb11);
342    global float4 * g = (global float4 *) (src2 + (i03%ne23)*nb23 + (i02%ne22)*nb22 + (i01%ne21)*nb21);
343
344    float sumf = 0;
345
346    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
347        sumf += dot(x[i00], x[i00]);
348    }
349    sumf = sub_group_reduce_add(sumf);
350
351    barrier(CLK_LOCAL_MEM_FENCE);
352
353    if (get_sub_group_local_id() == 0) {
354        sum[get_sub_group_id()] = sumf;
355    }
356
357    barrier(CLK_LOCAL_MEM_FENCE);
358
359    sumf = sum[get_sub_group_local_id()];
360    sumf = sub_group_reduce_add(sumf);
361
362    float mean  = sumf / ne00;
363    float scale = 1.0f/sqrt(mean + eps);
364
365    global float4 * y = (global float4 *) (dst + i03*nb3 + i02*nb2 + i01*nb1);
366    for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
367        y[i00] = ((x[i00] * scale) * f[i00%(ne10/4)] + g[i00%(ne20/4)]) * sc;
368    }
369}