char-fork/llama.cpp

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

Georgi Gerganovmetal : add concat support for quantized types (#28116)ab0b3bd3c

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
25.6 KiB480 linesraw
1#include "common.h"
2#include "dequantize.h"
3#include "quantize.h"
4
5template<typename T0, typename T1>
6kernel void kernel_cpy_t_t(
7        constant ggml_metal_kargs_cpy & args,
8        device  const char * src0,
9        device        char * dst,
10        uint3   tgpig[[threadgroup_position_in_grid]],
11        ushort3 tpitg[[thread_position_in_threadgroup]],
12        ushort3   ntg[[threads_per_threadgroup]]) {
13    const int32_t i03 = tgpig[2];
14    const int32_t i02 = tgpig[1];
15    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
16    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
17
18    if (i01 >= args.ne01) {
19        return;
20    }
21
22    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
23
24    const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
25    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
26    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
27    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
28
29    device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
30
31    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.ne00;) {
32        device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
33        dst_data[i00] = (T1) src[0];
34        break;
35    }
36}
37
38typedef decltype(kernel_cpy_t_t<float, float>) kernel_cpy_t;
39
40template [[host_name("kernel_cpy_f32_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   float>;
41template [[host_name("kernel_cpy_f32_f16")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   half>;
42template [[host_name("kernel_cpy_f32_i32")]]   kernel kernel_cpy_t kernel_cpy_t_t<float,   int32_t>;
43template [[host_name("kernel_cpy_i32_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<int32_t, float>;
44template [[host_name("kernel_cpy_i32_i32")]]   kernel kernel_cpy_t kernel_cpy_t_t<int32_t, int32_t>;
45#if defined(GGML_METAL_HAS_BF16)
46template [[host_name("kernel_cpy_f32_bf16")]]  kernel kernel_cpy_t kernel_cpy_t_t<float,   bfloat>;
47#endif
48template [[host_name("kernel_cpy_f16_f32")]]   kernel kernel_cpy_t kernel_cpy_t_t<half,    float>;
49template [[host_name("kernel_cpy_f16_f16")]]   kernel kernel_cpy_t kernel_cpy_t_t<half,    half>;
50#if defined(GGML_METAL_HAS_BF16)
51template [[host_name("kernel_cpy_bf16_f32")]]  kernel kernel_cpy_t kernel_cpy_t_t<bfloat,  float>;
52template [[host_name("kernel_cpy_bf16_bf16")]] kernel kernel_cpy_t kernel_cpy_t_t<bfloat,  bfloat>;
53#endif
54
55template<short QK,
56         typename block_q,
57         void (*quantize_func)(device const float *, device block_q &)>
58kernel void kernel_cpy_f32_q(
59        constant ggml_metal_kargs_cpy & args,
60        device const char * src0,
61        device char * dst,
62        uint3   tgpig[[threadgroup_position_in_grid]],
63        ushort3 tpitg[[thread_position_in_threadgroup]],
64        ushort3   ntg[[threads_per_threadgroup]]) {
65    const int32_t i03 = tgpig[2];
66    const int32_t i02 = tgpig[1];
67    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
68    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
69
70    if (i01 >= args.ne01) {
71        return;
72    }
73
74    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
75
76    const int32_t i3 = n / (args.ne2*args.ne1*args.ne0);
77    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0) / (args.ne1*args.ne0);
78    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0) / args.ne0;
79    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0)/QK;
80
81    device block_q * dst_data = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
82
83    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
84        device const float * src = (device const float *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + (i00*QK)*args.nb00);
85
86        quantize_func(src, dst_data[i00]);
87
88        break;
89    }
90}
91
92typedef decltype(kernel_cpy_f32_q<QK8_0,  block_q8_0,  quantize_q8_0>)  cpy_f_q_t;
93
94template [[host_name("kernel_cpy_f32_q8_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK8_0,  block_q8_0,   quantize_q8_0>;
95template [[host_name("kernel_cpy_f32_q1_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK1_0,  block_q1_0,   quantize_q1_0>;
96template [[host_name("kernel_cpy_f32_q2_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK2_0,  block_q2_0,   quantize_q2_0>;
97template [[host_name("kernel_cpy_f32_q4_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_0,  block_q4_0,   quantize_q4_0>;
98template [[host_name("kernel_cpy_f32_q4_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_1,  block_q4_1,   quantize_q4_1>;
99template [[host_name("kernel_cpy_f32_q5_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0,  block_q5_0,   quantize_q5_0>;
100template [[host_name("kernel_cpy_f32_q5_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_1,  block_q5_1,   quantize_q5_1>;
101template [[host_name("kernel_cpy_f32_iq4_nl")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_NL, block_iq4_nl, quantize_iq4_nl>;
102template [[host_name("kernel_cpy_f32_tq2_0")]]  kernel cpy_f_q_t kernel_cpy_f32_q<QK_K,   block_tq2_0,  quantize_tq2_0>;
103
104template<typename T4x4, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread T4x4 &)>
105kernel void kernel_cpy_q_f32(
106        constant ggml_metal_kargs_cpy & args,
107        device  const char * src0,
108        device        char * dst,
109        uint3   tgpig[[threadgroup_position_in_grid]],
110        ushort3 tpitg[[thread_position_in_threadgroup]],
111        ushort3   ntg[[threads_per_threadgroup]]) {
112    const int32_t i03 = tgpig[2];
113    const int32_t i02 = tgpig[1];
114    const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
115    const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
116
117    if (i01 >= args.ne01) {
118        return;
119    }
120
121    const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
122
123    const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
124    const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
125    const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
126    const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
127
128    device const block_q * src_data = (device const block_q *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
129    device       T4x4    * dst_data = (device       T4x4    *)(dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1 + i0*args.nb0);
130
131    for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
132        T4x4 temp;
133        dequantize_func(src_data + i00/nl, i00%nl, temp);
134        dst_data[i00] = temp;
135
136        break;
137    }
138}
139
140typedef decltype(kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>) cpy_q_f_t;
141
142template [[host_name("kernel_cpy_q1_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q1_0, 8, dequantize_q1_0>;
143template [[host_name("kernel_cpy_q2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q2_0, 4, dequantize_q2_0>;
144template [[host_name("kernel_cpy_q4_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>;
145template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_1, 2, dequantize_q4_1>;
146template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_0, 2, dequantize_q5_0>;
147template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_1, 2, dequantize_q5_1>;
148template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
149
150template [[host_name("kernel_cpy_tq2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_tq2_0, QK_NL, dequantize_tq2_0>;
151
152template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
153template [[host_name("kernel_cpy_q2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q2_0, 4, dequantize_q2_0>;
154template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
155template [[host_name("kernel_cpy_q4_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_1, 2, dequantize_q4_1>;
156template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_0, 2, dequantize_q5_0>;
157template [[host_name("kernel_cpy_q5_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_1, 2, dequantize_q5_1>;
158template [[host_name("kernel_cpy_q8_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q8_0, 2, dequantize_q8_0>;
159
160template [[host_name("kernel_cpy_tq2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_tq2_0, QK_NL, dequantize_tq2_0>;
161
162template<typename T>
163kernel void kernel_concat(
164        constant ggml_metal_kargs_concat & args,
165        device  const char * src0,
166        device  const char * src1,
167        device        char * dst,
168        uint3   tgpig[[threadgroup_position_in_grid]],
169        ushort3 tpitg[[thread_position_in_threadgroup]],
170        ushort3   ntg[[threads_per_threadgroup]]) {
171
172    const int i3 = tgpig.z;
173    const int i2 = tgpig.y;
174    const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
175
176    if (i1 >= args.ne1) {
177        return;
178    }
179
180    int o[4] = {0, 0, 0, 0};
181    o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
182
183    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
184        device const T * x;
185
186        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
187            x = (device const T *)(src0 + (i3       )*args.nb03 + (i2       )*args.nb02 + (i1       )*args.nb01 + (i0       )*args.nb00);
188        } else {
189            x = (device const T *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
190        }
191
192        device T * y = (device T *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
193
194        *y = *x;
195    }
196}
197
198typedef decltype(kernel_concat<float>) kernel_concat_t;
199
200template [[host_name("kernel_concat_f32")]]  kernel kernel_concat_t kernel_concat<float>;
201template [[host_name("kernel_concat_f16")]]  kernel kernel_concat_t kernel_concat<half>;
202#if defined(GGML_METAL_HAS_BF16)
203template [[host_name("kernel_concat_bf16")]] kernel kernel_concat_t kernel_concat<bfloat>;
204#endif
205template [[host_name("kernel_concat_i8")]]   kernel kernel_concat_t kernel_concat<char>;
206template [[host_name("kernel_concat_i16")]]  kernel kernel_concat_t kernel_concat<short>;
207template [[host_name("kernel_concat_i32")]]  kernel kernel_concat_t kernel_concat<int>;
208template [[host_name("kernel_concat_i64")]]  kernel kernel_concat_t kernel_concat<long>;
209
210template<typename block_q>
211kernel void kernel_concat_q(
212        constant ggml_metal_kargs_concat & args,
213        device  const char * src0,
214        device  const char * src1,
215        device        char * dst,
216        uint3   tgpig[[threadgroup_position_in_grid]],
217        ushort3 tpitg[[thread_position_in_threadgroup]],
218        ushort3   ntg[[threads_per_threadgroup]]) {
219
220    // note: for quantized types, the args are in units of blocks (nb0 == type_size)
221    const int i3 = tgpig.z;
222    const int i2 = tgpig.y;
223    const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
224
225    if (i1 >= args.ne1) {
226        return;
227    }
228
229    int o[4] = {0, 0, 0, 0};
230    o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
231
232    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
233        device const block_q * x;
234
235        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
236            x = (device const block_q *)(src0 + (i3       )*args.nb03 + (i2       )*args.nb02 + (i1       )*args.nb01 + (i0       )*args.nb00);
237        } else {
238            x = (device const block_q *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
239        }
240
241        device block_q * y = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
242
243        *y = *x;
244    }
245}
246
247typedef decltype(kernel_concat_q<block_q4_0>) kernel_concat_q_t;
248
249template [[host_name("kernel_concat_q4_0")]] kernel kernel_concat_q_t kernel_concat_q<block_q4_0>;
250template [[host_name("kernel_concat_q4_1")]] kernel kernel_concat_q_t kernel_concat_q<block_q4_1>;
251template [[host_name("kernel_concat_q5_0")]] kernel kernel_concat_q_t kernel_concat_q<block_q5_0>;
252template [[host_name("kernel_concat_q5_1")]] kernel kernel_concat_q_t kernel_concat_q<block_q5_1>;
253template [[host_name("kernel_concat_q8_0")]] kernel kernel_concat_q_t kernel_concat_q<block_q8_0>;
254
255template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
256kernel void kernel_get_rows_q(
257        constant ggml_metal_kargs_get_rows & args,
258        device const void * src0,
259        device const void * src1,
260        device       void * dst,
261        uint3               tgpig[[threadgroup_position_in_grid]],
262        ushort              tiitg[[thread_index_in_threadgroup]],
263        ushort3             ntg  [[threads_per_threadgroup]]) {
264    const int32_t iw0 = tgpig.x/args.ne10;
265    const int32_t i10 = tgpig.x%args.ne10;
266    const int32_t i11 = tgpig.y;
267    const int32_t i12 = tgpig.z;
268
269    const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
270
271    const int32_t i02 = i11;
272    const int32_t i03 = i12;
273
274    auto psrc = (device const block_q *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 +   r*args.nb01);
275    auto pdst = (device      float4x4 *) ((      device char *) dst  + i12*args.nb3  + i11*args.nb2  + i10*args.nb1);
276
277    for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
278        float4x4 temp;
279        dequantize_func(psrc + ind/nl, ind%nl, temp);
280        pdst[ind] = temp;
281
282        break;
283    }
284}
285
286template<typename T0, typename T>
287kernel void kernel_get_rows_f(
288        constant ggml_metal_kargs_get_rows & args,
289        device const void * src0,
290        device const void * src1,
291        device       void * dst,
292        uint3               tgpig[[threadgroup_position_in_grid]],
293        ushort              tiitg[[thread_index_in_threadgroup]],
294        ushort3             ntg [[threads_per_threadgroup]]) {
295    const int32_t iw0 = tgpig.x/args.ne10;
296    const int32_t i10 = tgpig.x%args.ne10;
297    const int32_t i11 = tgpig.y;
298    const int32_t i12 = tgpig.z;
299
300    const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
301
302    const int32_t i02 = i11;
303    const int32_t i03 = i12;
304
305    auto psrc = (const device T0 *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 +   r*args.nb01);
306    auto pdst = (      device T  *) ((      device char *)  dst + i12*args.nb3  + i11*args.nb2  + i10*args.nb1);
307
308    for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
309        pdst[ind] = psrc[ind];
310
311        break;
312    }
313}
314
315typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
316
317template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
318template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
319template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
320#if defined(GGML_METAL_HAS_BF16)
321template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
322#endif
323
324typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
325
326template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
327template [[host_name("kernel_get_rows_q2_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_0,    4, dequantize_q2_0>;
328template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
329template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
330template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
331template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
332template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
333template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
334template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
335template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
336template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
337template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
338template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
339template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
340template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
341template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
342template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
343template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
344template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
345template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
346template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
347template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
348template [[host_name("kernel_get_rows_tq2_0")]]   kernel get_rows_q_t kernel_get_rows_q<block_tq2_0,   QK_NL, dequantize_tq2_0>;
349
350template<typename TS, typename TI, short QK, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
351kernel void kernel_set_rows_q(
352        constant ggml_metal_kargs_set_rows & args,
353        device const  void * src0,
354        device const  void * src1,
355        device       float * dst,
356        uint3                tgpig[[threadgroup_position_in_grid]],
357        uint                 tiitg[[thread_index_in_threadgroup]],
358        uint3                tptg [[threads_per_threadgroup]]) {
359    const int32_t i03 = tgpig.z;
360    const int32_t i02 = tgpig.y;
361
362    const int32_t i12 = i03%args.ne12;
363    const int32_t i11 = i02%args.ne11;
364
365    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
366    if (i01 >= args.ne01) {
367        return;
368    }
369
370    const int32_t i10 = i01;
371    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
372
373          device block_q * dst_row = (      device block_q *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
374    const device TS      * src_row = (const device TS      *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
375
376    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
377        quantize_func(src_row + QK*ind, dst_row[ind]);
378    }
379}
380
381template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
382kernel void kernel_set_rows_q32(
383        constant ggml_metal_kargs_set_rows & args,
384        device const  void * src0,
385        device const  void * src1,
386        device       float * dst,
387        uint3                tgpig[[threadgroup_position_in_grid]],
388        uint                 tiitg[[thread_index_in_threadgroup]],
389        uint3                tptg [[threads_per_threadgroup]]) {
390    const int32_t i03 = tgpig.z;
391    const int32_t i02 = tgpig.y;
392
393    const int32_t i12 = i03%args.ne12;
394    const int32_t i11 = i02%args.ne11;
395
396    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
397    if (i01 >= args.ne01) {
398        return;
399    }
400
401    const int32_t i10 = i01;
402    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
403
404          device block_q * dst_row = (      device block_q *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
405    const device TS      * src_row = (const device TS      *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
406
407    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
408        quantize_func(src_row + 32*ind, dst_row[ind]);
409    }
410}
411
412template<typename TS, typename TI, typename TD>
413kernel void kernel_set_rows_f(
414        constant ggml_metal_kargs_set_rows & args,
415        device const  void * src0,
416        device const  void * src1,
417        device       float * dst,
418        uint3                tgpig[[threadgroup_position_in_grid]],
419        uint                 tiitg[[thread_index_in_threadgroup]],
420        uint3                tptg [[threads_per_threadgroup]]) {
421    const int32_t i03 = tgpig.z;
422    const int32_t i02 = tgpig.y;
423
424    const int32_t i12 = i03%args.ne12;
425    const int32_t i11 = i02%args.ne11;
426
427    const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
428    if (i01 >= args.ne01) {
429        return;
430    }
431
432    const int32_t i10 = i01;
433    const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
434
435          device TD * dst_row = (      device TD *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
436    const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
437
438    for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
439        dst_row[ind] = (TD) src_row[ind];
440    }
441}
442
443typedef decltype(kernel_set_rows_f<float, int64_t, float>) set_rows_f_t;
444
445template [[host_name("kernel_set_rows_f32_i64_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
446template [[host_name("kernel_set_rows_f32_i32_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
447template [[host_name("kernel_set_rows_f32_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
448template [[host_name("kernel_set_rows_f32_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
449#if defined(GGML_METAL_HAS_BF16)
450template [[host_name("kernel_set_rows_f32_i64_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
451template [[host_name("kernel_set_rows_f32_i32_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
452#endif
453
454template [[host_name("kernel_set_rows_f16_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
455template [[host_name("kernel_set_rows_f16_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
456#if defined(GGML_METAL_HAS_BF16)
457template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
458template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
459#endif
460
461typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
462
463template [[host_name("kernel_set_rows_f32_i64_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0,   quantize_q8_0>;
464template [[host_name("kernel_set_rows_f32_i32_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0,   quantize_q8_0>;
465template [[host_name("kernel_set_rows_f32_i64_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0,   quantize_q4_0>;
466template [[host_name("kernel_set_rows_f32_i32_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0,   quantize_q4_0>;
467template [[host_name("kernel_set_rows_f32_i64_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1,   quantize_q4_1>;
468template [[host_name("kernel_set_rows_f32_i32_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1,   quantize_q4_1>;
469template [[host_name("kernel_set_rows_f32_i64_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0,   quantize_q5_0>;
470template [[host_name("kernel_set_rows_f32_i32_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0,   quantize_q5_0>;
471template [[host_name("kernel_set_rows_f32_i64_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1,   quantize_q5_1>;
472template [[host_name("kernel_set_rows_f32_i32_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1,   quantize_q5_1>;
473template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
474template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
475
476typedef decltype(kernel_set_rows_q<float, int64_t, QK_K, block_tq2_0, quantize_tq2_0>) set_rows_qK_t;
477
478template [[host_name("kernel_set_rows_f32_i64_tq2_0")]]  kernel set_rows_qK_t kernel_set_rows_q<float, int64_t, QK_K, block_tq2_0, quantize_tq2_0>;
479template [[host_name("kernel_set_rows_f32_i32_tq2_0")]]  kernel set_rows_qK_t kernel_set_rows_q<float, int32_t, QK_K, block_tq2_0, quantize_tq2_0>;
480