char-fork/llama.cpp

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

Georgi Gerganovmetal : single-source fusion table + fusion debug rework (#28164)a2878d30d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
26.4 KiB1353 linesraw
1#ifndef GGML_METAL_IMPL
2#define GGML_METAL_IMPL
3
4// kernel parameters for mat-mat threadgroups
5//
6// TODO: become function constants
7
8#define SZ_SIMDGROUP 16
9#define N_MM_NK 2
10#define N_MM_NK_TOTAL (SZ_SIMDGROUP * N_MM_NK)
11
12#define N_MM_BLOCK_X 4
13#define N_MM_BLOCK_Y 2
14#define N_MM_SIMD_GROUP_X 2
15#define N_MM_SIMD_GROUP_Y 2
16
17// kernel parameters for mat-vec threadgroups
18//
19// N_R0: number of src0 rows to process per simdgroup
20// N_SG: number of simdgroups per threadgroup
21//
22// TODO: for optimal performance, become function of the device and work size
23
24#define N_R0_Q1_0 8
25#define N_SG_Q1_0 2
26
27#define N_R0_Q2_0 8
28#define N_SG_Q2_0 2
29
30#define N_R0_Q4_0 4
31#define N_SG_Q4_0 2
32
33#define N_R0_Q4_1 4
34#define N_SG_Q4_1 2
35
36#define N_R0_Q5_0 4
37#define N_SG_Q5_0 2
38
39#define N_R0_Q5_1 4
40#define N_SG_Q5_1 2
41
42#define N_R0_Q8_0 2
43#define N_SG_Q8_0 4
44
45#define N_R0_MXFP4 2
46#define N_SG_MXFP4 2
47
48#define N_R0_Q2_K 4
49#define N_SG_Q2_K 2
50
51#define N_R0_Q3_K 2
52#define N_SG_Q3_K 2
53
54#define N_R0_Q4_K 2
55#define N_SG_Q4_K 2
56
57#define N_R0_Q5_K 1
58#define N_SG_Q5_K 2
59
60#define N_R0_Q6_K 2
61#define N_SG_Q6_K 2
62
63#define N_R0_IQ1_S 4
64#define N_SG_IQ1_S 2
65#define N_R0_IQ1_S_SPLIT 8
66
67#define N_R0_IQ1_M 4
68#define N_SG_IQ1_M 2
69#define N_R0_IQ1_M_SPLIT 8
70
71#define N_R0_IQ2_XXS 4
72#define N_SG_IQ2_XXS 2
73#define N_R0_IQ2_XXS_SPLIT 8
74
75#define N_R0_IQ2_XS 4
76#define N_SG_IQ2_XS 2
77#define N_R0_IQ2_XS_SPLIT 8
78
79#define N_R0_IQ2_S 4
80#define N_SG_IQ2_S 2
81#define N_R0_IQ2_S_SPLIT 8
82
83#define N_R0_IQ3_XXS 4
84#define N_SG_IQ3_XXS 2
85#define N_R0_IQ3_XXS_SPLIT 8
86
87#define N_R0_IQ3_S 4
88#define N_SG_IQ3_S 2
89#define N_R0_IQ3_S_SPLIT 8
90
91#define N_R0_IQ4_NL 2
92#define N_SG_IQ4_NL 2
93
94#define N_R0_IQ4_XS 2
95#define N_SG_IQ4_XS 2
96
97#define N_R0_TQ2_0 4
98#define N_SG_TQ2_0 2
99
100// function constants offsets
101#define FC_FLASH_ATTN_EXT_PAD          100
102#define FC_FLASH_ATTN_EXT_BLK          200
103#define FC_FLASH_ATTN_EXT              300
104#define FC_FLASH_ATTN_EXT_VEC          400
105#define FC_FLASH_ATTN_EXT_VEC_REDUCE   500
106#define FC_MUL_MV                      600
107#define FC_MUL_MM                      700
108#define FC_ROPE                        800
109#define FC_SSM_CONV                    900
110#define FC_SOLVE_TRI                   1000
111#define FC_COUNT_EQUAL                 1100
112#define FC_UNARY                       1200
113#define FC_BIN                         1300
114#define FC_SUM_ROWS                    1400
115#define FC_UPSCALE                     1500
116#define FC_GATED_DELTA_NET             1600
117
118// op-specific constants
119#define OP_FLASH_ATTN_EXT_NQPSG 8
120#define OP_FLASH_ATTN_EXT_NCPSG 64
121
122#define OP_FLASH_ATTN_EXT_VEC_NQPSG 1
123#define OP_FLASH_ATTN_EXT_VEC_NCPSG 32
124
125#define OP_LIGHTNING_INDEXER_DK    128
126#define OP_LIGHTNING_INDEXER_NH     64
127#define OP_LIGHTNING_INDEXER_NHPTG   8
128#define OP_LIGHTNING_INDEXER_NKPSG   8
129#define OP_LIGHTNING_INDEXER_NSG     8
130#define OP_LIGHTNING_INDEXER_NBPTG   8
131
132#define OP_UNARY_NUM_SCALE      10
133#define OP_UNARY_NUM_FILL       11
134#define OP_UNARY_NUM_CLAMP      12
135#define OP_UNARY_NUM_SQR        13
136#define OP_UNARY_NUM_SQRT       14
137#define OP_UNARY_NUM_SIN        15
138#define OP_UNARY_NUM_COS        16
139#define OP_UNARY_NUM_LOG        17
140#define OP_UNARY_NUM_LEAKY_RELU 18
141
142#define OP_UNARY_NUM_TANH        100
143#define OP_UNARY_NUM_RELU        101
144#define OP_UNARY_NUM_SIGMOID     102
145#define OP_UNARY_NUM_GELU        103
146#define OP_UNARY_NUM_GELU_ERF    104
147#define OP_UNARY_NUM_GELU_QUICK  105
148#define OP_UNARY_NUM_SILU        106
149#define OP_UNARY_NUM_ELU         107
150#define OP_UNARY_NUM_NEG         108
151#define OP_UNARY_NUM_ABS         109
152#define OP_UNARY_NUM_SGN         110
153#define OP_UNARY_NUM_STEP        111
154#define OP_UNARY_NUM_HARDSWISH   112
155#define OP_UNARY_NUM_HARDSIGMOID 113
156#define OP_UNARY_NUM_EXP         114
157#define OP_UNARY_NUM_SOFTPLUS    115
158#define OP_UNARY_NUM_EXPM1       116
159#define OP_UNARY_NUM_FLOOR       117
160#define OP_UNARY_NUM_CEIL        118
161#define OP_UNARY_NUM_ROUND       119
162#define OP_UNARY_NUM_TRUNC       120
163#define OP_UNARY_NUM_XIELU       121
164
165#define OP_SUM_ROWS_NUM_SUM_ROWS 10
166#define OP_SUM_ROWS_NUM_MEAN     11
167
168#define OP_SSM_SCAN_SSD_CS  64 // Metal-specific; Chunk Size; 64 is largest multiple of 8 (simdgroup tile) fitting into 32 KiB Metal threadgroup mem limit (~26.75 KiB shared mem; see smem layout comment in kernel_ssm_scan_ssd_mma_f32)
169#define OP_SSM_SCAN_SSD_HD  64 // Metal-specific; Head Dim the MMA kernel is specialized for (Mamba-2); use_mma gates on d_inner == this
170#define OP_SSM_SCAN_SSD_NSG 4  // Metal-specific; Number of SimdGroups per threadgroup; NSG*32 == threads dispatched per threadgroup
171
172// kernel argument structs
173//
174// - element counters (e.g. ne00) typically use int32_t to reduce register usage
175//   however, be careful from int overflows when using those in the kernel implementation
176//
177// - strides (e.g. nb00) use uint64_t
178
179typedef struct {
180    int32_t  ne00;
181    int32_t  ne01;
182    int32_t  ne02;
183    int32_t  ne03;
184    uint64_t nb00;
185    uint64_t nb01;
186    uint64_t nb02;
187    uint64_t nb03;
188    int32_t  ne10;
189    int32_t  ne11;
190    int32_t  ne12;
191    int32_t  ne13;
192    uint64_t nb10;
193    uint64_t nb11;
194    uint64_t nb12;
195    uint64_t nb13;
196    int32_t  ne0;
197    int32_t  ne1;
198    int32_t  ne2;
199    int32_t  ne3;
200    uint64_t nb0;
201    uint64_t nb1;
202    uint64_t nb2;
203    uint64_t nb3;
204    int32_t  dim;
205} ggml_metal_kargs_concat;
206
207typedef struct {
208    int32_t  ne00;
209    int32_t  ne01;
210    int32_t  ne02;
211    int32_t  ne03;
212    uint64_t nb00;
213    uint64_t nb01;
214    uint64_t nb02;
215    uint64_t nb03;
216    int32_t  ne0;
217    int32_t  ne1;
218    int32_t  ne2;
219    int32_t  ne3;
220    uint64_t nb0;
221    uint64_t nb1;
222    uint64_t nb2;
223    uint64_t nb3;
224    float    slope;
225    float    scale;
226    float    bias;
227    float    val;
228    float    min;
229    float    max;
230} ggml_metal_kargs_unary;
231
232typedef struct {
233    int32_t  ne00;
234    int32_t  ne01;
235    int32_t  ne02;
236    int32_t  ne03;
237    uint64_t nb00;
238    uint64_t nb01;
239    uint64_t nb02;
240    uint64_t nb03;
241    int32_t  ne10;
242    int32_t  ne11;
243    int32_t  ne12;
244    int32_t  ne13;
245    uint64_t nb10;
246    uint64_t nb11;
247    uint64_t nb12;
248    uint64_t nb13;
249    int32_t  ne0;
250    int32_t  ne1;
251    int32_t  ne2;
252    int32_t  ne3;
253    uint64_t nb0;
254    uint64_t nb1;
255    uint64_t nb2;
256    uint64_t nb3;
257    uint64_t offs;
258    uint64_t o1[8];
259} ggml_metal_kargs_bin;
260
261typedef struct {
262    int64_t ne0;
263    int64_t ne1;
264    size_t nb01;
265    size_t nb02;
266    size_t nb11;
267    size_t nb21;
268} ggml_metal_kargs_add_id;
269
270typedef struct {
271    int32_t  ne00;
272    int32_t  ne01;
273    int32_t  ne02;
274    int32_t  ne03;
275    uint64_t nb00;
276    uint64_t nb01;
277    uint64_t nb02;
278    uint64_t nb03;
279    int32_t  ne0;
280    int32_t  ne1;
281    int32_t  ne2;
282    int32_t  ne3;
283    uint64_t nb0;
284    uint64_t nb1;
285    uint64_t nb2;
286    uint64_t nb3;
287} ggml_metal_kargs_repeat;
288
289typedef struct {
290    int64_t  nk0;
291    int64_t  ne00;
292    int64_t  ne01;
293    int64_t  ne02;
294    int64_t  ne03;
295    uint64_t nb00;
296    uint64_t nb01;
297    uint64_t nb02;
298    uint64_t nb03;
299    int64_t  ne0;
300    int64_t  ne1;
301    int64_t  ne2;
302    int64_t  ne3;
303    uint64_t nb0;
304    uint64_t nb1;
305    uint64_t nb2;
306    uint64_t nb3;
307} ggml_metal_kargs_cpy;
308
309typedef struct {
310    int64_t  ne10;
311    int64_t  ne11;
312    int64_t  ne12;
313    uint64_t nb10;
314    uint64_t nb11;
315    uint64_t nb12;
316    uint64_t nb13;
317    uint64_t nb1;
318    uint64_t nb2;
319    uint64_t nb3;
320    uint64_t offs;
321    bool     inplace;
322} ggml_metal_kargs_set;
323
324typedef struct {
325    int32_t  ne00;
326    int32_t  ne01;
327    int32_t  ne02;
328    int32_t  ne03;
329    uint64_t nb00;
330    uint64_t nb01;
331    uint64_t nb02;
332    uint64_t nb03;
333    int32_t  ne0;
334    int32_t  ne1;
335    int32_t  ne2;
336    int32_t  ne3;
337    uint64_t nb0;
338    uint64_t nb1;
339    uint64_t nb2;
340    uint64_t nb3;
341    int32_t  n_past;
342    int32_t  n_dims;
343    int32_t  n_offs;
344    int32_t  n_ctx_orig;
345    float    freq_base;
346    float    freq_scale;
347    float    ext_factor;
348    float    attn_factor;
349    float    beta_fast;
350    float    beta_slow;
351    int32_t  sect_0;
352    int32_t  sect_1;
353    int32_t  sect_2;
354    int32_t  sect_3;
355    bool     src2;
356    bool     inplace;
357} ggml_metal_kargs_rope;
358
359typedef struct {
360    int32_t  ne0;
361    int32_t  ne1;
362    int32_t  ne2;
363    int32_t  ne3;
364    uint64_t nb0;
365    uint64_t nb1;
366    uint64_t nb2;
367    uint64_t nb3;
368    int32_t  nblocks;
369} ggml_metal_kargs_flash_attn_ext_kv_f16;
370
371typedef struct {
372    int32_t  ne11;
373    int32_t  ne_12_2; // assume K and V are same shape
374    int32_t  ne_12_3;
375    uint64_t nb11;
376    uint64_t nb12;
377    uint64_t nb13;
378    uint64_t nb21;
379    uint64_t nb22;
380    uint64_t nb23;
381    int32_t  ne31;
382    int32_t  ne32;
383    int32_t  ne33;
384    uint64_t nb31;
385    uint64_t nb32;
386    uint64_t nb33;
387} ggml_metal_kargs_flash_attn_ext_pad;
388
389typedef struct {
390    int32_t  ne01;
391    int32_t  ne30;
392    int32_t  ne31;
393    int32_t  ne32;
394    int32_t  ne33;
395    uint64_t nb31;
396    uint64_t nb32;
397    uint64_t nb33;
398} ggml_metal_kargs_flash_attn_ext_blk;
399
400typedef struct {
401    int32_t  ne01;
402    int32_t  ne02;
403    int32_t  ne03;
404    uint64_t nb01;
405    uint64_t nb02;
406    uint64_t nb03;
407    int32_t  ne11;
408    int32_t  ne_12_2; // assume K and V are same shape
409    int32_t  ne_12_3;
410    int32_t  ns10;
411    uint64_t nb11;
412    uint64_t nb12;
413    uint64_t nb13;
414    int32_t  ns20;
415    uint64_t nb21;
416    uint64_t nb22;
417    uint64_t nb23;
418    int32_t  ne31;
419    int32_t  ne32;
420    int32_t  ne33;
421    uint64_t nb31;
422    uint64_t nb32;
423    uint64_t nb33;
424    int32_t  ne1;
425    int32_t  ne2;
426    int32_t  ne3;
427    float    scale;
428    float    max_bias;
429    float    m0;
430    float    m1;
431    int32_t  n_head_log2;
432    float    logit_softcap;
433} ggml_metal_kargs_flash_attn_ext;
434
435typedef struct {
436    int32_t  ne01;
437    int32_t  ne02;
438    int32_t  ne03;
439    uint64_t nb01;
440    uint64_t nb02;
441    uint64_t nb03;
442    int32_t  ne11;
443    int32_t  ne_12_2; // assume K and V are same shape
444    int32_t  ne_12_3;
445    int32_t  ns10;
446    uint64_t nb11;
447    uint64_t nb12;
448    uint64_t nb13;
449    int32_t  ns20;
450    uint64_t nb21;
451    uint64_t nb22;
452    uint64_t nb23;
453    int32_t  ne31;
454    int32_t  ne32;
455    int32_t  ne33;
456    uint64_t nb31;
457    uint64_t nb32;
458    uint64_t nb33;
459    int32_t  ne1;
460    int32_t  ne2;
461    int32_t  ne3;
462    float    scale;
463    float    max_bias;
464    float    m0;
465    float    m1;
466    int32_t  n_head_log2;
467    float    logit_softcap;
468    int32_t  n_kv_max_padded;
469} ggml_metal_kargs_flash_attn_ext_vec;
470
471typedef struct {
472    int32_t  ne30;
473    int32_t  ne31;
474    int32_t  ne32;
475    int32_t  ne33;
476    uint64_t nb31;
477    uint64_t nb32;
478    uint64_t nb33;
479    int32_t  n_kv_max;
480    int32_t  n_kv_max_padded;
481} ggml_metal_kargs_flash_attn_ext_vec_idx;
482
483typedef struct {
484    int32_t  nrows;
485} ggml_metal_kargs_flash_attn_ext_vec_reduce;
486
487typedef struct {
488    int32_t  ne00;
489    int32_t  ne02;
490    uint64_t nb01;
491    uint64_t nb02;
492    uint64_t nb03;
493    int32_t  ne12;
494    uint64_t nb10;
495    uint64_t nb11;
496    uint64_t nb12;
497    uint64_t nb13;
498    int32_t  ne0;
499    int32_t  ne1;
500    int16_t  r2;
501    int16_t  r3;
502} ggml_metal_kargs_mul_mm;
503
504typedef struct {
505    int32_t  ne00;
506    int32_t  ne01;
507    int32_t  ne02;
508    uint64_t nb00;
509    uint64_t nb01;
510    uint64_t nb02;
511    uint64_t nb03;
512    int32_t  ne10;
513    int32_t  ne11;
514    int32_t  ne12;
515    uint64_t nb10;
516    uint64_t nb11;
517    uint64_t nb12;
518    uint64_t nb13;
519    int32_t  ne0;
520    int32_t  ne1;
521    int32_t  nr0;
522    int16_t  r2;
523    int16_t  r3;
524} ggml_metal_kargs_mul_mv;
525
526typedef struct {
527    int32_t  ne00;
528    int32_t  ne01;
529    int32_t  ne02;
530    uint64_t nb00;
531    uint64_t nb01;
532    uint64_t nb02;
533    uint64_t nb03;
534    int32_t  ne10;
535    int32_t  ne11;
536    int32_t  ne12;
537    uint64_t nb10;
538    uint64_t nb11;
539    uint64_t nb12;
540    uint64_t nb13;
541    int32_t  ne0;
542    int32_t  ne1;
543    int16_t  r2;
544    int16_t  r3;
545} ggml_metal_kargs_mul_mv_ext;
546
547typedef struct {
548    int32_t  ne02;
549    int32_t  ne10;
550    int32_t  ne11;  // n_expert_used (bcast)
551    uint64_t nb11;
552    uint64_t nb12;
553    int32_t  ne21; // n_tokens
554    int32_t  ne20;  // n_expert_used
555    uint64_t nb21;
556} ggml_metal_kargs_mul_mm_id_map0;
557
558typedef struct {
559    int32_t  ne00;
560    int32_t  ne02;
561    uint64_t nb01;
562    uint64_t nb02;
563    uint64_t nb03;
564    int32_t  ne11;
565    uint64_t nb10;
566    uint64_t nb11;
567    uint64_t nb12;
568    uint64_t nb13;
569    int32_t  ne20;
570    int32_t  ne21;
571    int32_t  ne0;
572    int32_t  ne1;
573    int16_t  r2;
574    int16_t  r3;
575} ggml_metal_kargs_mul_mm_id;
576
577typedef struct {
578    int32_t  nei0;
579    int32_t  nei1;
580    uint64_t nbi1;
581    int32_t  ne00;
582    int32_t  ne01;
583    int32_t  ne02;
584    uint64_t nb00;
585    uint64_t nb01;
586    uint64_t nb02;
587    int32_t  ne10;
588    int32_t  ne11;
589    int32_t  ne12;
590    int32_t  ne13;
591    uint64_t nb10;
592    uint64_t nb11;
593    uint64_t nb12;
594    int32_t  ne0;
595    int32_t  ne1;
596    uint64_t nb1;
597    int32_t  nr0;
598} ggml_metal_kargs_mul_mv_id;
599
600// NORM
601// RMS_NORM
602typedef struct {
603    int32_t  ne00;
604    int32_t  ne00_t;
605    uint64_t nb1;
606    uint64_t nb2;
607    uint64_t nb3;
608    float    eps;
609    int32_t  nef1[3];
610    int32_t  nef2[3];
611    int32_t  nef3[3];
612    uint64_t nbf1[3];
613    uint64_t nbf2[3];
614    uint64_t nbf3[3];
615} ggml_metal_kargs_norm;
616
617typedef struct {
618    int32_t  ne00;
619    int32_t  ne01;
620    int32_t  ne02;
621    int32_t  ne03;
622    uint64_t nb00;
623    uint64_t nb01;
624    uint64_t nb02;
625    uint64_t nb03;
626    int32_t  ne0;
627    int32_t  ne1;
628    int32_t  ne2;
629    int32_t  ne3;
630    uint64_t nb0;
631    uint64_t nb1;
632    uint64_t nb2;
633    uint64_t nb3;
634    float    eps;
635} ggml_metal_kargs_l2_norm;
636
637typedef struct {
638    int64_t  ne00;
639    int64_t  ne01;
640    int64_t  ne02;
641    uint64_t nb00;
642    uint64_t nb01;
643    uint64_t nb02;
644    int32_t  ngrp;
645    float    eps;
646} ggml_metal_kargs_group_norm;
647
648typedef struct {
649    int32_t  IC;
650    int32_t  IL;
651    int32_t  K;
652    int32_t  s0;
653    uint64_t nb0;
654    uint64_t nb1;
655} ggml_metal_kargs_conv_transpose_1d;
656
657typedef struct {
658    int32_t  T_in;
659    int32_t  T_out;
660    int32_t  OC;
661    int32_t  K;
662    int32_t  K_OC;
663    int32_t  s0;
664    int32_t  p0;
665} ggml_metal_kargs_col2im_1d;
666
667typedef struct {
668    int32_t T;
669    int32_t C;
670} ggml_metal_kargs_snake;
671
672typedef struct {
673    int32_t  IC;
674    int32_t  IH;
675    int32_t  IW;
676    int32_t  KH;
677    int32_t  KW;
678    int32_t  OC;
679    int32_t  s0;
680    uint64_t nb0;
681    uint64_t nb1;
682    uint64_t nb2;
683    uint64_t nb3;
684} ggml_metal_kargs_conv_transpose_2d;
685
686typedef struct {
687    uint64_t nb00;
688    uint64_t nb01;
689    uint64_t nb02;
690    uint64_t nb03;
691    uint64_t nb10;
692    uint64_t nb11;
693    uint64_t nb12;
694    uint64_t nb13;
695    uint64_t nb0;
696    uint64_t nb1;
697    uint64_t nb2;
698    uint64_t nb3;
699    int32_t  IW;
700    int32_t  IH;
701    int32_t  KW;
702    int32_t  KH;
703    int32_t  IC;
704    int32_t  OC;
705    int32_t  OW;
706    int32_t  OH;
707    int32_t  N;
708    int32_t  s0;
709    int32_t  s1;
710    int32_t  p0;
711    int32_t  p1;
712    int32_t  d0;
713    int32_t  d1;
714} ggml_metal_kargs_conv_2d;
715
716typedef struct {
717    uint64_t nb00;  // kernel strides
718    uint64_t nb01;
719    uint64_t nb02;
720    uint64_t nb10;  // input strides
721    uint64_t nb11;
722    uint64_t nb12;
723    uint64_t nb13;
724    uint64_t nb0;   // output strides
725    uint64_t nb1;
726    uint64_t nb2;
727    uint64_t nb3;
728    int32_t  IW;    // input width
729    int32_t  IH;    // input height
730    int32_t  KW;    // kernel width
731    int32_t  KH;    // kernel height
732    int32_t  C;     // channels (IC == OC for depthwise)
733    int32_t  OW;    // output width
734    int32_t  OH;    // output height
735    int32_t  N;     // batch size
736    int32_t  s0;    // stride x
737    int32_t  s1;    // stride y
738    int32_t  p0;    // padding x
739    int32_t  p1;    // padding y
740    int32_t  d0;    // dilation x
741    int32_t  d1;    // dilation y
742} ggml_metal_kargs_conv_2d_dw;
743
744typedef struct {
745    uint64_t  ofs0;
746    uint64_t  ofs1;
747    int32_t  IW;
748    int32_t  IH;
749    int32_t  CHW;
750    int32_t  s0;
751    int32_t  s1;
752    int32_t  p0;
753    int32_t  p1;
754    int32_t  d0;
755    int32_t  d1;
756    int32_t  N;
757    int32_t  KH;
758    int32_t  KW;
759    int32_t  KHW; // KH * KW, pre-computed on CPU to save GPU resources
760} ggml_metal_kargs_im2col;
761
762typedef struct {
763    int32_t  IW;
764    int32_t  IH;
765    int32_t  ID;
766    int32_t  OW;
767    int32_t  OH;
768    int32_t  OD;
769    int32_t  KW;
770    int32_t  KH;
771    int32_t  KD;
772    int32_t  s0;
773    int32_t  s1;
774    int32_t  s2;
775    int32_t  p0;
776    int32_t  p1;
777    int32_t  p2;
778    int32_t  d0;
779    int32_t  d1;
780    int32_t  d2;
781    int32_t  IC;
782    int32_t  N;
783    int32_t  OC;
784    uint64_t nb00;
785    uint64_t nb01;
786    uint64_t nb02;
787    uint64_t nb03;
788    uint64_t nb10;
789    uint64_t nb11;
790    uint64_t nb12;
791    uint64_t nb13;
792    uint64_t nb0;
793    uint64_t nb1;
794    uint64_t nb2;
795    uint64_t nb3;
796} ggml_metal_kargs_conv_3d;
797
798typedef struct{
799    int32_t  ne00;
800    uint64_t nb01;
801    int32_t  ne10;
802    uint64_t nb11;
803    int32_t  ne0;
804    uint64_t nb1;
805    int32_t  i00;
806    int32_t  i10;
807    float    alpha;
808    float    limit;
809} ggml_metal_kargs_glu;
810
811typedef struct {
812    uint64_t np;
813} ggml_metal_kargs_sum;
814
815typedef struct {
816    int64_t  ne00;
817    int64_t  ne01;
818    int64_t  ne02;
819    int64_t  ne03;
820    uint64_t nb00;
821    uint64_t nb01;
822    uint64_t nb02;
823    uint64_t nb03;
824    int64_t  ne0;
825    int64_t  ne1;
826    int64_t  ne2;
827    int64_t  ne3;
828    uint64_t nb0;
829    uint64_t nb1;
830    uint64_t nb2;
831    uint64_t nb3;
832} ggml_metal_kargs_sum_rows;
833
834typedef struct {
835    int64_t  ne00;
836    int64_t  ne01;
837    int64_t  ne02;
838    int64_t  ne03;
839    uint64_t nb00;
840    uint64_t nb01;
841    uint64_t nb02;
842    uint64_t nb03;
843    int64_t  net0;
844    int64_t  net1;
845    int64_t  net2;
846    int64_t  net3;
847    uint64_t nbt0;
848    uint64_t nbt1;
849    uint64_t nbt2;
850    uint64_t nbt3;
851    bool     outb;
852} ggml_metal_kargs_cumsum_blk;
853
854typedef struct {
855    int64_t  ne00;
856    int64_t  ne01;
857    int64_t  ne02;
858    int64_t  ne03;
859    uint64_t nb00;
860    uint64_t nb01;
861    uint64_t nb02;
862    uint64_t nb03;
863    int64_t  net0;
864    int64_t  net1;
865    int64_t  net2;
866    int64_t  net3;
867    uint64_t nbt0;
868    uint64_t nbt1;
869    uint64_t nbt2;
870    uint64_t nbt3;
871} ggml_metal_kargs_cumsum_add;
872
873typedef struct {
874    int32_t  ne00;
875    int32_t  ne01;
876    int32_t  ne02;
877    uint64_t nb01;
878    uint64_t nb02;
879    uint64_t nb03;
880    int32_t  ne11;
881    int32_t  ne12;
882    int32_t  ne13;
883    uint64_t nb11;
884    uint64_t nb12;
885    uint64_t nb13;
886    uint64_t nb1;
887    uint64_t nb2;
888    uint64_t nb3;
889    float    scale;
890    float    max_bias;
891    float    m0;
892    float    m1;
893    int32_t  n_head_log2;
894} ggml_metal_kargs_soft_max;
895
896typedef struct {
897    int64_t  ne00;
898    int64_t  ne01;
899    int64_t  ne02;
900    uint64_t nb00;
901    uint64_t nb01;
902    uint64_t nb02;
903    int64_t  ne10;
904    int64_t  ne11;
905    uint64_t nb10;
906    uint64_t nb11;
907    int64_t  ne0;
908    int64_t  ne1;
909    int64_t  ne2;
910    uint64_t nb0;
911    uint64_t nb1;
912    uint64_t nb2;
913} ggml_metal_kargs_ssm_conv;
914
915typedef struct {
916    int64_t  d_state;
917    int64_t  d_inner;
918    int64_t  n_head;
919    int64_t  n_group;
920    int64_t  n_seq_tokens;
921    int64_t  n_seq_tokens_total;
922    int64_t  token_offset;
923    int64_t  n_seqs;
924    int64_t  K;
925    uint64_t s_off;
926    uint64_t nb00;
927    uint64_t nb01;
928    uint64_t nb02;
929    uint64_t nb03;
930    uint64_t nb10;
931    uint64_t nb11;
932    uint64_t nb12;
933    uint64_t ns12;
934    uint64_t nb13;
935    uint64_t nb20;
936    uint64_t nb21;
937    uint64_t ns21;
938    uint64_t nb22;
939    int64_t  ne30;
940    uint64_t nb31;
941    uint64_t nb41;
942    uint64_t nb42;
943    uint64_t ns42;
944    uint64_t nb43;
945    uint64_t nb51;
946    uint64_t nb52;
947    uint64_t ns52;
948    uint64_t nb53;
949    uint64_t nb0;
950} ggml_metal_kargs_ssm_scan;
951
952typedef struct {
953    int32_t  ne00;
954    int32_t  ne01;
955    int32_t  ne02;
956    int32_t  ne03;
957    uint64_t nb00;
958    uint64_t nb01;
959    uint64_t nb02;
960    uint64_t nb03;
961    int32_t  ne10;
962    int32_t  ne11;
963    int32_t  ne12;
964    int32_t  ne13;
965    uint64_t nb10;
966    uint64_t nb11;
967    uint64_t nb12;
968    uint64_t nb13;
969    int32_t  ne20;
970    int32_t  ne21;
971    int32_t  ne22;
972    int32_t  ne23;
973    uint64_t nb20;
974    uint64_t nb21;
975    uint64_t nb22;
976    uint64_t nb23;
977    int32_t  ns02;
978    int32_t  ns12;
979    int32_t  ns22;
980    int32_t  ne0;
981    int32_t  ne1;
982    int32_t  ne2;
983    int32_t  ne3;
984    uint64_t nb0;
985    uint64_t nb1;
986    uint64_t nb2;
987    uint64_t nb3;
988    uint64_t nb_out; // 0 => snapshots are appended after the attn scores (unfused)
989} ggml_metal_kargs_gated_delta_net;
990
991typedef struct {
992    int32_t  ne00;
993    int32_t  ne01;
994    int32_t  ne02;
995    int32_t  ne03;
996    uint64_t nb00;
997    uint64_t nb01;
998    uint64_t nb02;
999    uint64_t nb03;
1000    int32_t  ne10;
1001    int32_t  ne11;
1002    int32_t  ne12;
1003    int32_t  ne13;
1004    uint64_t nb10;
1005    uint64_t nb11;
1006    uint64_t nb12;
1007    uint64_t nb13;
1008    int32_t  ne0;
1009    int32_t  ne1;
1010    int32_t  ne2;
1011    int32_t  ne3;
1012    uint64_t nb0;
1013    uint64_t nb1;
1014    uint64_t nb2;
1015    uint64_t nb3;
1016} ggml_metal_kargs_solve_tri;
1017
1018typedef struct {
1019    int32_t  ne00t;
1020    int32_t  ne00;
1021    uint64_t nb01;
1022    uint64_t nb02;
1023    uint64_t nb03;
1024    int32_t  ne10;
1025    uint64_t nb10;
1026    uint64_t nb11;
1027    uint64_t nb12;
1028    uint64_t nb1;
1029    uint64_t nb2;
1030    uint64_t nb3;
1031} ggml_metal_kargs_get_rows;
1032
1033typedef struct {
1034    int32_t  nk0;
1035    int32_t  ne01;
1036    uint64_t nb01;
1037    uint64_t nb02;
1038    uint64_t nb03;
1039    int32_t  ne11;
1040    int32_t  ne12;
1041    uint64_t nb10;
1042    uint64_t nb11;
1043    uint64_t nb12;
1044    uint64_t nb1;
1045    uint64_t nb2;
1046    uint64_t nb3;
1047} ggml_metal_kargs_set_rows;
1048
1049typedef struct {
1050    int32_t  ne00;
1051    int32_t  ne01;
1052    int32_t  ne02;
1053    int32_t  ne03;
1054    uint64_t nb00;
1055    uint64_t nb01;
1056    uint64_t nb02;
1057    uint64_t nb03;
1058    int32_t  ne0;
1059    int32_t  ne1;
1060    int32_t  ne2;
1061    int32_t  ne3;
1062    uint64_t nb0;
1063    uint64_t nb1;
1064    uint64_t nb2;
1065    uint64_t nb3;
1066} ggml_metal_kargs_diag;
1067
1068typedef struct {
1069    int64_t  ne00;
1070    int64_t  ne01;
1071    int64_t  ne02;
1072    int64_t  ne03;
1073    uint64_t nb00;
1074    uint64_t nb01;
1075    uint64_t nb02;
1076    uint64_t nb03;
1077    int64_t  ne0;
1078    int64_t  ne1;
1079    int64_t  ne2;
1080    int64_t  ne3;
1081    uint64_t nb0;
1082    uint64_t nb1;
1083    uint64_t nb2;
1084    uint64_t nb3;
1085    float    sf0;
1086    float    sf1;
1087    float    sf2;
1088    float    sf3;
1089    float    poffs;
1090} ggml_metal_kargs_upscale;
1091
1092typedef struct {
1093    int64_t  ne00;
1094    int64_t  ne01;
1095    int64_t  ne02;
1096    int64_t  ne03;
1097    uint64_t nb00;
1098    uint64_t nb01;
1099    uint64_t nb02;
1100    uint64_t nb03;
1101    int64_t  ne0;
1102    int64_t  ne1;
1103    int64_t  ne2;
1104    int64_t  ne3;
1105    uint64_t nb0;
1106    uint64_t nb1;
1107    uint64_t nb2;
1108    uint64_t nb3;
1109} ggml_metal_kargs_pad;
1110
1111typedef struct {
1112    int64_t  ne00;
1113    int64_t  ne01;
1114    int64_t  ne02;
1115    int64_t  ne03;
1116    uint64_t nb00;
1117    uint64_t nb01;
1118    uint64_t nb02;
1119    uint64_t nb03;
1120    int64_t  ne0;
1121    int64_t  ne1;
1122    int64_t  ne2;
1123    int64_t  ne3;
1124    uint64_t nb0;
1125    uint64_t nb1;
1126    uint64_t nb2;
1127    uint64_t nb3;
1128    int32_t  p0;
1129    int32_t  p1;
1130} ggml_metal_kargs_pad_reflect_1d;
1131
1132typedef struct {
1133    int64_t  ne00;
1134    int64_t  ne01;
1135    int64_t  ne02;
1136    int64_t  ne03;
1137    uint64_t nb00;
1138    uint64_t nb01;
1139    uint64_t nb02;
1140    uint64_t nb03;
1141    int64_t  ne0;
1142    int64_t  ne1;
1143    int64_t  ne2;
1144    int64_t  ne3;
1145    uint64_t nb0;
1146    uint64_t nb1;
1147    uint64_t nb2;
1148    uint64_t nb3;
1149    int32_t  s0;
1150    int32_t  s1;
1151    int32_t  s2;
1152    int32_t  s3;
1153} ggml_metal_kargs_roll;
1154
1155typedef struct {
1156    uint64_t nb1;
1157    int      dim;
1158    int      max_period;
1159} ggml_metal_kargs_timestep_embedding;
1160
1161typedef struct {
1162    int32_t  ne00;
1163    int32_t  ne01;
1164    int32_t  ne02;
1165    int32_t  ne03;
1166    uint64_t nb00;
1167    uint64_t nb01;
1168    uint64_t nb02;
1169    uint64_t nb03;
1170    int32_t  ne0;
1171    int32_t  ne1;
1172    int32_t  ne2;
1173    int32_t  ne3;
1174    uint64_t nb0;
1175    uint64_t nb1;
1176    uint64_t nb2;
1177    uint64_t nb3;
1178} ggml_metal_kargs_tri;
1179
1180typedef struct {
1181    int32_t  ne00;
1182    int32_t  ne01;
1183    int32_t  ne02;
1184    int32_t  ne03;
1185    uint64_t nb00;
1186    uint64_t nb01;
1187    uint64_t nb02;
1188    uint64_t nb03;
1189    int32_t  ne0;
1190    int32_t  ne1;
1191    int32_t  ne2;
1192    int32_t  ne3;
1193    int32_t  top_k;
1194} ggml_metal_kargs_argsort;
1195
1196typedef struct {
1197    int64_t  ne00;
1198    int64_t  ne01;
1199    int64_t  ne02;
1200    int64_t  ne03;
1201    uint64_t nb00;
1202    uint64_t nb01;
1203    uint64_t nb02;
1204    uint64_t nb03;
1205    int32_t  ne0;
1206    int32_t  ne1;
1207    int32_t  ne2;
1208    int32_t  ne3;
1209    int32_t  top_k;
1210    int32_t  len;
1211} ggml_metal_kargs_argsort_merge;
1212
1213typedef struct {
1214    int32_t  ne00;   // number of columns (elements per row)
1215    int32_t  ne01;   // rows
1216    int32_t  ne02;
1217    int32_t  ne03;
1218    uint64_t nb01;   // row stride in src0
1219    uint64_t nb02;
1220    uint64_t nb03;
1221    int32_t  top_k;  // k
1222} ggml_metal_kargs_top_k;
1223
1224typedef struct {
1225    int32_t nrows;
1226} ggml_metal_kargs_fwht;
1227
1228typedef struct {
1229    int64_t  ne0;
1230    float    start;
1231    float    step;
1232} ggml_metal_kargs_arange;
1233
1234typedef struct {
1235    int64_t val;
1236} ggml_metal_kargs_memset;
1237
1238typedef struct {
1239    int32_t  n_kv;
1240    int32_t  n_batch;
1241    int32_t  mask_ne3;
1242    uint64_t nb1;
1243    uint64_t nb3;
1244    uint64_t nbq1;
1245    uint64_t nbq2;
1246    uint64_t nbq3;
1247    uint64_t nbk2;
1248    uint64_t nbk3;
1249    uint64_t nbw1;
1250    uint64_t nbw3;
1251    uint64_t nbm1;
1252    uint64_t nbm3;
1253} ggml_metal_kargs_lightning_indexer;
1254
1255typedef struct {
1256    int32_t  n_tokens;
1257    int32_t  n_iter;
1258    uint64_t nb_m0;
1259    uint64_t nb_m1;
1260    uint64_t nb_s0;
1261    uint64_t nb_b0;
1262    uint64_t nb_d0;
1263    uint64_t nb_d1;
1264    uint64_t nb_d2;
1265    float    eps;
1266} ggml_metal_kargs_dsv4_hc_comb;
1267
1268typedef struct {
1269    int32_t  n_embd;
1270    int32_t  n_tokens;
1271    uint64_t nb_x0;
1272    uint64_t nb_x1;
1273    uint64_t nb_x2;
1274    uint64_t nb_w0;
1275    uint64_t nb_w1;
1276    uint64_t nb_d0;
1277    uint64_t nb_d1;
1278} ggml_metal_kargs_dsv4_hc_pre;
1279
1280typedef struct {
1281    int32_t  n_embd;
1282    int32_t  n_tokens;
1283    uint64_t nb_x0;
1284    uint64_t nb_x1;
1285    uint64_t nb_r0;
1286    uint64_t nb_r1;
1287    uint64_t nb_r2;
1288    uint64_t nb_p0;
1289    uint64_t nb_p1;
1290    uint64_t nb_c0;
1291    uint64_t nb_c1;
1292    uint64_t nb_c2;
1293    uint64_t nb_d0;
1294    uint64_t nb_d1;
1295    uint64_t nb_d2;
1296} ggml_metal_kargs_dsv4_hc_post;
1297
1298typedef struct {
1299    int32_t  ne00;
1300    int32_t  ne01;
1301    int32_t  ne02;
1302    int32_t  ne03;
1303    uint64_t nb00;
1304    uint64_t nb01;
1305    uint64_t nb02;
1306    uint64_t nb03;
1307    uint64_t nb10;
1308    uint64_t nb11;
1309    uint64_t nb12;
1310    uint64_t nb13;
1311} ggml_metal_kargs_count_equal;
1312
1313typedef struct {
1314    int32_t  k0;
1315    int32_t  k1;
1316    int32_t  s0;
1317    int32_t  s1;
1318    int32_t  p0;
1319    int32_t  p1;
1320    int64_t  IH;
1321    int64_t  IW;
1322    int64_t  OH;
1323    int64_t  OW;
1324    int64_t  np;
1325} ggml_metal_kargs_pool_2d;
1326
1327typedef struct {
1328    int32_t  k0;
1329    int32_t  s0;
1330    int32_t  p0;
1331    int64_t  IW;
1332    int64_t  OW;
1333    int64_t  np;
1334} ggml_metal_kargs_pool_1d;
1335
1336typedef struct {
1337     int64_t ne00;
1338    uint64_t nb01;
1339} ggml_metal_kargs_argmax;
1340
1341typedef struct {
1342    int64_t  np;
1343} ggml_metal_kargs_opt_step_adamw;
1344
1345typedef struct {
1346    int64_t  np;
1347} ggml_metal_kargs_opt_step_sgd;
1348
1349typedef struct {
1350    int64_t ne;
1351} ggml_metal_kargs_silu_back;
1352
1353#endif // GGML_METAL_IMPL