char-fork/llama.cpp

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

BartowskiAVX2: Speed up large batch size prompt processing of IQ models (#27402)85c55223c

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
47.2 KiB1253 linesraw
1#define GGML_COMMON_IMPL_CPP
2#define GGML_COMMON_DECL_CPP
3#include "ggml-common.h"
4
5#include "ggml-impl.h"
6#include "ggml-cpu.h"
7#include "ggml-cpu-impl.h"
8#include "simd-mappings.h"
9#include "traits.h"
10
11#include <cassert>
12#include <cstdlib>
13#include <cstring>
14
15#include "iqp.h"
16
17#define UNUSED GGML_UNUSED
18
19// smallest src1 batch for which the decode pays for itself
20#define GGML_IQP_MIN_BATCH 8
21
22// same, per expert, for MUL_MAT_ID
23#define GGML_IQP_MIN_BATCH_ID 8
24
25bool ggml_cpu_iqp_mul_mat_id_min_batch(int64_t cne1) {
26    return cne1 >= GGML_IQP_MIN_BATCH_ID;
27}
28
29// src0 rows interleaved per panel
30#define IQP_NB_ROWS 8
31
32#define IQP_SB_SIZE 16                    // weights per sub-block
33#define IQP_NSB     (QK_K / IQP_SB_SIZE)  // sub-blocks per super-block
34
35// one super-block of a grid based IQ type decoded to int8, 8 rows interleaved:
36// dfac[row] * iscales[sb*8 + row] * qs is bit identical to dequantize_row_iq*
37struct block_iqp_x8 {
38    float   dfac[8];               // f32 super-block scale, d * 2^-k
39    int32_t bias[8];               // 128 * sum(qs * iscale), see GGML_IQP_USE_BIAS
40    int8_t  iscales[IQP_NSB * 8];  // integer sub-block scales, in [-32, 31]
41    int8_t  qs[QK_K * 8];          // qs[sb*128 + g*32 + row*4 + k] = column sb*16 + g*4 + k
42};
43
44static_assert(sizeof(block_iqp_x8) == 8 * sizeof(float) + 8 * sizeof(int32_t) + IQP_NSB * 8 + QK_K * 8,
45              "wrong iqp_x8 block size/padding");
46
47// feed the activations to VNNI as unsigned bytes (y + 128) and correct with bias[]; without VNNI the kernels use the maddubs sign trick instead and bias[] is not filled
48#if defined(__AVX2__) && ((defined(__AVX512VNNI__) && defined(__AVX512VL__)) || defined(__AVXVNNI__))
49#    define GGML_IQP_USE_BIAS 1
50#else
51#    define GGML_IQP_USE_BIAS 0
52#endif
53
54static inline size_t ggml_cpu_iqp_row_size(const struct ggml_tensor * dst) {
55    return ggml_row_size(GGML_TYPE_Q8_K, dst->src[1]->ne[0]);
56}
57
58// the low 7 bits of v are the first 7 signs and the 8th is their parity (cf. unpack_ksigns in the CUDA backend)
59static inline uint8_t iqp_unpack_ksigns(uint32_t v) {
60    uint32_t p = v ^ (v >> 4);
61
62    p ^= p >> 2;
63    p ^= p >> 1;
64
65    return (uint8_t) (v ^ ((p & 1) << 7));
66}
67
68#if defined(__AVX2__)
69
70// 0xFF in every byte whose sign bit is set; sv holds each sign byte broadcast over the 8 bytes it governs
71static inline __m256i iqp_sign_mask(__m256i sv) {
72    const __m256i sel = _mm256_set1_epi64x((int64_t) 0x8040201008040201ULL);
73
74#    if defined(__GFNI__)
75    // computes the and + compare in one instruction
76    return _mm256_gf2p8affine_epi64_epi8(sel, sv, 0);
77#    else
78    return _mm256_cmpeq_epi8(_mm256_and_si256(sv, sel), sel);
79#    endif
80}
81
82// signs holds four sign bytes, byte l governing values 8*l .. 8*l+7 - spread each over its 8 lanes
83static inline __m256i iqp_sign_bytes(uint32_t signs) {
84    const __m256i bcast = _mm256_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,  //
85                                           2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3);
86
87    return _mm256_shuffle_epi8(_mm256_set1_epi32((int32_t) signs), bcast);
88}
89
90// x ^ m - m negates the lanes where m is 0xFF
91static inline __m256i iqp_apply_signs(__m256i x, __m256i m) {
92    return _mm256_sub_epi8(_mm256_xor_si256(x, m), m);
93}
94
95#endif
96
97// 32 values from four 8 byte grid entries, sign byte l of signs applied to group l
98static inline void iqp_store_signed_x8(int8_t * GGML_RESTRICT dst,
99                                       uint64_t               g0,
100                                       uint64_t               g1,
101                                       uint64_t               g2,
102                                       uint64_t               g3,
103                                       uint32_t               signs) {
104#if defined(__AVX2__)
105    const __m256i g = _mm256_set_epi64x((int64_t) g3, (int64_t) g2, (int64_t) g1, (int64_t) g0);
106    const __m256i m = iqp_sign_mask(iqp_sign_bytes(signs));
107
108    _mm256_storeu_si256((__m256i *) dst, iqp_apply_signs(g, m));
109#else
110    const uint64_t g[4] = { g0, g1, g2, g3 };
111
112    for (int l = 0; l < 4; ++l) {
113        const uint8_t * grid = (const uint8_t *) &g[l];
114        const uint8_t   s    = (uint8_t) (signs >> 8 * l);
115
116        for (int j = 0; j < 8; ++j) {
117            dst[8 * l + j] = s & kmask_iq2xs[j] ? -grid[j] : grid[j];
118        }
119    }
120#endif
121}
122
123// same, but the eight values of group l come from two 4 byte grid entries
124static inline void iqp_store_signed_x4(int8_t * GGML_RESTRICT dst,
125                                       uint32_t               g0a,
126                                       uint32_t               g0b,
127                                       uint32_t               g1a,
128                                       uint32_t               g1b,
129                                       uint32_t               g2a,
130                                       uint32_t               g2b,
131                                       uint32_t               g3a,
132                                       uint32_t               g3b,
133                                       uint32_t               signs) {
134#if defined(__AVX2__)
135    const __m256i g = _mm256_setr_epi32((int32_t) g0a, (int32_t) g0b, (int32_t) g1a, (int32_t) g1b, (int32_t) g2a,
136                                        (int32_t) g2b, (int32_t) g3a, (int32_t) g3b);
137    const __m256i m = iqp_sign_mask(iqp_sign_bytes(signs));
138
139    _mm256_storeu_si256((__m256i *) dst, iqp_apply_signs(g, m));
140#else
141    const uint32_t ga[4] = { g0a, g1a, g2a, g3a };
142    const uint32_t gb[4] = { g0b, g1b, g2b, g3b };
143
144    for (int l = 0; l < 4; ++l) {
145        const uint8_t * grid1 = (const uint8_t *) &ga[l];
146        const uint8_t * grid2 = (const uint8_t *) &gb[l];
147        const uint8_t   s     = (uint8_t) (signs >> 8 * l);
148
149        for (int j = 0; j < 4; ++j) {
150            dst[8 * l + j + 0] = s & kmask_iq2xs[j + 0] ? -grid1[j] : grid1[j];
151            dst[8 * l + j + 4] = s & kmask_iq2xs[j + 4] ? -grid2[j] : grid2[j];
152        }
153    }
154#endif
155}
156
157// 32 values of 8 * grid + delta from four 8 byte grid entries (grid bytes are in {-1, 0, 1}), byte l of deltas applying to group l
158static inline void iqp_store_iq1_x8(int8_t * GGML_RESTRICT dst,
159                                    uint64_t               g0,
160                                    uint64_t               g1,
161                                    uint64_t               g2,
162                                    uint64_t               g3,
163                                    uint32_t               deltas) {
164#if defined(__AVX2__)
165    __m256i g = _mm256_set_epi64x((int64_t) g3, (int64_t) g2, (int64_t) g1, (int64_t) g0);
166
167    // no byte shift in AVX2
168    g = _mm256_add_epi8(g, g);
169    g = _mm256_add_epi8(g, g);
170    g = _mm256_add_epi8(g, g);
171
172    _mm256_storeu_si256((__m256i *) dst, _mm256_add_epi8(g, iqp_sign_bytes(deltas)));
173#else
174    const uint64_t g[4] = { g0, g1, g2, g3 };
175
176    for (int l = 0; l < 4; ++l) {
177        const int8_t * grid  = (const int8_t *) &g[l];
178        const int8_t   delta = (int8_t) (deltas >> 8 * l);
179
180        for (int j = 0; j < 8; ++j) {
181            dst[8 * l + j] = 8 * grid[j] + delta;
182        }
183    }
184#endif
185}
186
187// 32 values from 16 packed nibbles through the kvalues_iq4nl lookup: low nibbles first, then high
188static inline void iqp_store_iq4_x32(int8_t * GGML_RESTRICT dst, const uint8_t * GGML_RESTRICT qs) {
189#if defined(__AVX2__)
190    const __m128i q   = _mm_loadu_si128((const __m128i *) qs);
191    const __m128i lut = _mm_loadu_si128((const __m128i *) kvalues_iq4nl);
192    const __m128i m4  = _mm_set1_epi8(0xf);
193
194    _mm_storeu_si128((__m128i *) (dst + 0), _mm_shuffle_epi8(lut, _mm_and_si128(q, m4)));
195    _mm_storeu_si128((__m128i *) (dst + 16), _mm_shuffle_epi8(lut, _mm_and_si128(_mm_srli_epi16(q, 4), m4)));
196#else
197    for (int j = 0; j < 16; ++j) {
198        dst[j + 0]  = kvalues_iq4nl[qs[j] & 0xf];
199        dst[j + 16] = kvalues_iq4nl[qs[j] >> 4];
200    }
201#endif
202}
203
204#if GGML_IQP_USE_BIAS
205
206// sum of qs * iscale over one super-block, at most 256 * 127 * 32 = 1.04e6
207static inline int32_t iqp_weighted_sum(const int8_t * GGML_RESTRICT vals, const int8_t * GGML_RESTRICT iscales) {
208#if defined(__AVX2__)
209    static_assert(IQP_SB_SIZE == 16, "the vector path folds two sub-blocks per 32 byte load");
210
211    const __m256i ones8  = _mm256_set1_epi8(1);
212    const __m256i ones16 = _mm256_set1_epi16(1);
213
214    __m256i acc = _mm256_setzero_si256();
215
216    for (int i = 0; i < QK_K / 32; ++i) {
217        // sum groups of 4 bytes into int32, the low four lanes cover sub-block 2*i and the high four 2*i + 1
218        const __m256i v = _mm256_loadu_si256((const __m256i *) (vals + 32 * i));
219        const __m256i p = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, v), ones16);
220
221        const __m256i s = _mm256_set_m128i(_mm_set1_epi32(iscales[2 * i + 1]), _mm_set1_epi32(iscales[2 * i + 0]));
222
223        acc = _mm256_add_epi32(acc, _mm256_mullo_epi32(p, s));
224    }
225
226    __m128i sum = _mm_add_epi32(_mm256_castsi256_si128(acc), _mm256_extracti128_si256(acc, 1));
227
228    sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, _MM_SHUFFLE(1, 0, 3, 2)));
229    sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, _MM_SHUFFLE(2, 3, 0, 1)));
230
231    return _mm_cvtsi128_si32(sum);
232#else
233    int32_t wsum = 0;
234
235    for (int sb = 0; sb < IQP_NSB; ++sb) {
236        int32_t vsum = 0;
237
238        for (int k = 0; k < IQP_SB_SIZE; ++k) {
239            vsum += vals[sb * IQP_SB_SIZE + k];
240        }
241
242        wsum += iscales[sb] * vsum;
243    }
244
245    return wsum;
246#endif
247}
248
249#endif  // GGML_IQP_USE_BIAS
250
251static void iqp_decode_iq2_xxs(const void * GGML_RESTRICT vx,
252                               int8_t * GGML_RESTRICT     vals,
253                               int8_t * GGML_RESTRICT     iscales,
254                               float * GGML_RESTRICT      dfac) {
255    const block_iq2_xxs * x = (const block_iq2_xxs *) vx;
256
257    // db = d * (0.5 + ls) * 0.25 = (d / 8) * (2 * ls + 1), ls 4 bit
258    *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f;
259
260    uint32_t        aux32[2];
261    const uint8_t * aux8 = (const uint8_t *) aux32;
262
263    for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) {
264        memcpy(aux32, x->qs + 4 * ib32, 2 * sizeof(uint32_t));
265        const int8_t ls = (int8_t) (2 * (aux32[1] >> 28) + 1);
266
267        iscales[2 * ib32 + 0] = ls;
268        iscales[2 * ib32 + 1] = ls;
269
270        const uint32_t signs = (uint32_t) iqp_unpack_ksigns((aux32[1] >> 0) & 127) |
271                               (uint32_t) iqp_unpack_ksigns((aux32[1] >> 7) & 127) << 8 |
272                               (uint32_t) iqp_unpack_ksigns((aux32[1] >> 14) & 127) << 16 |
273                               (uint32_t) iqp_unpack_ksigns((aux32[1] >> 21) & 127) << 24;
274
275        iqp_store_signed_x8(vals + 32 * ib32, iq2xxs_grid[aux8[0]], iq2xxs_grid[aux8[1]], iq2xxs_grid[aux8[2]],
276                            iq2xxs_grid[aux8[3]], signs);
277    }
278}
279
280static void iqp_decode_iq2_xs(const void * GGML_RESTRICT vx,
281                              int8_t * GGML_RESTRICT     vals,
282                              int8_t * GGML_RESTRICT     iscales,
283                              float * GGML_RESTRICT      dfac) {
284    const block_iq2_xs * x = (const block_iq2_xs *) vx;
285
286    *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f;
287
288    for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) {
289        iscales[2 * ib32 + 0] = (int8_t) (2 * (x->scales[ib32] & 0xf) + 1);
290        iscales[2 * ib32 + 1] = (int8_t) (2 * (x->scales[ib32] >> 4) + 1);
291
292        const uint16_t * q = x->qs + 4 * ib32;
293
294        const uint32_t signs = (uint32_t) iqp_unpack_ksigns(q[0] >> 9) | (uint32_t) iqp_unpack_ksigns(q[1] >> 9) << 8 |
295                               (uint32_t) iqp_unpack_ksigns(q[2] >> 9) << 16 |
296                               (uint32_t) iqp_unpack_ksigns(q[3] >> 9) << 24;
297
298        iqp_store_signed_x8(vals + 32 * ib32, iq2xs_grid[q[0] & 511], iq2xs_grid[q[1] & 511], iq2xs_grid[q[2] & 511],
299                            iq2xs_grid[q[3] & 511], signs);
300    }
301}
302
303static void iqp_decode_iq2_s(const void * GGML_RESTRICT vx,
304                             int8_t * GGML_RESTRICT     vals,
305                             int8_t * GGML_RESTRICT     iscales,
306                             float * GGML_RESTRICT      dfac) {
307    const block_iq2_s * x = (const block_iq2_s *) vx;
308
309    const uint8_t * qs    = x->qs;
310    const uint8_t * qh    = x->qh;
311    const uint8_t * signs = qs + QK_K / 8;
312
313    *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f;
314
315    for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) {
316        iscales[2 * ib32 + 0] = (int8_t) (2 * (x->scales[ib32] & 0xf) + 1);
317        iscales[2 * ib32 + 1] = (int8_t) (2 * (x->scales[ib32] >> 4) + 1);
318
319        const uint32_t sbits =
320            (uint32_t) signs[0] | (uint32_t) signs[1] << 8 | (uint32_t) signs[2] << 16 | (uint32_t) signs[3] << 24;
321
322        iqp_store_signed_x8(vals + 32 * ib32, iq2s_grid[qs[0] | (qh[ib32] << 8 & 0x300)],
323                            iq2s_grid[qs[1] | (qh[ib32] << 6 & 0x300)], iq2s_grid[qs[2] | (qh[ib32] << 4 & 0x300)],
324                            iq2s_grid[qs[3] | (qh[ib32] << 2 & 0x300)], sbits);
325        qs += 4;
326        signs += 4;
327    }
328}
329
330static void iqp_decode_iq3_xxs(const void * GGML_RESTRICT vx,
331                               int8_t * GGML_RESTRICT     vals,
332                               int8_t * GGML_RESTRICT     iscales,
333                               float * GGML_RESTRICT      dfac) {
334    const block_iq3_xxs * x = (const block_iq3_xxs *) vx;
335
336    const uint8_t * qs               = x->qs;
337    const uint8_t * scales_and_signs = qs + QK_K / 4;
338
339    // db = d * (0.5 + ls) * 0.5 = (d / 4) * (2 * ls + 1), ls 4 bit
340    *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.25f;
341
342    uint32_t aux32;
343
344    for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) {
345        memcpy(&aux32, scales_and_signs + 4 * ib32, sizeof(uint32_t));
346        const int8_t ls = (int8_t) (2 * (aux32 >> 28) + 1);
347
348        iscales[2 * ib32 + 0] = ls;
349        iscales[2 * ib32 + 1] = ls;
350
351        const uint32_t signs = (uint32_t) iqp_unpack_ksigns((aux32 >> 0) & 127) |
352                               (uint32_t) iqp_unpack_ksigns((aux32 >> 7) & 127) << 8 |
353                               (uint32_t) iqp_unpack_ksigns((aux32 >> 14) & 127) << 16 |
354                               (uint32_t) iqp_unpack_ksigns((aux32 >> 21) & 127) << 24;
355
356        iqp_store_signed_x4(vals + 32 * ib32, iq3xxs_grid[qs[0]], iq3xxs_grid[qs[1]], iq3xxs_grid[qs[2]],
357                            iq3xxs_grid[qs[3]], iq3xxs_grid[qs[4]], iq3xxs_grid[qs[5]], iq3xxs_grid[qs[6]],
358                            iq3xxs_grid[qs[7]], signs);
359        qs += 8;
360    }
361}
362
363static void iqp_decode_iq3_s(const void * GGML_RESTRICT vx,
364                             int8_t * GGML_RESTRICT     vals,
365                             int8_t * GGML_RESTRICT     iscales,
366                             float * GGML_RESTRICT      dfac) {
367    const block_iq3_s * x = (const block_iq3_s *) vx;
368
369    const uint8_t * qs    = x->qs;
370    const uint8_t * qh    = x->qh;
371    const uint8_t * signs = x->signs;
372
373    // db = d * (1 + 2 * ls), ls 4 bit
374    *dfac = GGML_CPU_FP16_TO_FP32(x->d);
375
376    int k = 0;
377
378    for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) {
379        const int8_t db1 = (int8_t) (1 + 2 * (x->scales[ib32 / 2] & 0xf));
380        const int8_t db2 = (int8_t) (1 + 2 * (x->scales[ib32 / 2] >> 4));
381
382        iscales[2 * ib32 + 0] = db1;
383        iscales[2 * ib32 + 1] = db1;
384        iscales[2 * ib32 + 2] = db2;
385        iscales[2 * ib32 + 3] = db2;
386
387        for (int h = 0; h < 2; ++h) {
388            const uint32_t sbits =
389                (uint32_t) signs[0] | (uint32_t) signs[1] << 8 | (uint32_t) signs[2] << 16 | (uint32_t) signs[3] << 24;
390
391            iqp_store_signed_x4(vals + k, iq3s_grid[qs[0] | ((qh[h] << 8) & 256)],
392                                iq3s_grid[qs[1] | ((qh[h] << 7) & 256)], iq3s_grid[qs[2] | ((qh[h] << 6) & 256)],
393                                iq3s_grid[qs[3] | ((qh[h] << 5) & 256)], iq3s_grid[qs[4] | ((qh[h] << 4) & 256)],
394                                iq3s_grid[qs[5] | ((qh[h] << 3) & 256)], iq3s_grid[qs[6] | ((qh[h] << 2) & 256)],
395                                iq3s_grid[qs[7] | ((qh[h] << 1) & 256)], sbits);
396
397            k += 32;
398            qs += 8;
399            signs += 4;
400        }
401        qh += 2;
402    }
403}
404
405// dequantize_row_iq1_* computes y = dl * (grid[j] + delta) with delta = +-1/8, so the panel stores 8 * grid[j] +- 1 and folds the /8 into dfac
406static void iqp_decode_iq1_s(const void * GGML_RESTRICT vx,
407                             int8_t * GGML_RESTRICT     vals,
408                             int8_t * GGML_RESTRICT     iscales,
409                             float * GGML_RESTRICT      dfac) {
410    const block_iq1_s * x = (const block_iq1_s *) vx;
411
412    const uint8_t *  qs = x->qs;
413    const uint16_t * qh = x->qh;
414
415    // dl = d * (2 * ls + 1) * 0.125, ls 3 bit
416    *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f;
417
418    for (int ib = 0; ib < QK_K / 32; ++ib) {
419        const int8_t dl    = (int8_t) (2 * ((qh[ib] >> 12) & 7) + 1);
420        const int8_t delta = qh[ib] & 0x8000 ? -1 : 1;
421
422        iscales[2 * ib + 0] = dl;
423        iscales[2 * ib + 1] = dl;
424
425        iqp_store_iq1_x8(vals + 32 * ib, iq1s_grid[qs[0] | (((qh[ib] >> 0) & 7) << 8)],
426                         iq1s_grid[qs[1] | (((qh[ib] >> 3) & 7) << 8)], iq1s_grid[qs[2] | (((qh[ib] >> 6) & 7) << 8)],
427                         iq1s_grid[qs[3] | (((qh[ib] >> 9) & 7) << 8)], ((uint8_t) delta) * 0x01010101u);
428        qs += 4;
429    }
430}
431
432static void iqp_decode_iq1_m(const void * GGML_RESTRICT vx,
433                             int8_t * GGML_RESTRICT     vals,
434                             int8_t * GGML_RESTRICT     iscales,
435                             float * GGML_RESTRICT      dfac) {
436    const block_iq1_m * x = (const block_iq1_m *) vx;
437
438    // block_iq1_m has no d field - the fp16 super-block scale is spread over the top nibbles of the four scale words
439    const uint16_t * sc = (const uint16_t *) x->scales;
440
441    iq1m_scale_t scale;
442    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
443
444    *dfac = GGML_CPU_FP16_TO_FP32(scale.f16) * 0.125f;
445
446    const uint8_t * qs = x->qs;
447    const uint8_t * qh = x->qh;
448
449    for (int ib = 0; ib < QK_K / 32; ++ib) {
450        iscales[2 * ib + 0] = (int8_t) (2 * ((sc[ib / 2] >> (6 * (ib % 2) + 0)) & 0x7) + 1);
451        iscales[2 * ib + 1] = (int8_t) (2 * ((sc[ib / 2] >> (6 * (ib % 2) + 3)) & 0x7) + 1);
452
453        const uint16_t idx[4] = {
454            (uint16_t) (qs[0] | ((qh[0] << 8) & 0x700)),
455            (uint16_t) (qs[1] | ((qh[0] << 4) & 0x700)),
456            (uint16_t) (qs[2] | ((qh[1] << 8) & 0x700)),
457            (uint16_t) (qs[3] | ((qh[1] << 4) & 0x700)),
458        };
459        const uint32_t deltas = (uint32_t) (qh[0] & 0x08 ? 0xff : 0x01) | (uint32_t) (qh[0] & 0x80 ? 0xff : 0x01) << 8 |
460                                (uint32_t) (qh[1] & 0x08 ? 0xff : 0x01) << 16 |
461                                (uint32_t) (qh[1] & 0x80 ? 0xff : 0x01) << 24;
462
463        iqp_store_iq1_x8(vals + 32 * ib, iq1s_grid[idx[0]], iq1s_grid[idx[1]], iq1s_grid[idx[2]], iq1s_grid[idx[3]],
464                         deltas);
465        qs += 4;
466        qh += 2;
467    }
468}
469
470static void iqp_decode_iq4_xs(const void * GGML_RESTRICT vx,
471                              int8_t * GGML_RESTRICT     vals,
472                              int8_t * GGML_RESTRICT     iscales,
473                              float * GGML_RESTRICT      dfac) {
474    const block_iq4_xs * x = (const block_iq4_xs *) vx;
475
476    const uint8_t * qs = x->qs;
477
478    // dl = d * (ls - 32), ls 6 bit, so the integer scale is in [-32, 31]
479    *dfac = GGML_CPU_FP16_TO_FP32(x->d);
480
481    for (int ib = 0; ib < QK_K / 32; ++ib) {
482        const int    ls = ((x->scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((x->scales_h >> 2 * ib) & 3) << 4);
483        const int8_t dl = (int8_t) (ls - 32);
484
485        iscales[2 * ib + 0] = dl;
486        iscales[2 * ib + 1] = dl;
487
488        iqp_store_iq4_x32(vals + 32 * ib, qs);
489        qs += 16;
490    }
491}
492
493// expanded by the eligibility test and the decode dispatch
494#define IQP_TYPE_LIST(T) \
495    T(IQ2_XXS, iq2_xxs)  \
496    T(IQ2_XS, iq2_xs)    \
497    T(IQ2_S, iq2_s)      \
498    T(IQ3_XXS, iq3_xxs)  \
499    T(IQ3_S, iq3_s)      \
500    T(IQ1_S, iq1_s)      \
501    T(IQ1_M, iq1_m)      \
502    T(IQ4_XS, iq4_xs)
503
504static bool iqp_decode_superblock(enum ggml_type             type,
505                                  const void * GGML_RESTRICT vx,
506                                  int8_t * GGML_RESTRICT     vals,
507                                  int8_t * GGML_RESTRICT     iscales,
508                                  float * GGML_RESTRICT      dfac) {
509    switch (type) {
510#define IQP_CASE(E, name)                           \
511    case GGML_TYPE_##E:                             \
512        iqp_decode_##name(vx, vals, iscales, dfac); \
513        return true;
514        IQP_TYPE_LIST(IQP_CASE)
515#undef IQP_CASE
516        default:
517            return false;
518    }
519}
520
521#if defined(__AVX2__)
522
523// 8x8 int32 transpose of the 32 column group starting at column off
524static inline void iqp_interleave_x8(int8_t * GGML_RESTRICT dst, const int8_t (*vals)[QK_K], int off) {
525    static_assert(IQP_NB_ROWS == 8, "the transpose is 8x8");
526
527    __m256i v[IQP_NB_ROWS];
528
529    for (int r = 0; r < IQP_NB_ROWS; ++r) {
530        v[r] = _mm256_loadu_si256((const __m256i *) (vals[r] + off));
531    }
532
533    // pair rows into dword couples, then into qword quadruples, then swap the 128 bit lanes
534    const __m256i a0 = _mm256_unpacklo_epi32(v[0], v[1]);
535    const __m256i a1 = _mm256_unpackhi_epi32(v[0], v[1]);
536    const __m256i a2 = _mm256_unpacklo_epi32(v[2], v[3]);
537    const __m256i a3 = _mm256_unpackhi_epi32(v[2], v[3]);
538    const __m256i a4 = _mm256_unpacklo_epi32(v[4], v[5]);
539    const __m256i a5 = _mm256_unpackhi_epi32(v[4], v[5]);
540    const __m256i a6 = _mm256_unpacklo_epi32(v[6], v[7]);
541    const __m256i a7 = _mm256_unpackhi_epi32(v[6], v[7]);
542
543    const __m256i b0 = _mm256_unpacklo_epi64(a0, a2);
544    const __m256i b1 = _mm256_unpackhi_epi64(a0, a2);
545    const __m256i b2 = _mm256_unpacklo_epi64(a1, a3);
546    const __m256i b3 = _mm256_unpackhi_epi64(a1, a3);
547    const __m256i b4 = _mm256_unpacklo_epi64(a4, a6);
548    const __m256i b5 = _mm256_unpackhi_epi64(a4, a6);
549    const __m256i b6 = _mm256_unpacklo_epi64(a5, a7);
550    const __m256i b7 = _mm256_unpackhi_epi64(a5, a7);
551
552    _mm256_storeu_si256((__m256i *) (dst + 0 * 32), _mm256_permute2x128_si256(b0, b4, 0x20));
553    _mm256_storeu_si256((__m256i *) (dst + 1 * 32), _mm256_permute2x128_si256(b1, b5, 0x20));
554    _mm256_storeu_si256((__m256i *) (dst + 2 * 32), _mm256_permute2x128_si256(b2, b6, 0x20));
555    _mm256_storeu_si256((__m256i *) (dst + 3 * 32), _mm256_permute2x128_si256(b3, b7, 0x20));
556    _mm256_storeu_si256((__m256i *) (dst + 4 * 32), _mm256_permute2x128_si256(b0, b4, 0x31));
557    _mm256_storeu_si256((__m256i *) (dst + 5 * 32), _mm256_permute2x128_si256(b1, b5, 0x31));
558    _mm256_storeu_si256((__m256i *) (dst + 6 * 32), _mm256_permute2x128_si256(b2, b6, 0x31));
559    _mm256_storeu_si256((__m256i *) (dst + 7 * 32), _mm256_permute2x128_si256(b3, b7, 0x31));
560}
561
562#endif
563
564// decode IQP_NB_ROWS consecutive source rows (starting at src, row stride nb01) into a panel of nblocks block_iqp_x8
565static void iqp_decode_panel_8(enum ggml_type               type,
566                               const char * GGML_RESTRICT   src,
567                               size_t                       nb01,
568                               int64_t                      nblocks,
569                               block_iqp_x8 * GGML_RESTRICT dst) {
570    const size_t bsize = ggml_type_size(type);
571
572    int8_t vals[IQP_NB_ROWS][QK_K];
573    int8_t iscales[IQP_NB_ROWS][IQP_NSB];
574    float  dfac[IQP_NB_ROWS];
575
576    for (int64_t x = 0; x < nblocks; x++) {
577        for (int r = 0; r < IQP_NB_ROWS; r++) {
578            const char * blk = src + r * nb01 + x * bsize;
579
580            const bool ok = iqp_decode_superblock(type, blk, vals[r], iscales[r], &dfac[r]);
581            GGML_ASSERT(ok);
582
583#ifdef GGML_IQP_VERIFY
584            // check that the panel reproduces the reference dequantization bit exactly
585            float ref[QK_K];
586            ggml_get_type_traits(type)->to_float(blk, ref, QK_K);
587            for (int j = 0; j < QK_K; j++) {
588                const float scale = dfac[r] * iscales[r][j / IQP_SB_SIZE];
589                GGML_ASSERT(scale * vals[r][j] == ref[j]);
590            }
591#endif
592        }
593
594        for (int r = 0; r < IQP_NB_ROWS; r++) {
595            dst->dfac[r] = dfac[r];
596
597            for (int sb = 0; sb < IQP_NSB; sb++) {
598                dst->iscales[sb * IQP_NB_ROWS + r] = iscales[r][sb];
599            }
600
601#if GGML_IQP_USE_BIAS
602            dst->bias[r] = 128 * iqp_weighted_sum(vals[r], iscales[r]);
603#endif
604        }
605
606#if defined(__AVX2__)
607        for (int grp = 0; grp < QK_K / 32; grp++) {
608            iqp_interleave_x8(dst->qs + grp * 256, vals, grp * 32);
609        }
610#else
611        for (int r = 0; r < IQP_NB_ROWS; r++) {
612            for (int sb = 0; sb < IQP_NSB; sb++) {
613                for (int g = 0; g < IQP_SB_SIZE / 4; g++) {
614                    memcpy(dst->qs + sb * 128 + g * 32 + r * 4, vals[r] + sb * IQP_SB_SIZE + g * 4, 4);
615                }
616            }
617        }
618#endif
619
620        dst++;
621    }
622}
623
624// gemm/gemv kernels: vx points at block_iqp_x8, vy at plain (non interleaved) block_q8_K rows
625
626static void iqp_gemv_8x8_q8_K_generic(int                        n,
627                                      float * GGML_RESTRICT      s,
628                                      size_t                     bs,
629                                      const void * GGML_RESTRICT vx,
630                                      const void * GGML_RESTRICT vy,
631                                      int                        nr,
632                                      int                        nc) {
633    const int nb                = n / QK_K;
634    const int ncols_interleaved = 8;
635
636    assert(n % QK_K == 0);
637    assert(nc % ncols_interleaved == 0);
638
639    UNUSED(bs);
640    UNUSED(nr);
641
642    const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx;
643    const block_q8_K *   a_ptr       = (const block_q8_K *) vy;
644
645    for (int x = 0; x < nc / ncols_interleaved; x++) {
646        const block_iqp_x8 * b_ptr = b_ptr_start + x * nb;
647
648        float sumf[8] = { 0 };
649
650        for (int l = 0; l < nb; l++) {
651            int32_t sumi[8] = { 0 };
652
653            for (int sb = 0; sb < IQP_NSB; sb++) {
654                int32_t isum[8] = { 0 };
655
656                for (int g = 0; g < 4; g++) {
657                    for (int j = 0; j < ncols_interleaved; j++) {
658                        for (int k = 0; k < 4; k++) {
659                            isum[j] += b_ptr[l].qs[sb * 128 + g * 32 + j * 4 + k] * a_ptr[l].qs[sb * 16 + g * 4 + k];
660                        }
661                    }
662                }
663
664                for (int j = 0; j < ncols_interleaved; j++) {
665                    sumi[j] += isum[j] * b_ptr[l].iscales[sb * 8 + j];
666                }
667            }
668
669            for (int j = 0; j < ncols_interleaved; j++) {
670                sumf[j] += (float) sumi[j] * (b_ptr[l].dfac[j] * a_ptr[l].d);
671            }
672        }
673
674        for (int j = 0; j < ncols_interleaved; j++) {
675            s[x * ncols_interleaved + j] = sumf[j];
676        }
677    }
678}
679
680// one 4 row x nc column tile; s points at the first of the four output rows, bs floats apart
681static void iqp_gemm_tile_4_generic(int                                nb,
682                                    float * GGML_RESTRICT              s,
683                                    size_t                             bs,
684                                    const block_iqp_x8 * GGML_RESTRICT b_ptr_start,
685                                    const block_q8_K * const           a_ptr[4],
686                                    int                                nc) {
687    const int ncols_interleaved = 8;
688
689    for (int x = 0; x < nc / ncols_interleaved; x++) {
690        const block_iqp_x8 * b_ptr = b_ptr_start + x * nb;
691
692        float sumf[4][8];
693        for (int m = 0; m < 4; m++) {
694            for (int j = 0; j < ncols_interleaved; j++) {
695                sumf[m][j] = 0.0f;
696            }
697        }
698
699        for (int l = 0; l < nb; l++) {
700            for (int m = 0; m < 4; m++) {
701                int32_t sumi[8] = { 0 };
702
703                for (int sb = 0; sb < IQP_NSB; sb++) {
704                    int32_t isum[8] = { 0 };
705
706                    for (int g = 0; g < 4; g++) {
707                        for (int j = 0; j < ncols_interleaved; j++) {
708                            for (int k = 0; k < 4; k++) {
709                                isum[j] +=
710                                    b_ptr[l].qs[sb * 128 + g * 32 + j * 4 + k] * a_ptr[m][l].qs[sb * 16 + g * 4 + k];
711                            }
712                        }
713                    }
714
715                    for (int j = 0; j < ncols_interleaved; j++) {
716                        sumi[j] += isum[j] * b_ptr[l].iscales[sb * 8 + j];
717                    }
718                }
719
720                for (int j = 0; j < ncols_interleaved; j++) {
721                    sumf[m][j] += (float) sumi[j] * (b_ptr[l].dfac[j] * a_ptr[m][l].d);
722                }
723            }
724        }
725
726        for (int m = 0; m < 4; m++) {
727            for (int j = 0; j < ncols_interleaved; j++) {
728                s[m * bs + x * ncols_interleaved + j] = sumf[m][j];
729            }
730        }
731    }
732}
733
734static void iqp_gemm_8x8_q8_K_generic(int                        n,
735                                      float * GGML_RESTRICT      s,
736                                      size_t                     bs,
737                                      const void * GGML_RESTRICT vx,
738                                      const void * GGML_RESTRICT vy,
739                                      int                        nr,
740                                      int                        nc) {
741    const int nb = n / QK_K;
742
743    assert(n % QK_K == 0);
744    assert(nr % 4 == 0);
745    assert(nc % 8 == 0);
746
747    const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx;
748    const block_q8_K *   a_ptr_start = (const block_q8_K *) vy;
749
750    for (int y = 0; y < nr / 4; y++) {
751        const block_q8_K * a_ptr[4];
752        for (int m = 0; m < 4; m++) {
753            a_ptr[m] = a_ptr_start + (y * 4 + m) * nb;
754        }
755
756        iqp_gemm_tile_4_generic(nb, s + y * 4 * bs, bs, b_ptr_start, a_ptr, nc);
757    }
758}
759
760static void iqp_gemm_8x8_q8_K_p4_generic(int                                n,
761                                         float * GGML_RESTRICT              s,
762                                         size_t                             bs,
763                                         const void * GGML_RESTRICT         vx,
764                                         const void * const * GGML_RESTRICT vy,
765                                         int                                nc) {
766    const int nb = n / QK_K;
767
768    assert(n % QK_K == 0);
769    assert(nc % 8 == 0);
770
771    const block_q8_K * a_ptr[4];
772    for (int m = 0; m < 4; m++) {
773        a_ptr[m] = (const block_q8_K *) vy[m];
774    }
775
776    iqp_gemm_tile_4_generic(nb, s, bs, (const block_iqp_x8 *) vx, a_ptr, nc);
777}
778
779#if defined(__AVX2__)
780
781// add int16_t pairwise and return as 256 bit int vector, then add the accumulator
782static inline __m256i sum_i16_pairs_acc_int32x8(const __m256i acc, const __m256i x) {
783    const __m256i ones = _mm256_set1_epi16(1);
784    return _mm256_add_epi32(acc, _mm256_madd_epi16(ones, x));
785}
786
787static inline __m256i mul_sum_us8_pairs_acc_int32x8(const __m256i acc, const __m256i ax, const __m256i sy) {
788#    if defined(__AVX512VNNI__) && defined(__AVX512VL__)
789    return _mm256_dpbusd_epi32(acc, ax, sy);
790#    elif defined(__AVXVNNI__)
791    return _mm256_dpbusd_avx_epi32(acc, ax, sy);
792#    else
793    // Perform multiplication and create 16-bit values
794    const __m256i dot = _mm256_maddubs_epi16(ax, sy);
795    return sum_i16_pairs_acc_int32x8(acc, dot);
796#    endif
797}
798
799// Integer variant of the function defined in ggml-quants.c
800// multiply int8_t, add results pairwise twice and return as 256 bit int vector, then add the accumulator
801static inline __m256i mul_sum_i8_pairs_acc_int32x8(const __m256i acc, const __m256i x, const __m256i y) {
802#    if defined(__AVXVNNIINT8__)
803    return _mm256_dpbssd_epi32(acc, x, y);
804#    else
805    // Get absolute values of x vectors
806    const __m256i ax = _mm256_sign_epi8(x, x);
807    // Sign the values of the y vectors
808    const __m256i sy = _mm256_sign_epi8(y, x);
809    return mul_sum_us8_pairs_acc_int32x8(acc, ax, sy);
810#    endif
811}
812
813// load the 16 activations of one sub-block, offset by 128 when they are fed to dpbusd as unsigned bytes
814static inline __m256i iqp_load_y(const int8_t * GGML_RESTRICT qs) {
815    __m128i y = _mm_loadu_si128((const __m128i *) qs);
816#    if GGML_IQP_USE_BIAS
817    y = _mm_xor_si128(y, _mm_set1_epi8((char) 0x80));
818#    endif
819    return _mm256_broadcastsi128_si256(y);
820}
821
822// xv: 8 rows x 4 signed weights, yb: the matching 4 activation bytes broadcast to all 8 lanes
823static inline __m256i iqp_dot4(const __m256i acc, const __m256i xv, const __m256i yb) {
824#    if GGML_IQP_USE_BIAS
825    return mul_sum_us8_pairs_acc_int32x8(acc, yb, xv);
826#    else
827    return mul_sum_i8_pairs_acc_int32x8(acc, xv, yb);
828#    endif
829}
830
831static inline __m256i iqp_load_iscales(const int8_t * GGML_RESTRICT iscales) {
832    return _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *) iscales));
833}
834
835// accumulate one super-block of 8 interleaved rows against one q8_K row in int32; worst case 16 * 32 * 16 * 255 * 127 = 2.65e8 plus a bias of at most 1.33e8 does not overflow
836static inline __m256i iqp_acc_block(const block_iqp_x8 * GGML_RESTRICT b, const block_q8_K * GGML_RESTRICT a) {
837    __m256i sumi = _mm256_setzero_si256();
838
839    for (int sb = 0; sb < IQP_NSB; sb++) {
840        const int8_t * qs = b->qs + sb * 128;
841
842        const __m256i yv = iqp_load_y(a->qs + sb * 16);
843
844        __m256i isum = _mm256_setzero_si256();
845
846        isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 0)), _mm256_shuffle_epi32(yv, 0x00));
847        isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 32)), _mm256_shuffle_epi32(yv, 0x55));
848        isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 64)), _mm256_shuffle_epi32(yv, 0xAA));
849        isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 96)), _mm256_shuffle_epi32(yv, 0xFF));
850
851        sumi = _mm256_add_epi32(sumi, _mm256_mullo_epi32(isum, iqp_load_iscales(b->iscales + sb * 8)));
852    }
853
854#    if GGML_IQP_USE_BIAS
855    sumi = _mm256_sub_epi32(sumi, _mm256_loadu_si256((const __m256i *) b->bias));
856#    endif
857
858    return sumi;
859}
860
861// one 4 row x nc column tile; s points at the first of the four output rows, bs floats apart
862static inline void iqp_gemm_tile_4(int                                nb,
863                                   float * GGML_RESTRICT              s,
864                                   size_t                             bs,
865                                   const block_iqp_x8 * GGML_RESTRICT b_ptr_start,
866                                   const block_q8_K * const           a_ptr[4],
867                                   int                                nc) {
868    const int ncols_interleaved = 8;
869
870    for (int x = 0; x < nc / ncols_interleaved; x++) {
871        const block_iqp_x8 * b_ptr = b_ptr_start + x * nb;
872
873        __m256 sumf[4];
874        for (int m = 0; m < 4; m++) {
875            sumf[m] = _mm256_setzero_ps();
876        }
877
878        for (int l = 0; l < nb; l++) {
879            __m256i sumi[4];
880            for (int m = 0; m < 4; m++) {
881                sumi[m] = _mm256_setzero_si256();
882            }
883
884            for (int sb = 0; sb < IQP_NSB; sb++) {
885                const int8_t * qs = b_ptr[l].qs + sb * 128;
886
887                __m256i yv[4];
888                __m256i isum[4];
889                for (int m = 0; m < 4; m++) {
890                    yv[m]   = iqp_load_y(a_ptr[m][l].qs + sb * 16);
891                    isum[m] = _mm256_setzero_si256();
892                }
893
894                const __m256i xv0 = _mm256_loadu_si256((const __m256i *) (qs + 0));
895                const __m256i xv1 = _mm256_loadu_si256((const __m256i *) (qs + 32));
896                const __m256i xv2 = _mm256_loadu_si256((const __m256i *) (qs + 64));
897                const __m256i xv3 = _mm256_loadu_si256((const __m256i *) (qs + 96));
898
899                for (int m = 0; m < 4; m++) {
900                    isum[m] = iqp_dot4(isum[m], xv0, _mm256_shuffle_epi32(yv[m], 0x00));
901                    isum[m] = iqp_dot4(isum[m], xv1, _mm256_shuffle_epi32(yv[m], 0x55));
902                    isum[m] = iqp_dot4(isum[m], xv2, _mm256_shuffle_epi32(yv[m], 0xAA));
903                    isum[m] = iqp_dot4(isum[m], xv3, _mm256_shuffle_epi32(yv[m], 0xFF));
904                }
905
906                const __m256i isc = iqp_load_iscales(b_ptr[l].iscales + sb * 8);
907                for (int m = 0; m < 4; m++) {
908                    sumi[m] = _mm256_add_epi32(sumi[m], _mm256_mullo_epi32(isum[m], isc));
909                }
910            }
911
912#    if GGML_IQP_USE_BIAS
913            const __m256i bias = _mm256_loadu_si256((const __m256i *) b_ptr[l].bias);
914            for (int m = 0; m < 4; m++) {
915                sumi[m] = _mm256_sub_epi32(sumi[m], bias);
916            }
917#    endif
918
919            const __m256 dfac = _mm256_loadu_ps(b_ptr[l].dfac);
920            for (int m = 0; m < 4; m++) {
921                sumf[m] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(sumi[m]),
922                                          _mm256_mul_ps(dfac, _mm256_set1_ps(a_ptr[m][l].d)), sumf[m]);
923            }
924        }
925
926        for (int m = 0; m < 4; m++) {
927            _mm256_storeu_ps(s + m * bs + x * ncols_interleaved, sumf[m]);
928        }
929    }
930}
931
932#endif  // __AVX2__
933
934static void iqp_gemv_8x8_q8_K(int                        n,
935                              float * GGML_RESTRICT      s,
936                              size_t                     bs,
937                              const void * GGML_RESTRICT vx,
938                              const void * GGML_RESTRICT vy,
939                              int                        nr,
940                              int                        nc) {
941    const int nb                = n / QK_K;
942    const int ncols_interleaved = 8;
943
944    assert(n % QK_K == 0);
945    assert(nc % ncols_interleaved == 0);
946
947    UNUSED(bs);
948    UNUSED(nr);
949    UNUSED(nb);
950    UNUSED(ncols_interleaved);
951
952#if defined(__AVX2__)
953    const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx;
954    const block_q8_K *   a_ptr       = (const block_q8_K *) vy;
955
956    for (int x = 0; x < nc / ncols_interleaved; x++) {
957        const block_iqp_x8 * b_ptr = b_ptr_start + x * nb;
958
959        __m256 sumf = _mm256_setzero_ps();
960
961        for (int l = 0; l < nb; l++) {
962            const __m256 dv = _mm256_mul_ps(_mm256_loadu_ps(b_ptr[l].dfac), _mm256_set1_ps(a_ptr[l].d));
963
964            sumf = _mm256_fmadd_ps(_mm256_cvtepi32_ps(iqp_acc_block(b_ptr + l, a_ptr + l)), dv, sumf);
965        }
966
967        _mm256_storeu_ps(s + x * ncols_interleaved, sumf);
968    }
969
970    return;
971#endif
972
973    iqp_gemv_8x8_q8_K_generic(n, s, bs, vx, vy, nr, nc);
974}
975
976static void iqp_gemm_8x8_q8_K(int                        n,
977                              float * GGML_RESTRICT      s,
978                              size_t                     bs,
979                              const void * GGML_RESTRICT vx,
980                              const void * GGML_RESTRICT vy,
981                              int                        nr,
982                              int                        nc) {
983    const int nb                = n / QK_K;
984    const int ncols_interleaved = 8;
985
986    assert(n % QK_K == 0);
987    assert(nr % 4 == 0);
988    assert(nc % ncols_interleaved == 0);
989
990    UNUSED(nb);
991    UNUSED(ncols_interleaved);
992
993#if defined(__AVX2__)
994    const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx;
995    const block_q8_K *   a_ptr_start = (const block_q8_K *) vy;
996
997    for (int y = 0; y < nr / 4; y++) {
998        const block_q8_K * a_ptr[4];
999        for (int m = 0; m < 4; m++) {
1000            a_ptr[m] = a_ptr_start + (y * 4 + m) * nb;
1001        }
1002
1003        iqp_gemm_tile_4(nb, s + y * 4 * bs, bs, b_ptr_start, a_ptr, nc);
1004    }
1005
1006    return;
1007#endif
1008
1009    iqp_gemm_8x8_q8_K_generic(n, s, bs, vx, vy, nr, nc);
1010}
1011
1012// same as iqp_gemm_8x8_q8_K with nr = 4, but the activation rows are passed as separate pointers (for the scattered rows of MUL_MAT_ID)
1013static void iqp_gemm_8x8_q8_K_p4(int                                n,
1014                                 float * GGML_RESTRICT              s,
1015                                 size_t                             bs,
1016                                 const void * GGML_RESTRICT         vx,
1017                                 const void * const * GGML_RESTRICT vy,
1018                                 int                                nc) {
1019    const int nb                = n / QK_K;
1020    const int ncols_interleaved = 8;
1021
1022    assert(n % QK_K == 0);
1023    assert(nc % ncols_interleaved == 0);
1024
1025    UNUSED(nb);
1026    UNUSED(ncols_interleaved);
1027
1028#if defined(__AVX2__)
1029    const block_q8_K * a_ptr[4];
1030    for (int m = 0; m < 4; m++) {
1031        a_ptr[m] = (const block_q8_K *) vy[m];
1032    }
1033
1034    iqp_gemm_tile_4(nb, s, bs, (const block_iqp_x8 *) vx, a_ptr, nc);
1035
1036    return;
1037#endif
1038
1039    iqp_gemm_8x8_q8_K_p4_generic(n, s, bs, vx, vy, nc);
1040}
1041
1042static bool iqp_type_supported(enum ggml_type type) {
1043    switch (type) {
1044#define IQP_CASE(E, name) case GGML_TYPE_##E:
1045        IQP_TYPE_LIST(IQP_CASE)
1046#undef IQP_CASE
1047        return true;
1048        default:
1049            return false;
1050    }
1051}
1052
1053static bool iqp_supported_common(const struct ggml_tensor * dst) {
1054    const struct ggml_tensor * src0 = dst->src[0];
1055    const struct ggml_tensor * src1 = dst->src[1];
1056
1057    if (!iqp_type_supported(src0->type)) {
1058        return false;
1059    }
1060
1061    // the path assumes the src1 conversion type is q8_K
1062    if (ggml_get_type_traits_cpu(src0->type)->vec_dot_type != GGML_TYPE_Q8_K) {
1063        return false;
1064    }
1065
1066    // escape hatch to A/B the panel against the plain vec_dot path without rebuilding (--no-repack does not cover this path)
1067    static const bool disabled = getenv("GGML_NO_IQ_PANEL") != nullptr;
1068    if (disabled) {
1069        return false;
1070    }
1071
1072    if (!ggml_cpu_has_avx2()) {
1073        return false;
1074    }
1075
1076    if (src1->type != GGML_TYPE_F32) {
1077        return false;
1078    }
1079
1080    if (src0->ne[0] % QK_K != 0 || src0->ne[1] % IQP_NB_ROWS != 0) {
1081        return false;
1082    }
1083
1084    if (src0->ne[3] != 1 || src1->ne[3] != 1 || !ggml_is_contiguous(src0)) {
1085        return false;
1086    }
1087
1088    if (dst->type != GGML_TYPE_F32 || dst->nb[0] != sizeof(float)) {
1089        return false;
1090    }
1091
1092    return true;
1093}
1094
1095bool ggml_cpu_iqp_supports_mul_mat(const struct ggml_tensor * dst) {
1096    const struct ggml_tensor * src0 = dst->src[0];
1097    const struct ggml_tensor * src1 = dst->src[1];
1098
1099    if (!iqp_supported_common(dst)) {
1100        return false;
1101    }
1102
1103    if (src1->ne[1] < GGML_IQP_MIN_BATCH) {
1104        return false;
1105    }
1106
1107    // plain 2D weight matmuls only (src1 may still be batched over ne12)
1108    if (src0->ne[2] != 1) {
1109        return false;
1110    }
1111
1112    return true;
1113}
1114
1115bool ggml_cpu_iqp_supports_mul_mat_id(const struct ggml_tensor * dst) {
1116    const struct ggml_tensor * ids = dst->src[2];
1117
1118    if (!iqp_supported_common(dst)) {
1119        return false;
1120    }
1121
1122    // skip the node entirely (work buffer included) if no expert can reach the per expert threshold
1123    if (!ggml_cpu_iqp_mul_mat_id_min_batch(ids->ne[0] * ids->ne[1])) {
1124        return false;
1125    }
1126
1127    return true;
1128}
1129
1130void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * params,
1131                                         struct ggml_tensor *               dst,
1132                                         int64_t                            cur_a,
1133                                         int64_t                            cne1,
1134                                         const int32_t *                    expert_rows,
1135                                         void *                             panels) {
1136    const struct ggml_tensor * src0 = dst->src[0];
1137    const struct ggml_tensor * src1 = dst->src[1];
1138
1139    GGML_TENSOR_BINARY_OP_LOCALS
1140
1141    const int ith = params->ith;
1142    const int nth = params->nth;
1143
1144    const int64_t nblocks = ne00 / QK_K;
1145
1146    const size_t nbw1 = ggml_cpu_iqp_row_size(dst);
1147
1148    block_iqp_x8 * panel = (block_iqp_x8 *) ((char *) panels + (size_t) ith * ggml_cpu_iqp_scratch_size(dst));
1149
1150    const char * src0_cur = (const char *) src0->data + cur_a * nb02;
1151
1152    const int64_t ngroups = ne01 / IQP_NB_ROWS;
1153
1154    const int64_t g0 = (ngroups * ith) / nth;
1155    const int64_t g1 = (ngroups * (ith + 1)) / nth;
1156
1157    for (int64_t g = g0; g < g1; g++) {
1158        const int64_t r = g * IQP_NB_ROWS;
1159
1160        iqp_decode_panel_8(src0->type, src0_cur + r * nb01, nb01, nblocks, panel);
1161
1162        // the dst rows are scattered, so the gemm writes into tmp and it is copied out row by row
1163        float tmp[4 * IQP_NB_ROWS];
1164
1165        for (int64_t k = 0; k < cne1; k += 4) {
1166            const int64_t nrows = MIN(4, cne1 - k);
1167
1168            // a short tail tile duplicates its last row into the unused slots; the padding is never copied out
1169            const void * rows[4];
1170
1171            for (int64_t m = 0; m < 4; m++) {
1172                const int64_t kk = k + MIN(m, nrows - 1);
1173
1174                rows[m] = (const char *) params->wdata +
1175                          ((expert_rows[2 * kk + 0] % ne11) + expert_rows[2 * kk + 1] * ne11) * nbw1;
1176            }
1177
1178            iqp_gemm_8x8_q8_K_p4(ne00, tmp, IQP_NB_ROWS, panel, rows, IQP_NB_ROWS);
1179
1180            for (int64_t m = 0; m < nrows; m++) {
1181                float * dst_col = (float *) ((char *) dst->data + expert_rows[2 * (k + m) + 0] * nb1 +
1182                                             expert_rows[2 * (k + m) + 1] * nb2);
1183                memcpy(dst_col + r, tmp + m * IQP_NB_ROWS, IQP_NB_ROWS * sizeof(float));
1184            }
1185        }
1186    }
1187}
1188
1189size_t ggml_cpu_iqp_scratch_size(const struct ggml_tensor * dst) {
1190    return GGML_PAD((dst->src[0]->ne[0] / QK_K) * sizeof(block_iqp_x8), 64);
1191}
1192
1193void ggml_compute_forward_mul_mat_iqp(const struct ggml_compute_params * params, struct ggml_tensor * dst) {
1194    const struct ggml_tensor * src0 = dst->src[0];
1195    const struct ggml_tensor * src1 = dst->src[1];
1196
1197    GGML_TENSOR_BINARY_OP_LOCALS
1198
1199    const int ith = params->ith;
1200    const int nth = params->nth;
1201
1202    const int64_t nblocks = ne00 / QK_K;
1203
1204    const size_t nbw1 = ggml_row_size(GGML_TYPE_Q8_K, ne10);
1205    const size_t nbw2 = nbw1 * ne11;
1206
1207    const size_t scratch_size = ggml_cpu_iqp_scratch_size(dst);
1208
1209    const size_t scratch_offset = GGML_PAD(nbw2 * ne12, 64);
1210
1211    GGML_ASSERT(scratch_offset + (size_t) nth * scratch_size <= params->wsize);
1212
1213    block_iqp_x8 * panel = (block_iqp_x8 *) ((char *) params->wdata + scratch_offset + (size_t) ith * scratch_size);
1214
1215    const int64_t nrows = ne11;
1216
1217    const int64_t ngroups = ne01 / IQP_NB_ROWS;
1218
1219    // aim for 4 chunks per thread; the caller has already reset the chunk counter
1220    // on NUMA systems fall back to one chunk per thread
1221    const int64_t chunks_per_thread = ggml_is_numa() ? 1 : 4;
1222    const int64_t groups_per_chunk  = MAX(1, (ngroups + nth * chunks_per_thread - 1) / (nth * chunks_per_thread));
1223    const int64_t nchunk            = (ngroups + groups_per_chunk - 1) / groups_per_chunk;
1224
1225    int current_chunk = ith;
1226
1227    while (current_chunk < nchunk) {
1228        const int64_t g0 = current_chunk * groups_per_chunk;
1229        const int64_t g1 = MIN(g0 + groups_per_chunk, ngroups);
1230
1231        for (int64_t g = g0; g < g1; g++) {
1232            const int64_t r = g * IQP_NB_ROWS;
1233
1234            iqp_decode_panel_8(src0->type, (const char *) src0->data + r * nb01, nb01, nblocks, panel);
1235
1236            for (int64_t i12 = 0; i12 < ne12; i12++) {
1237                const char * src1_ptr = (const char *) params->wdata + i12 * nbw2;
1238                char *       dst_ptr  = (char *) dst->data + i12 * nb2;
1239
1240                if (nrows > 3) {
1241                    iqp_gemm_8x8_q8_K(ne00, (float *) dst_ptr + r, nb1 / nb0, panel, src1_ptr, nrows - (nrows % 4),
1242                                      IQP_NB_ROWS);
1243                }
1244                for (int64_t iter = nrows - (nrows % 4); iter < nrows; iter++) {
1245                    iqp_gemv_8x8_q8_K(ne00, (float *) (dst_ptr + iter * nb1) + r, ne01, panel, src1_ptr + nbw1 * iter,
1246                                      1 /* nrows */, IQP_NB_ROWS);
1247                }
1248            }
1249        }
1250
1251        current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);
1252    }
1253}