char-fork/llama.cpp

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

alex-spacemitggml-cpu: Add IME2 Instruction Support for the SpacemiT Backend (#22863)81b0d882a

char/aliceai-t5-moe
284.9 KiB5768 linesraw
1#include "ggml-impl.h"
2#include "ggml.h"
3#include "ime_kernels.h"
4#include "rvv_kernels.h"
5#include "string.h"
6
7#include <algorithm>
8#include <cmath>
9#include <stdexcept>
10
11#if !defined(__riscv_v) || !defined(__riscv_v_intrinsic)
12#    error "riscv v extension or v_intrinsic not enabled"
13#else
14#    include <riscv_vector.h>
15#endif
16
17#if !defined(__riscv_zfh)
18#    error "riscv zfh extension not enabled"
19#endif
20
21#if defined(RISCV64_SPACEMIT_IME2)
22#else
23#    error "RISCV64_SPACEMIT_IME2 not defined"
24#endif
25
26#if defined(__GNUC__)
27#    pragma GCC diagnostic ignored "-Woverlength-strings"
28#    pragma GCC diagnostic ignored "-Wcast-qual"
29#    pragma GCC diagnostic ignored "-Wunused-parameter"
30#endif
31
32namespace spacemit_kernels {
33namespace ime2 {
34
35template <size_t MB_ROWS, size_t NB_COLS>
36void gemm_kernel_i8i2k_mrow_ref(size_t          blk_len,
37                                const uint8_t * quant_a_ptr,
38                                const uint8_t * quant_b_data,
39                                float *         c_ptr,
40                                size_t          count_m,
41                                size_t          count_n,
42                                size_t          k_blks,
43                                size_t          ldc) {
44    using blk_type                 = nrow_block_q2_k<NB_COLS>;
45    constexpr float refactor_scale = 16.0f;
46    constexpr float factor_scale   = 1.0f / refactor_scale;
47
48    int64_t a_blk_stride        = q8k_blk_size(256);
49    int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS;
50    int64_t b_ncol_block_stride = sizeof(blk_type);
51
52    float      output[MB_ROWS * NB_COLS]     = { 0 };
53    _Float16   output_f16[MB_ROWS * NB_COLS] = { 0 };
54    blk_type * quant_b_blk_data              = (blk_type *) (quant_b_data);
55
56    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
57        size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
58
59        int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS * 16;
60
61        for (size_t mi = 0; mi < MB_ROWS; mi++) {
62            for (size_t ci = 0; ci < NB_COLS; ci++) {
63                output[ci + mi * NB_COLS] = 0;
64            }
65        }
66
67        for (size_t ki = 0; ki < k_blks; ki++, quant_b_blk_data++, a_data += a_nrow_block_stride) {
68            uint8_t * b_data   = quant_b_blk_data->qs;
69            uint8_t * scales   = quant_b_blk_data->scales;
70            uint8_t * scales16 = (uint8_t *) (quant_b_blk_data->scales16);
71            uint8_t * zeros16  = (uint8_t *) (quant_b_blk_data->zeros16);
72
73            _Float16 * scales_fp16 = (_Float16 *) scales16;
74            _Float16 * zeros_fp16  = (_Float16 *) zeros16;
75
76            float *   a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS * 16);
77            int16_t * a_sum_row   = (int16_t *) (a_data - sizeof(int16_t) * MB_ROWS * 16);
78
79            memset(output_f16, 0, sizeof(output_f16));
80
81            uint8_t * scales_temp = scales;
82            uint8_t * zps_temp    = scales;
83            for (size_t kii = 0; kii < 16; kii++, scales_temp += NB_COLS, zps_temp++) {
84                size_t b_shift = (kii % 4) * 2;
85
86                uint8_t * b_data_col = b_data + (kii / 4) * NB_COLS * 16;
87
88                for (size_t mi = 0; mi < MB_ROWS; mi++) {
89                    int16_t a_sum = a_sum_row[mi * 16 + kii];
90                    for (size_t ci = 0; ci < NB_COLS; ci++) {
91                        _Float16 acc_0 = 0.0;
92
93                        uint8_t b_zp    = zps_temp[ci * 16] >> 4;
94                        uint8_t b_scale = scales_temp[ci] & 0x0F;
95                        for (size_t bi = 0; bi < 16; bi++) {
96                            int8_t  a0 = a_data[mi * 256 + bi + kii * 16];
97                            uint8_t b0 = b_data_col[ci * 16 + bi];
98                            acc_0 += static_cast<int16_t>(a0) * static_cast<int16_t>((b0 >> b_shift) & 0x03);
99                        }
100
101                        _Float16 scale_item =
102                            static_cast<_Float16>(b_scale) * static_cast<_Float16>(factor_scale) * scales_fp16[ci];
103
104                        output_f16[ci + mi * NB_COLS] += acc_0 * scale_item;
105                        output[ci + mi * NB_COLS] += b_zp * a_sum * a_scale_row[mi] * zeros_fp16[ci];
106                    }
107                }
108            }
109
110            for (size_t mi = 0; mi < MB_ROWS; mi++) {
111                auto a_scale = a_scale_row[mi] * refactor_scale;
112                for (size_t ci = 0; ci < NB_COLS; ci++) {
113                    output[ci + mi * NB_COLS] += output_f16[ci + mi * NB_COLS] * a_scale;
114                }
115            }
116        }
117
118        for (size_t mi = 0; mi < MB_ROWS; mi++) {
119            for (size_t ci = 0; ci < nb_real; ci++) {
120                c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
121            }
122        }
123    }
124}
125
126template <size_t MB_ROWS, size_t NB_COLS>
127void gemm_kernel_i8i3k_mrow_ref(size_t          blk_len,
128                                const uint8_t * quant_a_ptr,
129                                const uint8_t * quant_b_data,
130                                float *         c_ptr,
131                                size_t          count_m,
132                                size_t          count_n,
133                                size_t          k_blks,
134                                size_t          ldc) {
135    using blk_type                 = nrow_block_q2_k<NB_COLS>;
136    constexpr float refactor_scale = 16.0f;
137    constexpr float factor_scale   = 1.0f / refactor_scale;
138
139    int64_t a_blk_stride        = q8k_blk_size(256);
140    int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS;
141    int64_t b_ncol_block_stride = sizeof(blk_type);
142
143    float    output[MB_ROWS * NB_COLS]     = { 0 };
144    _Float16 output_f16[MB_ROWS * NB_COLS] = { 0 };
145
146    blk_type * quant_b_blk_data = (blk_type *) (quant_b_data);
147
148    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
149        size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
150
151        int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS * 16;
152
153        for (size_t mi = 0; mi < MB_ROWS; mi++) {
154            for (size_t ci = 0; ci < NB_COLS; ci++) {
155                output[ci + mi * NB_COLS] = 0;
156            }
157        }
158
159        for (size_t ki = 0; ki < k_blks; ki++, quant_b_blk_data++, a_data += a_nrow_block_stride) {
160            uint8_t * b_data   = quant_b_blk_data->qs;
161            uint8_t * b_hmask  = quant_b_blk_data->hmask;
162            int8_t *  scales   = quant_b_blk_data->scales;
163            uint8_t * scales16 = (uint8_t *) (quant_b_blk_data->scales16);
164
165            _Float16 * scales_fp16 = (_Float16 *) scales16;
166
167            float *   a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS * 16);
168            int16_t * a_sum_row   = (int16_t *) (a_data - sizeof(int16_t) * MB_ROWS * 16);
169
170            memset(output_f16, 0, sizeof(output_f16));
171
172            int8_t *   scales_temp = scales;
173            uint16_t * b_mask_col  = (uint16_t *) b_hmask;
174
175            float acc_0_max = 0.0f;
176            for (size_t kii = 0; kii < 16; kii++, scales_temp += NB_COLS, b_mask_col += NB_COLS) {
177                size_t b_shift = (kii % 4) * 2;
178
179                uint8_t * b_data_col = b_data + (kii / 4) * NB_COLS * 16;
180
181                for (size_t mi = 0; mi < MB_ROWS; mi++) {
182                    for (size_t ci = 0; ci < NB_COLS; ci++) {
183                        _Float16 acc_0        = 0;
184                        // blk 2 * kii + 0
185                        uint16_t b_shift_mask = 1;
186                        for (size_t bi = 0; bi < 16; bi++, b_shift_mask <<= 1) {
187                            int8_t a0 = a_data[mi * 256 + bi + kii * 16];
188                            int8_t b0 = static_cast<int8_t>((b_data_col[ci * 16 + bi] >> b_shift) & 0x03);
189                            b0 -= b_mask_col[ci] & b_shift_mask ? 0 : 4;
190                            acc_0 += static_cast<int16_t>(a0) * static_cast<int16_t>(b0);
191                        }
192
193                        _Float16 scale_item = static_cast<_Float16>(scales_temp[ci]) * scales_fp16[ci] *
194                                              static_cast<_Float16>(factor_scale);
195
196                        output_f16[ci + mi * NB_COLS] += acc_0 * scale_item;
197                    }
198                }
199            }
200
201            for (size_t mi = 0; mi < MB_ROWS; mi++) {
202                auto a_scale = a_scale_row[mi] * refactor_scale;
203                for (size_t ci = 0; ci < NB_COLS; ci++) {
204                    output[ci + mi * NB_COLS] += output_f16[ci + mi * NB_COLS] * a_scale;
205                }
206            }
207        }
208
209        for (size_t mi = 0; mi < MB_ROWS; mi++) {
210            for (size_t ci = 0; ci < nb_real; ci++) {
211                c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
212            }
213        }
214    }
215}
216
217template <size_t MB_ROWS, size_t NB_COLS>
218void gemm_kernel_i8i4_mrow_ref(size_t          blk_len,
219                               const uint8_t * quant_a_ptr,
220                               const uint8_t * quant_b_data,
221                               const uint8_t * quant_b_zp,
222                               float *         c_ptr,
223                               size_t          count_m,
224                               size_t          count_n,
225                               size_t          k_blks,
226                               size_t          ldc) {
227    constexpr size_t kblks_per_blk = 16;
228    GGML_ASSERT(k_blks % kblks_per_blk == 0);
229
230    int64_t b_blk_stride        = (sizeof(_Float16) + (blk_len / 2) + (quant_b_zp ? sizeof(uint8_t) : 0));
231    int64_t b_stride            = k_blks * b_blk_stride;
232    int64_t a_blk_stride        = q8_blk_size(blk_len, true);
233    int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS;
234    int64_t b_ncol_block_stride = b_blk_stride * NB_COLS;
235
236    float    output[MB_ROWS * NB_COLS]     = { 0 };
237    _Float16 output_f16[MB_ROWS * NB_COLS] = { 0 };
238
239    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
240        size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
241        uint8_t * b_data  = (uint8_t *) quant_b_data + ni * b_stride + NB_COLS * sizeof(_Float16);
242        if (quant_b_zp) {
243            b_data += NB_COLS * sizeof(uint8_t);
244        }
245
246        int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS;
247
248        for (size_t mi = 0; mi < MB_ROWS; mi++) {
249            for (size_t ci = 0; ci < NB_COLS; ci++) {
250                output[ci + mi * NB_COLS]     = 0.0f;
251                output_f16[ci + mi * NB_COLS] = static_cast<_Float16>(0.0f);
252            }
253        }
254
255        size_t kii = 0;
256        for (size_t ki = 0; ki < k_blks; ki++, a_data += a_nrow_block_stride, b_data += b_ncol_block_stride) {
257            _Float16 * b_scale_fp16 = (_Float16 *) (b_data - NB_COLS * sizeof(_Float16));
258            uint8_t *  b_zp         = nullptr;
259            if (quant_b_zp) {
260                b_scale_fp16 = (_Float16 *) (b_data - NB_COLS * sizeof(_Float16) - NB_COLS * sizeof(uint8_t));
261                b_zp         = (uint8_t *) (b_data - NB_COLS * sizeof(uint8_t));
262            }
263
264            float *   a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS);
265            int16_t * a_sum_row   = (int16_t *) (a_data - sizeof(int16_t) * MB_ROWS);
266
267            for (size_t mi = 0; mi < MB_ROWS; mi++) {
268                _Float16 a_scale = a_scale_row[mi];
269                int16_t  a_sum   = a_sum_row[mi];
270
271                for (size_t ci = 0; ci < NB_COLS; ci++) {
272                    _Float16 b_scale = b_scale_fp16[ci];
273                    int32_t  acc     = 0;
274                    if (b_zp) {
275                        acc += a_sum * b_zp[ci];
276                    } else {
277                        acc += a_sum * 8;
278                    }
279                    for (size_t bi = 0; bi < blk_len / 2; bi++) {
280                        int8_t  a0 = a_data[mi * blk_len + 2 * bi];
281                        int8_t  a1 = a_data[mi * blk_len + 2 * bi + 1];
282                        uint8_t b  = b_data[ci * blk_len / 2 + bi];
283                        int8_t  b0 = static_cast<int8_t>(b & 0x0F);
284                        int8_t  b1 = static_cast<int8_t>((b & 0xF0) >> 4);
285                        acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
286                               static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
287                    }
288                    output_f16[ci + mi * NB_COLS] +=
289                        static_cast<float>(acc) * static_cast<float>(a_scale) * static_cast<float>(b_scale);
290                }
291            }
292
293            if (kii == kblks_per_blk - 1) {
294                for (size_t mi = 0; mi < MB_ROWS; mi++) {
295                    for (size_t ci = 0; ci < NB_COLS; ci++) {
296                        output[ci + mi * NB_COLS] += static_cast<float>(output_f16[ci + mi * NB_COLS]);
297                        output_f16[ci + mi * NB_COLS] = 0.0f;
298                    }
299                }
300                kii = 0;
301            } else {
302                kii++;
303            }
304        }
305
306        if (kii == kblks_per_blk - 1) {
307            for (size_t mi = 0; mi < MB_ROWS; mi++) {
308                for (size_t ci = 0; ci < NB_COLS; ci++) {
309                    output[ci + mi * NB_COLS] += static_cast<float>(output_f16[ci + mi * NB_COLS]);
310                    output_f16[ci + mi * NB_COLS] = 0.0f;
311                }
312            }
313            kii = 0;
314        }
315
316        for (size_t mi = 0; mi < MB_ROWS; mi++) {
317            for (size_t ci = 0; ci < nb_real; ci++) {
318                c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
319            }
320        }
321    }
322}
323
324template <size_t MB_ROWS, size_t NB_COLS>
325void gemm_kernel_i8i4_hp_mrow_ref(size_t          blk_len,
326                                  const uint8_t * quant_a_ptr,
327                                  const uint8_t * quant_b_data,
328                                  const uint8_t * quant_b_zp,
329                                  float *         c_ptr,
330                                  size_t          count_m,
331                                  size_t          count_n,
332                                  size_t          k_blks,
333                                  size_t          ldc) {
334    constexpr size_t k_subblks_per_superblk = 8;
335
336    struct block_q4_0x32_layout {
337        _Float16 d[NB_COLS];
338        uint8_t  qs[16 * NB_COLS];
339    };
340
341    GGML_ASSERT(blk_len == 256);
342
343    const size_t b_superblk_stride = sizeof(block_q4_0x32_layout) * k_subblks_per_superblk +
344                                     (quant_b_zp ? NB_COLS * k_subblks_per_superblk * sizeof(uint8_t) : 0);
345    const size_t b_tile_stride = k_blks * b_superblk_stride;
346
347    const size_t a_nrow_block_stride = q8_hp_blk_size(blk_len, true, true) * MB_ROWS;
348    const size_t a_subblk_stride     = q8_hp_blk_size(32, false, false) * MB_ROWS;
349
350    float output[MB_ROWS * NB_COLS] = { 0 };
351    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
352        size_t          nb_real     = std::min<size_t>(NB_COLS, count_n - ni);
353        const uint8_t * b_tile_base = quant_b_data + (ni / NB_COLS) * b_tile_stride;
354        int8_t *        a_data      = (int8_t *) quant_a_ptr;
355
356        for (size_t mi = 0; mi < MB_ROWS; mi++) {
357            for (size_t ci = 0; ci < NB_COLS; ci++) {
358                output[ci + mi * NB_COLS] = 0.0f;
359            }
360        }
361
362        for (size_t ki = 0; ki < k_blks; ki++, a_data += a_nrow_block_stride) {
363            _Float16 output_f16[MB_ROWS * NB_COLS] = { 0 };
364
365            const uint8_t *              b_superblk_ptr = b_tile_base + ki * b_superblk_stride;
366            const block_q4_0x32_layout * b_blocks = reinterpret_cast<const block_q4_0x32_layout *>(b_superblk_ptr);
367            const uint8_t *              b_zps =
368                quant_b_zp ? b_superblk_ptr + sizeof(block_q4_0x32_layout) * k_subblks_per_superblk : nullptr;
369
370            _Float16 * a_sum_row       = (_Float16 *) (a_data + a_subblk_stride * k_subblks_per_superblk);
371            _Float16 * a_scale_avg_row = (_Float16 *) (a_data + a_nrow_block_stride - sizeof(_Float16) * MB_ROWS);
372            _Float16   scale_factor    = a_scale_avg_row[0];
373
374            for (size_t ksi = 0; ksi < k_subblks_per_superblk; ++ksi) {
375                const _Float16 * a_scale_row = reinterpret_cast<const _Float16 *>(a_data + a_subblk_stride * ksi);
376                int8_t *         a_subblk    = a_data + a_subblk_stride * ksi + MB_ROWS * sizeof(_Float16);
377                const _Float16   a_scale     = a_scale_row[0];
378                const block_q4_0x32_layout & b_block = b_blocks[ksi];
379
380                for (size_t mi = 0; mi < MB_ROWS; mi++) {
381                    for (size_t ci = 0; ci < NB_COLS; ci++) {
382                        const uint8_t * b_qs    = b_block.qs + ci * 16;
383                        _Float16        b_scale = b_block.d[ci] * a_scale;
384
385                        int16_t acc = 0;
386                        for (size_t bi = 0; bi < 16; bi++) {
387                            uint8_t b  = b_qs[bi];
388                            int8_t  b0 = static_cast<int8_t>(b & 0x0F);
389                            int8_t  b1 = static_cast<int8_t>((b & 0xF0) >> 4);
390
391                            acc += static_cast<int16_t>(a_subblk[mi * 32 + 2 * bi]) * static_cast<int16_t>(b0) +
392                                   static_cast<int16_t>(a_subblk[mi * 32 + 2 * bi + 1]) * static_cast<int16_t>(b1);
393                        }
394
395                        const _Float16 scaled_acc = static_cast<_Float16>(acc) * b_scale;
396                        output_f16[ci + mi * NB_COLS] += scaled_acc;
397                    }
398                }
399            }
400
401            for (size_t ksi = 0; ksi < k_subblks_per_superblk; ++ksi) {
402                const _Float16 * a_scale_row = reinterpret_cast<const _Float16 *>(a_data + a_subblk_stride * ksi);
403                const block_q4_0x32_layout & b_block  = b_blocks[ksi];
404                const uint8_t *              b_zp_row = b_zps ? b_zps + ksi * NB_COLS : nullptr;
405                const _Float16               a_scale  = a_scale_row[0];
406
407                for (size_t mi = 0; mi < MB_ROWS; mi++) {
408                    const _Float16 a_sum = a_sum_row[mi * k_subblks_per_superblk + ksi];
409                    for (size_t ci = 0; ci < NB_COLS; ci++) {
410                        _Float16 b_scale   = b_block.d[ci] * a_scale;
411                        _Float16 a_sum_bzp = a_sum;
412                        if (b_zp_row) {
413                            a_sum_bzp = a_sum * static_cast<_Float16>(0.125f) * static_cast<_Float16>(b_zp_row[ci]);
414                        }
415
416                        const _Float16 scaled_acc = a_sum_bzp * b_scale;
417                        output[ci + mi * NB_COLS] += scaled_acc * scale_factor;
418                    }
419                }
420            }
421
422            for (size_t mi = 0; mi < MB_ROWS; mi++) {
423                for (size_t ci = 0; ci < NB_COLS; ci++) {
424                    auto val = static_cast<float>(output_f16[ci + mi * NB_COLS]) * static_cast<float>(scale_factor);
425                    output[ci + mi * NB_COLS] += val;
426                }
427            }
428        }
429
430        for (size_t mi = 0; mi < MB_ROWS; mi++) {
431            for (size_t ci = 0; ci < nb_real; ci++) {
432                c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
433            }
434        }
435    }
436}
437
438template <size_t MB_ROWS, size_t NB_COLS>
439void moe_gemm_kernel_i8i4_mrow_ref(size_t           blk_len,
440                                   const uint8_t ** quant_a_ptr,
441                                   const uint8_t *  quant_b_data,
442                                   const uint8_t *  quant_b_zp,
443                                   float **         c_ptr,
444                                   size_t           count_m,
445                                   size_t           count_n,
446                                   size_t           k_blks,
447                                   size_t           ldc) {
448    int64_t b_blk_stride        = (sizeof(ggml_fp16_t) + (blk_len / 2) + (quant_b_zp ? sizeof(uint8_t) : 0));
449    int64_t b_stride            = k_blks * b_blk_stride;
450    int64_t a_blk_stride        = q8_blk_size(blk_len, true);
451    int64_t b_ncol_block_stride = b_blk_stride * NB_COLS;
452
453    float                         output[MB_ROWS * NB_COLS] = { 0 };
454    std::array<int8_t *, MB_ROWS> a_data;
455    std::array<float *, MB_ROWS>  c_data;
456
457    for (size_t mi = 0; mi < MB_ROWS; mi++) {
458        c_data[mi] = c_ptr[mi];
459    }
460
461    for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
462        size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
463        uint8_t * b_data  = (uint8_t *) quant_b_data + ni * b_stride + NB_COLS * sizeof(ggml_fp16_t);
464        if (quant_b_zp) {
465            b_data += NB_COLS * sizeof(uint8_t);
466        }
467
468        for (size_t mi = 0; mi < MB_ROWS; mi++) {
469            a_data[mi] = (int8_t *) quant_a_ptr[mi] + sizeof(float) + sizeof(int16_t);
470        }
471
472        for (size_t mi = 0; mi < MB_ROWS; mi++) {
473            for (size_t ci = 0; ci < NB_COLS; ci++) {
474                output[ci + mi * NB_COLS] = 0;
475            }
476        }
477
478        for (size_t ki = 0; ki < k_blks; ki++, b_data += b_ncol_block_stride) {
479            ggml_fp16_t * b_scale_fp16 = (ggml_fp16_t *) (b_data - NB_COLS * sizeof(ggml_fp16_t));
480            uint8_t *     b_zp         = nullptr;
481            if (quant_b_zp) {
482                b_scale_fp16 = (ggml_fp16_t *) (b_data - NB_COLS * sizeof(ggml_fp16_t) - NB_COLS * sizeof(uint8_t));
483                b_zp         = (uint8_t *) (b_data - NB_COLS * sizeof(uint8_t));
484            }
485
486            for (size_t mi = 0; mi < MB_ROWS; mi++) {
487                float *   a_scale_row = (float *) (a_data[mi] - sizeof(float) - sizeof(int16_t));
488                int16_t * a_sum_row   = (int16_t *) (a_data[mi] - sizeof(int16_t));
489
490                float   a_scale = *a_scale_row;
491                int16_t a_sum   = *a_sum_row;
492
493                for (size_t ci = 0; ci < NB_COLS; ci++) {
494                    float   b_scale = ggml_fp16_to_fp32(b_scale_fp16[ci]);
495                    int32_t acc     = 0;
496                    if (b_zp) {
497                        acc += a_sum * b_zp[ci];
498                    } else {
499                        acc += a_sum * 8;
500                    }
501                    for (size_t bi = 0; bi < blk_len / 2; bi++) {
502                        int8_t  a0 = (a_data[mi])[2 * bi];
503                        int8_t  a1 = (a_data[mi])[2 * bi + 1];
504                        uint8_t b  = b_data[ci * blk_len / 2 + bi];
505                        int8_t  b0 = static_cast<int8_t>(b & 0x0F);
506                        int8_t  b1 = static_cast<int8_t>((b & 0xF0) >> 4);
507                        acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
508                               static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
509                    }
510                    output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
511                }
512            }
513
514            for (size_t mi = 0; mi < MB_ROWS; mi++) {
515                a_data[mi] += a_blk_stride;
516            }
517        }
518
519        for (size_t mi = 0; mi < MB_ROWS; mi++) {
520            for (size_t ci = 0; ci < nb_real; ci++) {
521                (c_data[mi])[ci] = output[mi * NB_COLS + ci];
522            }
523        }
524
525        for (size_t mi = 0; mi < MB_ROWS; mi++) {
526            c_data[mi] += NB_COLS;
527        }
528    }
529}
530
531template <size_t MB_ROWS, size_t NB_COLS>
532void moe_gemm_kernel_i8i5_mrow_ref(size_t           blk_len,
533                                   const uint8_t ** quant_a_ptr,
534                                   const uint8_t *  quant_b_data,
535                                   const uint8_t *  quant_b_zp,
536                                   float **         c_ptr,
537                                   size_t           count_m,
538                                   size_t           count_n,
539                                   size_t           k_blks,
540                                   size_t           ldc) {
541    GGML_UNUSED(count_m);
542    GGML_UNUSED(ldc);
543
544    // blk_len is expected to be 32 for Q5 types.
545    int64_t a_blk_stride = q8_blk_size(blk_len, true);
546
547    float                         output[MB_ROWS * NB_COLS] = { 0 };
548    std::array<int8_t *, MB_ROWS> a_data;
549    std::array<float *, MB_ROWS>  c_data;
550
551    for (size_t mi = 0; mi < MB_ROWS; ++mi) {
552        c_data[mi] = c_ptr[mi];
553    }
554
555    if (quant_b_zp) {
556        using blk_type = nrow_block_q5_1<NB_COLS>;
557
558        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
559            size_t     nb_real          = std::min<size_t>(NB_COLS, count_n - ni);
560            blk_type * quant_b_blk_data = (blk_type *) quant_b_data + (ni / NB_COLS) * k_blks;
561
562            for (size_t mi = 0; mi < MB_ROWS; ++mi) {
563                a_data[mi] = (int8_t *) quant_a_ptr[mi] + sizeof(float) + sizeof(int16_t);
564            }
565
566            for (size_t mi = 0; mi < MB_ROWS; ++mi) {
567                for (size_t ci = 0; ci < NB_COLS; ++ci) {
568                    output[ci + mi * NB_COLS] = 0;
569                }
570            }
571
572            for (size_t ki = 0; ki < k_blks; ++ki, ++quant_b_blk_data) {
573                for (size_t mi = 0; mi < MB_ROWS; ++mi) {
574                    float *   a_scale_row = (float *) (a_data[mi] - sizeof(float) - sizeof(int16_t));
575                    int16_t * a_sum_row   = (int16_t *) (a_data[mi] - sizeof(int16_t));
576                    float     a_scale     = *a_scale_row;
577                    int16_t   a_sum       = *a_sum_row;
578
579                    for (size_t ci = 0; ci < NB_COLS; ++ci) {
580                        float   b_scale  = ggml_fp16_to_fp32(quant_b_blk_data->scales16[ci]);
581                        uint8_t b_zp_val = quant_b_blk_data->zp[ci];
582                        int32_t acc      = a_sum * static_cast<int32_t>(b_zp_val);
583
584                        for (size_t bi = 0; bi < blk_len / 2; ++bi) {
585                            int8_t  a0       = a_data[mi][2 * bi];
586                            int8_t  a1       = a_data[mi][2 * bi + 1];
587                            uint8_t qs_byte  = quant_b_blk_data->qs[ci * (blk_len / 2) + bi];
588                            int8_t  b0       = static_cast<int8_t>(qs_byte & 0x0F);
589                            int8_t  b1       = static_cast<int8_t>((qs_byte >> 4) & 0x0F);
590                            uint8_t qh_byte0 = quant_b_blk_data->qh[ci * 4 + (2 * bi) / 8];
591                            uint8_t qh_byte1 = quant_b_blk_data->qh[ci * 4 + (2 * bi + 1) / 8];
592                            uint8_t h0       = (qh_byte0 >> ((2 * bi) % 8)) & 1;
593                            uint8_t h1       = (qh_byte1 >> ((2 * bi + 1) % 8)) & 1;
594
595                            b0 |= (h0 << 4);
596                            b1 |= (h1 << 4);
597
598                            acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
599                                   static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
600                        }
601
602                        output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
603                    }
604
605                    a_data[mi] += a_blk_stride;
606                }
607            }
608
609            for (size_t mi = 0; mi < MB_ROWS; ++mi) {
610                for (size_t ci = 0; ci < nb_real; ++ci) {
611                    c_data[mi][ci] = output[mi * NB_COLS + ci];
612                }
613                c_data[mi] += NB_COLS;
614            }
615        }
616    } else {
617        using blk_type = nrow_block_q5_0<NB_COLS>;
618
619        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
620            size_t     nb_real          = std::min<size_t>(NB_COLS, count_n - ni);
621            blk_type * quant_b_blk_data = (blk_type *) quant_b_data + (ni / NB_COLS) * k_blks;
622
623            for (size_t mi = 0; mi < MB_ROWS; ++mi) {
624                a_data[mi] = (int8_t *) quant_a_ptr[mi] + sizeof(float) + sizeof(int16_t);
625            }
626
627            for (size_t mi = 0; mi < MB_ROWS; ++mi) {
628                for (size_t ci = 0; ci < NB_COLS; ++ci) {
629                    output[ci + mi * NB_COLS] = 0;
630                }
631            }
632
633            for (size_t ki = 0; ki < k_blks; ++ki, ++quant_b_blk_data) {
634                for (size_t mi = 0; mi < MB_ROWS; ++mi) {
635                    float *   a_scale_row = (float *) (a_data[mi] - sizeof(float) - sizeof(int16_t));
636                    int16_t * a_sum_row   = (int16_t *) (a_data[mi] - sizeof(int16_t));
637                    float     a_scale     = *a_scale_row;
638                    int16_t   a_sum       = *a_sum_row;
639
640                    for (size_t ci = 0; ci < NB_COLS; ++ci) {
641                        float   b_scale = ggml_fp16_to_fp32(quant_b_blk_data->scales16[ci]);
642                        int32_t acc     = a_sum * 16;
643
644                        for (size_t bi = 0; bi < blk_len / 2; ++bi) {
645                            int8_t  a0       = a_data[mi][2 * bi];
646                            int8_t  a1       = a_data[mi][2 * bi + 1];
647                            uint8_t qs_byte  = quant_b_blk_data->qs[ci * (blk_len / 2) + bi];
648                            int8_t  b0       = static_cast<int8_t>(qs_byte & 0x0F);
649                            int8_t  b1       = static_cast<int8_t>((qs_byte >> 4) & 0x0F);
650                            uint8_t qh_byte0 = quant_b_blk_data->qh[ci * 4 + (2 * bi) / 8];
651                            uint8_t qh_byte1 = quant_b_blk_data->qh[ci * 4 + (2 * bi + 1) / 8];
652                            uint8_t h0       = (qh_byte0 >> ((2 * bi) % 8)) & 1;
653                            uint8_t h1       = (qh_byte1 >> ((2 * bi + 1) % 8)) & 1;
654
655                            b0 |= (h0 << 4);
656                            b1 |= (h1 << 4);
657
658                            acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
659                                   static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
660                        }
661
662                        output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
663                    }
664
665                    a_data[mi] += a_blk_stride;
666                }
667            }
668
669            for (size_t mi = 0; mi < MB_ROWS; ++mi) {
670                for (size_t ci = 0; ci < nb_real; ++ci) {
671                    c_data[mi][ci] = output[mi * NB_COLS + ci];
672                }
673                c_data[mi] += NB_COLS;
674            }
675        }
676    }
677}
678
679template <size_t MB_ROWS, size_t NB_COLS>
680void gemm_kernel_i8i8_mrow_ref(size_t          blk_len,
681                               const uint8_t * quant_a_ptr,
682                               const uint8_t * quant_b_data,
683                               const uint8_t * quant_b_zp,
684                               float *         c_ptr,
685                               size_t          count_m,
686                               size_t          count_n,
687                               size_t          k_blks,
688                               size_t          ldc) {
689    int64_t b_blk_stride        = (sizeof(ggml_fp16_t) + blk_len);
690    int64_t b_stride            = k_blks * b_blk_stride;
691    int64_t a_blk_stride        = q8_blk_size(blk_len, true);
692    int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS;
693    int64_t b_ncol_block_stride = b_blk_stride * NB_COLS;
694
695    float output[MB_ROWS * NB_COLS] = { 0 };
696
697    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
698        size_t   nb_real = std::min<size_t>(NB_COLS, count_n - ni);
699        int8_t * b_data  = (int8_t *) quant_b_data + ni * b_stride + NB_COLS * sizeof(ggml_fp16_t);
700
701        int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS;
702
703        for (size_t mi = 0; mi < MB_ROWS; mi++) {
704            for (size_t ci = 0; ci < NB_COLS; ci++) {
705                output[ci + mi * NB_COLS] = 0;
706            }
707        }
708
709        for (size_t ki = 0; ki < k_blks; ki++, a_data += a_nrow_block_stride, b_data += b_ncol_block_stride) {
710            ggml_fp16_t * b_scale_fp16 = (ggml_fp16_t *) (b_data - NB_COLS * sizeof(ggml_fp16_t));
711
712            float * a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS);
713
714            for (size_t mi = 0; mi < MB_ROWS; mi++) {
715                float a_scale = a_scale_row[mi];
716                for (size_t ci = 0; ci < NB_COLS; ci++) {
717                    float   b_scale = ggml_fp16_to_fp32(b_scale_fp16[ci]);
718                    int32_t acc     = 0;
719                    for (size_t bi = 0; bi < blk_len; bi++) {
720                        int8_t a0 = a_data[mi * blk_len + bi];
721                        int8_t b0 = b_data[ci * blk_len + bi];
722                        acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0);
723                    }
724                    output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
725                }
726            }
727        }
728
729        for (size_t mi = 0; mi < MB_ROWS; mi++) {
730            for (size_t ci = 0; ci < nb_real; ci++) {
731                c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
732            }
733        }
734    }
735}
736
737template <size_t MB_ROWS, size_t NB_COLS>
738void gemm_kernel_i8i5_mrow_ref(size_t          blk_len,
739                               const uint8_t * quant_a_ptr,
740                               const uint8_t * quant_b_data,
741                               const uint8_t * quant_b_zp,
742                               float *         c_ptr,
743                               size_t          count_m,
744                               size_t          count_n,
745                               size_t          k_blks,
746                               size_t          ldc) {
747    // blk_len is expected to be 32 for Q5 types
748    // quant_b_zp != nullptr => nrow_block_q5_1<NB_COLS> (has zp)
749    // quant_b_zp == nullptr => nrow_block_q5_0<NB_COLS> (no zp)
750
751    int64_t a_blk_stride        = q8_blk_size(blk_len, true);
752    int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS;
753
754    float output[MB_ROWS * NB_COLS] = { 0 };
755
756    if (quant_b_zp) {
757        // nrow_block_q5_1<NB_COLS>: scales16[NB_COLS] + zp[NB_COLS] + qh[4*NB_COLS] + qs[16*NB_COLS]
758        using blk_type                 = nrow_block_q5_1<NB_COLS>;
759        int64_t    b_ncol_block_stride = sizeof(blk_type);
760        blk_type * quant_b_blk_data    = (blk_type *) quant_b_data;
761
762        for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
763            size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
764
765            int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS;
766
767            for (size_t mi = 0; mi < MB_ROWS; mi++) {
768                for (size_t ci = 0; ci < NB_COLS; ci++) {
769                    output[ci + mi * NB_COLS] = 0;
770                }
771            }
772
773            for (size_t ki = 0; ki < k_blks; ki++, quant_b_blk_data++, a_data += a_nrow_block_stride) {
774                float *   a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS);
775                int16_t * a_sum_row   = (int16_t *) (a_data - sizeof(int16_t) * MB_ROWS);
776
777                for (size_t mi = 0; mi < MB_ROWS; mi++) {
778                    float   a_scale = a_scale_row[mi];
779                    int16_t a_sum   = a_sum_row[mi];
780
781                    for (size_t ci = 0; ci < NB_COLS; ci++) {
782                        float   b_scale  = ggml_fp16_to_fp32(quant_b_blk_data->scales16[ci]);
783                        uint8_t b_zp_val = quant_b_blk_data->zp[ci];
784                        int32_t acc      = a_sum * static_cast<int32_t>(b_zp_val);
785
786                        for (size_t bi = 0; bi < blk_len / 2; bi++) {
787                            int8_t  a0      = a_data[mi * blk_len + 2 * bi];
788                            int8_t  a1      = a_data[mi * blk_len + 2 * bi + 1];
789                            uint8_t qs_byte = quant_b_blk_data->qs[ci * (blk_len / 2) + bi];
790                            int8_t  b0      = static_cast<int8_t>(qs_byte & 0x0F);
791                            int8_t  b1      = static_cast<int8_t>((qs_byte >> 4) & 0x0F);
792
793                            // Extract high bits from qh
794                            // qh is packed as 4 bytes per column (32 bits for 32 elements)
795                            uint8_t qh_byte0 = quant_b_blk_data->qh[ci * 4 + (2 * bi) / 8];
796                            uint8_t qh_byte1 = quant_b_blk_data->qh[ci * 4 + (2 * bi + 1) / 8];
797                            uint8_t h0       = (qh_byte0 >> ((2 * bi) % 8)) & 1;
798                            uint8_t h1       = (qh_byte1 >> ((2 * bi + 1) % 8)) & 1;
799
800                            b0 |= (h0 << 4);
801                            b1 |= (h1 << 4);
802
803                            acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
804                                   static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
805                        }
806                        output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
807                    }
808                }
809            }
810
811            for (size_t mi = 0; mi < MB_ROWS; mi++) {
812                for (size_t ci = 0; ci < nb_real; ci++) {
813                    c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
814                }
815            }
816        }
817    } else {
818        // nrow_block_q5_0<NB_COLS>: scales16[NB_COLS] + qh[4*NB_COLS] + qs[16*NB_COLS]
819        using blk_type                 = nrow_block_q5_0<NB_COLS>;
820        int64_t    b_ncol_block_stride = sizeof(blk_type);
821        blk_type * quant_b_blk_data    = (blk_type *) quant_b_data;
822
823        for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
824            size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
825
826            int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS;
827
828            for (size_t mi = 0; mi < MB_ROWS; mi++) {
829                for (size_t ci = 0; ci < NB_COLS; ci++) {
830                    output[ci + mi * NB_COLS] = 0;
831                }
832            }
833
834            for (size_t ki = 0; ki < k_blks; ki++, quant_b_blk_data++, a_data += a_nrow_block_stride) {
835                float *   a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS);
836                int16_t * a_sum_row   = (int16_t *) (a_data - sizeof(int16_t) * MB_ROWS);
837
838                for (size_t mi = 0; mi < MB_ROWS; mi++) {
839                    float   a_scale = a_scale_row[mi];
840                    int16_t a_sum   = a_sum_row[mi];
841
842                    for (size_t ci = 0; ci < NB_COLS; ci++) {
843                        float   b_scale = ggml_fp16_to_fp32(quant_b_blk_data->scales16[ci]);
844                        // Q5_0 has no zp, use default offset 16 (midpoint of 5-bit unsigned range)
845                        int32_t acc     = a_sum * 16;
846
847                        for (size_t bi = 0; bi < blk_len / 2; bi++) {
848                            int8_t  a0      = a_data[mi * blk_len + 2 * bi];
849                            int8_t  a1      = a_data[mi * blk_len + 2 * bi + 1];
850                            uint8_t qs_byte = quant_b_blk_data->qs[ci * (blk_len / 2) + bi];
851                            int8_t  b0      = static_cast<int8_t>(qs_byte & 0x0F);
852                            int8_t  b1      = static_cast<int8_t>((qs_byte >> 4) & 0x0F);
853
854                            // Extract high bits from qh
855                            uint8_t qh_byte0 = quant_b_blk_data->qh[ci * 4 + (2 * bi) / 8];
856                            uint8_t qh_byte1 = quant_b_blk_data->qh[ci * 4 + (2 * bi + 1) / 8];
857                            uint8_t h0       = (qh_byte0 >> ((2 * bi) % 8)) & 1;
858                            uint8_t h1       = (qh_byte1 >> ((2 * bi + 1) % 8)) & 1;
859
860                            b0 |= (h0 << 4);
861                            b1 |= (h1 << 4);
862
863                            acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
864                                   static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
865                        }
866                        output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
867                    }
868                }
869            }
870
871            for (size_t mi = 0; mi < MB_ROWS; mi++) {
872                for (size_t ci = 0; ci < nb_real; ci++) {
873                    c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
874                }
875            }
876        }
877    }
878}
879
880template <size_t MB_ROWS, size_t NB_COLS>
881void gemm_kernel_i8mxfp4_mrow_ref(size_t          blk_len,
882                                  const uint8_t * quant_a_ptr,
883                                  const uint8_t * quant_b_data,
884                                  const uint8_t * quant_b_zp,
885                                  float *         c_ptr,
886                                  size_t          count_m,
887                                  size_t          count_n,
888                                  size_t          k_blks,
889                                  size_t          ldc) {
890    // blk_len is expected to be 32 (QK_MXFP4)
891    // quant_b_zp is unused for MXFP4 (symmetric quantization)
892    GGML_UNUSED(quant_b_zp);
893
894    int64_t a_blk_stride        = q8_blk_size(blk_len, true);
895    int64_t a_nrow_block_stride = a_blk_stride * MB_ROWS;
896
897    float output[MB_ROWS * NB_COLS] = { 0 };
898
899    using blk_type              = nrow_block_mxfp4<NB_COLS>;
900    blk_type * quant_b_blk_data = (blk_type *) quant_b_data;
901
902    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
903        size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
904
905        int8_t * a_data = (int8_t *) quant_a_ptr + sizeof(float) * MB_ROWS + sizeof(int16_t) * MB_ROWS;
906
907        for (size_t mi = 0; mi < MB_ROWS; mi++) {
908            for (size_t ci = 0; ci < NB_COLS; ci++) {
909                output[ci + mi * NB_COLS] = 0;
910            }
911        }
912
913        for (size_t ki = 0; ki < k_blks; ki++, quant_b_blk_data++, a_data += a_nrow_block_stride) {
914            float *   a_scale_row = (float *) (a_data - sizeof(float) * MB_ROWS - sizeof(int16_t) * MB_ROWS);
915            int16_t * a_sum_row   = (int16_t *) (a_data - sizeof(int16_t) * MB_ROWS);
916
917            for (size_t mi = 0; mi < MB_ROWS; mi++) {
918                float a_scale = a_scale_row[mi];
919
920                for (size_t ci = 0; ci < NB_COLS; ci++) {
921                    float b_scale = GGML_E8M0_TO_FP32_HALF(quant_b_blk_data->e[ci]);
922
923                    // Read 32 sign bits for this column
924                    uint32_t sign_bits;
925                    memcpy(&sign_bits, &quant_b_blk_data->qh[ci * 4], 4);
926
927                    int32_t acc = 0;
928                    for (size_t bi = 0; bi < blk_len / 2; bi++) {
929                        int8_t a0 = a_data[mi * blk_len + 2 * bi];
930                        int8_t a1 = a_data[mi * blk_len + 2 * bi + 1];
931
932                        // qs[ci*16 + bi] stores abs(vals[bi*2]) in low 4 bits
933                        // and abs(vals[bi*2+1]) in high 4 bits
934                        uint8_t qs_byte = quant_b_blk_data->qs[ci * 16 + bi];
935                        int8_t  b_abs0  = static_cast<int8_t>(qs_byte & 0x0F);
936                        int8_t  b_abs1  = static_cast<int8_t>((qs_byte >> 4) & 0x0F);
937
938                        // Extract sign bits: bit (2*bi) for vals[2*bi], bit (2*bi+1) for vals[2*bi+1]
939                        int8_t b0 = (sign_bits >> (2 * bi)) & 1 ? -b_abs0 : b_abs0;
940                        int8_t b1 = (sign_bits >> (2 * bi + 1)) & 1 ? -b_abs1 : b_abs1;
941
942                        acc += static_cast<int32_t>(a0) * static_cast<int32_t>(b0) +
943                               static_cast<int32_t>(a1) * static_cast<int32_t>(b1);
944                    }
945                    output[ci + mi * NB_COLS] += static_cast<float>(acc) * a_scale * b_scale;
946                }
947            }
948        }
949
950        for (size_t mi = 0; mi < MB_ROWS; mi++) {
951            for (size_t ci = 0; ci < nb_real; ci++) {
952                c_ptr[mi * ldc + ci] = output[mi * NB_COLS + ci];
953            }
954        }
955    }
956}
957
958void gemm_kernel_i8i2k_m1(size_t          blk_len,
959                          const uint8_t * quant_a_ptr,
960                          const uint8_t * quant_b_data,
961                          float *         c_ptr,
962                          size_t          count_m,
963                          size_t          count_n,
964                          size_t          k_blks,
965                          size_t          ldc) {
966    constexpr size_t NB_COLS = 32;
967    using blk_type           = nrow_block_q2_k<NB_COLS>;
968
969    int64_t b_ncol_block_stride = sizeof(blk_type) * k_blks;
970
971    for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
972        uint8_t * b_data = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_ncol_block_stride;
973        int8_t *  a_data = (int8_t *) quant_a_ptr;
974        float *   dst_c  = (float *) c_ptr + ni;
975
976        asm volatile(
977            "vsetvli        t0, x0, e16, m1         \n\t"
978            "vxor.vv        v31, v31, v31           \n\t"
979            "mv             s1, %[BK]               \n\t"
980
981            ".align 4                               \n\t"
982            "BLK_LOOP%=:                            \n\t"
983            // load scale A
984            "flw            fa0, (%[A])             \n\t"
985            "addi           %[A], %[A], 4           \n\t"
986
987            "li             t1, 4                   \n\t"
988            "addi           t2, %[B], 512           \n\t"  // B data addr
989            "addi           t3, %[A], 32            \n\t"  // A data addr
990            "addi           s3, %[B], 0             \n\t"
991            "vxor.vv        v30, v29, v29           \n\t"  // tmp result
992
993            "INNER_K_LOOP%=:                        \n\t"
994            "vsetvli        t0, x0, e8, m1          \n\t"
995            "vxor.vv        v2, v2, v2              \n\t"
996            "vxor.vv        v3, v3, v3              \n\t"
997            "vxor.vv        v4, v4, v4              \n\t"
998            "vxor.vv        v5, v5, v5              \n\t"
999            "vxor.vv        v6, v6, v6              \n\t"
1000            "vxor.vv        v28, v28, v28           \n\t"
1001            "vxor.vv        v29, v29, v29           \n\t"
1002
1003            // load scale  B
1004            "vsetvli        t0, x0, e8, m1          \n\t"
1005            "vle8.v         v0, (%[B])              \n\t"
1006            "addi           %[B], %[B], 128         \n\t"
1007
1008            // A data, 1x64@i8
1009            "vsetivli       t0, 16, e8, mf4         \n\t"
1010            "vle8.v         v2, (t3)                \n\t"
1011            "addi           t3, t3, 16              \n\t"
1012
1013            "vsetivli       t0, 16, e8, mf4         \n\t"
1014            "vle8.v         v4, (t3)                \n\t"
1015            "addi           t3, t3, 16              \n\t"
1016
1017            "vsetivli       t0, 16, e8, mf4         \n\t"
1018            "vle8.v         v5, (t3)                \n\t"
1019            "addi           t3, t3, 16              \n\t"
1020
1021            "vsetivli       t0, 16, e8, mf4         \n\t"
1022            "vle8.v         v6, (t3)                \n\t"
1023            "addi           t3, t3, 16              \n\t"
1024
1025            "vsetvli        t0, x0, e64, mf2        \n\t"
1026            "vslideup.vi    v3, v4, 2               \n\t"
1027            "vslideup.vi    v28, v5, 4              \n\t"
1028            "vslideup.vi    v29, v6, 6              \n\t"
1029
1030            // init the accumu to zero
1031            "vsetvli        t0, x0, e16, m1         \n\t"
1032            "vxor.vv        v20, v18, v18           \n\t"
1033            "vxor.vv        v22, v18, v18           \n\t"
1034            "vxor.vv        v24, v18, v18           \n\t"
1035            "vxor.vv        v26, v18, v18           \n\t"
1036
1037            // B data, 32x64@i2
1038            "vsetvli        t0, x0, e8, m1          \n\t"
1039            "vl4r.v         v4, (t2)                \n\t"
1040            "addi           t2, t2, 512             \n\t"
1041            "vand.vi        v8, v4, 0x3             \n\t"  // 0-15
1042            "vsrl.vi        v9, v4, 2               \n\t"
1043            "vsrl.vi        v10, v4, 4              \n\t"
1044            "vsrl.vi        v11, v4, 6              \n\t"  // 48-63
1045            "vand.vi        v9, v9, 0x3             \n\t"  // 16-31
1046            "vand.vi        v10, v10, 0x3           \n\t"  // 32-47
1047
1048            "vand.vi        v12, v5, 0x3            \n\t"  // 0-15
1049            "vsrl.vi        v13, v5, 2              \n\t"
1050            "vsrl.vi        v14, v5, 4              \n\t"
1051            "vsrl.vi        v15, v5, 6              \n\t"  // 48-63
1052            "vand.vi        v13, v13, 0x3           \n\t"  // 16-31
1053            "vand.vi        v14, v14, 0x3           \n\t"  // 32-47
1054
1055            "vand.vi        v16, v6, 0x3            \n\t"  // 0-15
1056            "vsrl.vi        v17, v6, 2              \n\t"
1057            "vsrl.vi        v18, v6, 4              \n\t"
1058            "vsrl.vi        v19, v6, 6              \n\t"  // 48-63
1059            "vand.vi        v17, v17, 0x3           \n\t"  // 16-31
1060            "vand.vi        v18, v18, 0x3           \n\t"  // 32-47
1061
1062            "vand.vi        v4, v7, 0x3             \n\t"  // 0-15
1063            "vsrl.vi        v5, v7, 2               \n\t"
1064            "vsrl.vi        v6, v7, 4               \n\t"
1065            "vsrl.vi        v7, v7, 6               \n\t"  // 48-63
1066            "vand.vi        v5, v5, 0x3             \n\t"  // 16-31
1067            "vand.vi        v6, v6, 0x3             \n\t"  // 32-47
1068
1069            // i2 * i8 vmadot
1070            "vsetvli        t0, x0, e8, m1          \n\t"
1071            "vmadotsu       v20, v2, v8, i8         \n\t"
1072            "vmadotsu       v22, v2, v12, i8        \n\t"
1073            "vmadotsu       v24, v2, v16, i8        \n\t"
1074            "vmadotsu       v26, v2, v4, i8         \n\t"
1075
1076            "vmadotsu       v20, v3, v9, i8         \n\t"
1077            "vmadotsu       v22, v3, v13, i8        \n\t"
1078            "vmadotsu       v24, v3, v17, i8        \n\t"
1079            "vmadotsu       v26, v3, v5, i8         \n\t"
1080
1081            "vmadotsu       v20, v28, v10, i8       \n\t"
1082            "vmadotsu       v22, v28, v14, i8       \n\t"
1083            "vmadotsu       v24, v28, v18, i8       \n\t"
1084            "vmadotsu       v26, v28, v6, i8        \n\t"
1085
1086            "vmadotsu       v20, v29, v11, i8       \n\t"
1087            "vmadotsu       v22, v29, v15, i8       \n\t"
1088            "vmadotsu       v24, v29, v19, i8       \n\t"
1089            "vmadotsu       v26, v29, v7, i8        \n\t"
1090
1091            "vand.vi        v10, v0, 0xf            \n\t"  // scale
1092            "vwadd.vx       v12, v10, x0            \n\t"
1093            "vsetvli        t0, x0, e16, m2         \n\t"
1094            "vwadd.vx       v16, v12, x0            \n\t"
1095
1096            "vsetvli        t0, x0, e32, m1         \n\t"
1097            "vpack.vv       v2, v20, v22, 2         \n\t"
1098            "vpack.vv       v4, v24, v26, 2         \n\t"
1099            "vpack.vv       v6, v2, v4, 3           \n\t"  // 0,1
1100            "vpack.vv       v8, v3, v5, 3           \n\t"  // 2,3
1101
1102            // mul scale
1103            "vmacc.vv       v30, v6, v16            \n\t"
1104            "vmacc.vv       v30, v7, v17            \n\t"
1105            "vmacc.vv       v30, v8, v18            \n\t"
1106            "vmacc.vv       v30, v9, v19            \n\t"
1107
1108            "addi           t1, t1, -1              \n\t"
1109            "bgtz           t1, INNER_K_LOOP%=      \n\t"
1110
1111            // load zp B
1112            "vsetvli        t0, x0, e8, m4          \n\t"
1113            "vle8.v         v4, (s3)                \n\t"
1114            "vsrl.vi        v8, v4, 4               \n\t"  // zp
1115
1116            // asum * zp
1117            "vsetvli        t0, x0, e16, m1         \n\t"
1118            "vxor.vv        v20, v20, v20           \n\t"
1119            "vxor.vv        v22, v22, v22           \n\t"
1120            "vxor.vv        v24, v24, v24           \n\t"
1121            "vxor.vv        v26, v26, v26           \n\t"
1122
1123            "vsetvli        t0, x0, e16, mf4        \n\t"
1124            "vle16.v        v2, (%[A])              \n\t"
1125            "vsetvli        t0, x0, e8, mf4         \n\t"
1126            "vnsrl.wi       v12, v2, 0              \n\t"  // low 8
1127            "vnsra.wi       v13, v2, 8              \n\t"  // high 8
1128
1129            "vsetvli        t0, x0, e32, m1         \n\t"
1130            "vmadotsu       v20, v13, v8, i8        \n\t"
1131            "vmadotsu       v22, v13, v9, i8        \n\t"
1132            "vmadotsu       v24, v13, v10, i8       \n\t"
1133            "vmadotsu       v26, v13, v11, i8       \n\t"
1134
1135            "vsll.vi        v20, v20, 8             \n\t"
1136            "vsll.vi        v22, v22, 8             \n\t"
1137            "vsll.vi        v24, v24, 8             \n\t"
1138            "vsll.vi        v26, v26, 8             \n\t"
1139
1140            "vmadotu        v20, v12, v8, i8        \n\t"
1141            "vmadotu        v22, v12, v9, i8        \n\t"
1142            "vmadotu        v24, v12, v10, i8       \n\t"
1143            "vmadotu        v26, v12, v11, i8       \n\t"
1144
1145            "vpack.vv       v2, v20, v22, 2         \n\t"
1146            "vpack.vv       v4, v24, v26, 2         \n\t"
1147            "vpack.vv       v28, v2, v4, 3          \n\t"
1148
1149            "vsetvli        t0, x0, e16, mf2        \n\t"
1150            "vle16.v        v0, (t2)                \n\t"  // scale16
1151            "addi           t2, t2, 64              \n\t"
1152            "vle16.v        v1, (t2)                \n\t"  // zero16
1153            "vfwcvt.f.f.v   v2, v0                  \n\t"
1154            "vfwcvt.f.f.v   v4, v1                  \n\t"
1155            "vsetvli        t0, x0, e32, m1         \n\t"
1156            "vfcvt.f.x.v    v30, v30                \n\t"
1157            "vfcvt.f.x.v    v28, v28                \n\t"
1158            "addi           %[B], t2, 64            \n\t"
1159            "mv             %[A], t3                \n\t"
1160
1161            "vfmul.vv       v30, v30, v2            \n\t"  // mul scale16
1162            "vfmacc.vv      v30, v28, v4            \n\t"  // + mul zero16
1163            "vfmacc.vf      v31, fa0, v30           \n\t"
1164            "addi           s1, s1, -1              \n\t"
1165            "bgtz           s1, BLK_LOOP%=          \n\t"
1166
1167            // save
1168            "vsetvli        t0, x0, e32, m1         \n\t"
1169            "vse32.v        v31, (%[DST])           \n\t"
1170            : [A] "+r"(a_data), [B] "+r"(b_data)
1171            : [DST] "r"(dst_c), [BK] "r"(k_blks)
1172            : "t0", "t1", "t2", "t3", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12",
1173              "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27",
1174              "v28", "v29", "v30", "v31", "fa0", "t4", "t5", "t6", "s1", "s2", "s3");
1175    }
1176}
1177
1178void gemm_kernel_i8i2k_m4(size_t          blk_len,
1179                          const uint8_t * quant_a_ptr,
1180                          const uint8_t * quant_b_data,
1181                          float *         c_ptr,
1182                          size_t          count_m,
1183                          size_t          count_n,
1184                          size_t          k_blks,
1185                          size_t          ldc) {
1186    constexpr size_t NB_COLS = 32;
1187    using blk_type           = nrow_block_q2_k<NB_COLS>;
1188
1189    int64_t  b_ncol_block_stride = sizeof(blk_type) * k_blks;
1190    _Float16 scale               = 0.0625f;
1191    _Float16 scale_1             = 16.0f;
1192
1193    for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
1194        uint8_t * b_data = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_ncol_block_stride;
1195        int8_t *  a_data = (int8_t *) quant_a_ptr;
1196        float *   dst_c  = (float *) c_ptr + ni;
1197
1198        asm volatile(
1199            "vsetvli        t0, x0, e16, m1         \n\t"
1200            "vxor.vv        v28, v31, v31           \n\t"  // init result
1201            "vxor.vv        v29, v31, v31           \n\t"
1202            "vxor.vv        v30, v31, v31           \n\t"
1203            "vxor.vv        v31, v31, v31           \n\t"
1204            "mv             s1, %[BK]               \n\t"
1205
1206            ".align 4                               \n\t"
1207            "BLK_LOOP%=:                            \n\t"
1208            // load scale A
1209            "flw            fa0, (%[A])             \n\t"
1210            "flw            fa1, 4(%[A])            \n\t"
1211            "flw            fa2, 8(%[A])            \n\t"
1212            "flw            fa3, 12(%[A])           \n\t"
1213            "addi           %[A], %[A], 16          \n\t"
1214
1215            "li             t1, 4                   \n\t"
1216            "addi           t2, %[B], 512           \n\t"  // B data addr
1217            "addi           t3, %[A], 128           \n\t"  // A data addr
1218            "addi           s4, t2, 1024            \n\t"  // scale16 addr
1219            "addi           s4, s4, 1024            \n\t"  // TODO
1220            "addi           s3, %[B], 0             \n\t"
1221
1222            "vsetvli        t0, x0, e16, mf2        \n\t"
1223            "vle16.v        v1, (s4)                \n\t"  // load scale16
1224            "vsetvli        t0, x0, e16, m1         \n\t"
1225            "vpack.vv       v22, v1, v1, 3          \n\t"
1226
1227            "addi           s4, t3, 256             \n\t"  // addr 1
1228            "addi           s5, t3, 512             \n\t"  // addr 2
1229            "addi           s6, t3, 768             \n\t"  // addr 3
1230
1231            // init the accu to 0
1232            "vxor.vv        v24, v24, v24           \n\t"
1233            "vxor.vv        v25, v25, v25           \n\t"
1234            "vxor.vv        v26, v26, v26           \n\t"
1235            "vxor.vv        v27, v27, v27           \n\t"
1236
1237            "INNER_K_LOOP%=:                        \n\t"
1238            // load scale  B
1239            "vsetvli        t0, x0, e8, m1          \n\t"
1240            "vle8.v         v1, (%[B])              \n\t"
1241            "addi           %[B], %[B], 128         \n\t"
1242            "vand.vi        v1, v1, 0xf             \n\t"
1243
1244            "vfwcvt.f.x.v   v20, v1                 \n\t"  // f16 scale B
1245            "vsetvli        t0, x0, e16, m1         \n\t"
1246            "vfmul.vv       v0, v20, v22            \n\t"  // mul scale16
1247            "vfmul.vv       v1, v21, v22            \n\t"  // mul scale16
1248            "vfmul.vf       v0, v0, %[SCALE]        \n\t"  // mul magic
1249            "vfmul.vf       v1, v1, %[SCALE]        \n\t"  // mul magic
1250
1251            // A data, 4x64@i8
1252            "vsetvli        t0, x0, e8, mf2         \n\t"
1253            "vle8.v         v2, (t3)                \n\t"
1254            "addi           t3, t3, 64              \n\t"
1255            "vle8.v         v3, (s4)                \n\t"
1256            "addi           s4, s4, 64              \n\t"
1257            "vle8.v         v4, (s5)                \n\t"
1258            "addi           s5, s5, 64              \n\t"
1259            "vle8.v         v5, (s6)                \n\t"
1260            "addi           s6, s6, 64              \n\t"
1261
1262            // 4x64 => 4x16x4
1263            "vsetvli        t0, x0, e8, m1          \n\t"
1264            "vpack.vv       v6, v2, v3, 1           \n\t"
1265            "vpack.vv       v8, v4, v5, 1           \n\t"
1266            "vpack.vv       v2, v6, v8, 2           \n\t"  // 0, 2
1267
1268            "vpack.vv       v20, v2, v2, 3          \n\t"  // 1
1269            "vor.vv         v23, v21, v21           \n\t"
1270            "vpack.vv       v20, v3, v3, 3          \n\t"  // 3
1271
1272            // B data, 32x64@i2
1273            "vsetvli        t0, x0, e8, m1          \n\t"
1274            "vl4r.v         v4, (t2)                \n\t"
1275            "addi           t2, t2, 512             \n\t"
1276            "vand.vi        v8, v4, 0x3             \n\t"  // 0-15
1277            "vsrl.vi        v9, v4, 2               \n\t"
1278            "vsrl.vi        v10, v4, 4              \n\t"
1279            "vsrl.vi        v11, v4, 6              \n\t"  // 48-63
1280            "vand.vi        v9, v9, 0x3             \n\t"  // 16-31
1281            "vand.vi        v10, v10, 0x3           \n\t"  // 32-47
1282
1283            "vand.vi        v12, v5, 0x3            \n\t"  // 0-15
1284            "vsrl.vi        v13, v5, 2              \n\t"
1285            "vsrl.vi        v14, v5, 4              \n\t"
1286            "vsrl.vi        v15, v5, 6              \n\t"  // 48-63
1287            "vand.vi        v13, v13, 0x3           \n\t"  // 16-31
1288            "vand.vi        v14, v14, 0x3           \n\t"  // 32-47
1289
1290            "vand.vi        v16, v6, 0x3            \n\t"  // 0-15
1291            "vsrl.vi        v17, v6, 2              \n\t"
1292            "vsrl.vi        v18, v6, 4              \n\t"
1293            "vsrl.vi        v19, v6, 6              \n\t"  // 48-63
1294            "vand.vi        v17, v17, 0x3           \n\t"  // 16-31
1295            "vand.vi        v18, v18, 0x3           \n\t"  // 32-47
1296
1297            "vand.vi        v4, v7, 0x3             \n\t"  // 0-15
1298            "vsrl.vi        v5, v7, 2               \n\t"
1299            "vsrl.vi        v6, v7, 4               \n\t"
1300            "vsrl.vi        v7, v7, 6               \n\t"  // 48-63
1301            "vand.vi        v5, v5, 0x3             \n\t"  // 16-31
1302            "vand.vi        v6, v6, 0x3             \n\t"  // 32-47
1303
1304            // i2 * i8 vmadot
1305            "vsetvli        t0, x0, e8, m1          \n\t"
1306            "vmadotsu.hp    v24, v2, v8, v0, 0, i8  \n\t"
1307            "vmadotsu.hp    v25, v2, v12, v0, 1, i8 \n\t"
1308            "vmadotsu.hp    v26, v2, v16, v0, 2, i8 \n\t"
1309            "vmadotsu.hp    v27, v2, v4, v0, 3, i8  \n\t"
1310
1311            "vmadotsu.hp    v24, v23, v9, v0, 4, i8 \n\t"
1312            "vmadotsu.hp    v25, v23, v13, v0, 5, i8\n\t"
1313            "vmadotsu.hp    v26, v23, v17, v0, 6, i8\n\t"
1314            "vmadotsu.hp    v27, v23, v5, v0, 7, i8 \n\t"
1315
1316            "vmadotsu.hp    v24, v3, v10, v1, 0, i8 \n\t"
1317            "vmadotsu.hp    v25, v3, v14, v1, 1, i8 \n\t"
1318            "vmadotsu.hp    v26, v3, v18, v1, 2, i8 \n\t"
1319            "vmadotsu.hp    v27, v3, v6, v1, 3, i8  \n\t"
1320
1321            "vmadotsu.hp    v24, v21, v11, v1, 4, i8\n\t"
1322            "vmadotsu.hp    v25, v21, v15, v1, 5, i8\n\t"
1323            "vmadotsu.hp    v26, v21, v19, v1, 6, i8\n\t"
1324            "vmadotsu.hp    v27, v21, v7, v1, 7, i8 \n\t"
1325
1326            "addi           t1, t1, -1              \n\t"
1327            "bgtz           t1, INNER_K_LOOP%=      \n\t"
1328
1329            "vsetvli        t0, x0, e16, m1         \n\t"
1330            "vpack.vv       v2, v24, v25, 1         \n\t"
1331            "vpack.vv       v4, v26, v27, 1         \n\t"
1332            "vpack.vv       v6, v2, v4, 2           \n\t"  // 0,1,2,3
1333
1334            "vxor.vv        v18, v18, v18           \n\t"
1335            "vxor.vv        v20, v20, v20           \n\t"
1336            "vxor.vv        v22, v22, v22           \n\t"
1337            "vxor.vv        v24, v24, v24           \n\t"
1338            // load zp B, 16x8x4@int4
1339            "vsetvli        t0, x0, e8, m4          \n\t"
1340            "vle8.v         v0, (s3)                \n\t"
1341            "vsrl.vi        v0, v0, 4               \n\t"  // zp
1342
1343            // 4x16@int16
1344            "vsetvli        t0, x0, e16, m1         \n\t"  // a sum
1345            "vle16.v         v12, (%[A])             \n\t"
1346            "vsetvli        t0, x0, e8, m1          \n\t"
1347            "vnsrl.wi       v10, v12, 0             \n\t"  // low 8
1348            "vnsra.wi       v11, v12, 8             \n\t"  // high 8
1349
1350            // asum * zp
1351            "vsetvli        t0, x0, e32, m1          \n\t"
1352            "vmadotsu       v18, v11, v0, i8        \n\t"
1353            "vmadotsu       v20, v11, v1, i8        \n\t"
1354            "vmadotsu       v22, v11, v2, i8        \n\t"
1355            "vmadotsu       v24, v11, v3, i8        \n\t"
1356            "vsll.vi        v18, v18, 8             \n\t"
1357            "vsll.vi        v20, v20, 8             \n\t"
1358            "vsll.vi        v22, v22, 8             \n\t"
1359            "vsll.vi        v24, v24, 8             \n\t"
1360            "vmadotu        v18, v10, v0, i8        \n\t"
1361            "vmadotu        v20, v10, v1, i8        \n\t"
1362            "vmadotu        v22, v10, v2, i8        \n\t"
1363            "vmadotu        v24, v10, v3, i8        \n\t"
1364
1365            "vpack.vv       v10, v18, v20, 2        \n\t"
1366            "vpack.vv       v12, v22, v24, 2        \n\t"
1367            "vpack.vv       v14, v10, v12, 3        \n\t"
1368            "vpack.vv       v16, v11, v13, 3        \n\t"
1369
1370            "vsetvli        t0, x0, e16, mf2        \n\t"
1371            "addi           t2, t2, 64              \n\t"
1372            "vle16.v        v20, (t2)               \n\t"  // zero16
1373            "vfwcvt.f.f.v   v22, v20                \n\t"
1374
1375            // mul 1/magic
1376            "vsetvli        t0, x0, e16, m1         \n\t"
1377            "vfwmul.vf      v0, v6, %[SCALE_1]      \n\t"
1378            "vfwmul.vf      v2, v7, %[SCALE_1]      \n\t"
1379
1380            "vsetvli        t0, x0, e32, m1         \n\t"
1381            "vfcvt.f.x.v    v14, v14                \n\t"
1382            "vfcvt.f.x.v    v15, v15                \n\t"
1383            "vfcvt.f.x.v    v16, v16                \n\t"
1384            "vfcvt.f.x.v    v17, v17                \n\t"
1385
1386            "addi           %[B], t2, 64            \n\t"
1387            "mv             %[A], s6                \n\t"
1388
1389            "vfmacc.vv      v0, v14, v22            \n\t"  // + mul zero16
1390            "vfmacc.vv      v1, v15, v22            \n\t"
1391            "vfmacc.vv      v2, v16, v22            \n\t"
1392            "vfmacc.vv      v3, v17, v22            \n\t"
1393
1394            "vfmacc.vf      v28, fa0, v0            \n\t"  // mul a scale
1395            "vfmacc.vf      v29, fa1, v1            \n\t"
1396            "vfmacc.vf      v30, fa2, v2            \n\t"
1397            "vfmacc.vf      v31, fa3, v3            \n\t"
1398
1399            "addi           s1, s1, -1              \n\t"
1400            "bgtz           s1, BLK_LOOP%=          \n\t"
1401
1402            // save
1403            "vsetvli        t0, x0, e32, m1         \n\t"
1404            "add            t1, %[LDC], %[DST]      \n\t"
1405            "vse32.v        v28, (%[DST])           \n\t"
1406            "vse32.v        v29, (t1)               \n\t"
1407            "add            t1, t1, %[LDC]          \n\t"
1408            "vse32.v        v30, (t1)               \n\t"
1409            "add            t1, t1, %[LDC]          \n\t"
1410            "vse32.v        v31, (t1)               \n\t"
1411            : [A] "+r"(a_data), [B] "+r"(b_data)
1412            : [DST] "r"(dst_c), [BK] "r"(k_blks), [LDC] "r"(ldc * 4), [SCALE] "f"(scale), [SCALE_1] "f"(scale_1)
1413            : "t0", "t1", "t2", "t3", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12",
1414              "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27",
1415              "v28", "v29", "v30", "v31", "fa0", "t4", "t5", "t6", "s1", "s2", "s3", "s4", "s5", "s6");
1416    }
1417}
1418
1419void gemm_kernel_i8i3k_m1(size_t          blk_len,
1420                          const uint8_t * quant_a_ptr,
1421                          const uint8_t * quant_b_data,
1422                          float *         c_ptr,
1423                          size_t          count_m,
1424                          size_t          count_n,
1425                          size_t          k_blks,
1426                          size_t          ldc) {
1427    constexpr size_t NB_COLS = 32;  //only support 32 in ASM
1428    using blk_type           = nrow_block_q3_k<NB_COLS>;
1429
1430    const blk_type * b_base = reinterpret_cast<const blk_type *>(quant_b_data);
1431
1432    int64_t a_blk_stride        = q8k_blk_size(256);
1433    int64_t a_nrow_block_stride = a_blk_stride;
1434    int64_t b_ncol_block_stride = sizeof(blk_type);
1435
1436    // Constants used by q3_k scaling in HP branch:
1437    // - k_q3k_scale_step: per-nibble scale factor (1/16).
1438    // - k_a_scale_post_mul: A_scale needs an extra *16 at the end (pairs with 1/16 above).
1439    const _Float16 k_q3k_scale_step   = (_Float16) 0.0625f;  // 1 / 16
1440    const float    k_a_scale_post_mul = 16.0f;
1441
1442    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
1443        size_t           nb_real          = std::min<size_t>(NB_COLS, count_n - ni);
1444        const blk_type * quant_b_blk_data = b_base + (ni / NB_COLS) * k_blks;
1445#if 0
1446        //------------------------------------------------------------------------------
1447        // A format
1448        // Ascale   fp32 * 1    32bit
1449        // Asum     int16 * 16  256bit
1450        // A M1K256 int8        2048bit
1451        //------------------------------------------------------------------------------
1452        // B format
1453        // B_scl    uint8*N32*16    4096bit
1454        // B_Hmask  N32K16*16 1bit  8192bit
1455        // B_Qs     N32K16*16 2bit  16384bit
1456        // B scl16  fp16 * N32      512bit;
1457        //------------------------------------------------------------------------------
1458        //bias always be nullptr
1459        __asm__ volatile(
1460            // t2 = k_blks (each is K256 superblock)
1461            "mv           t2, %[KBLKS]            \n\t"
1462            // t3 = 256/64 = 4 (K64 iterations per superblock)
1463            "li           t3, 4                   \n\t"
1464            "mv           s2, %[pA]               \n\t"  // s2 = pASCL
1465            "addi         s3, %[pA], 4+32         \n\t"  // s3 = pAData, (pA+AScl+ASum)
1466
1467            // B block layout for nrow_block_q3_k<32>:
1468            // scales: 512B, hmask: 1024B, qs: 2048B, scales16: 64B
1469            "addi         s5, %[pB], 32*16        \n\t"  // s5 = pB_hmask
1470            "mv           s4, %[pB]               \n\t"  // s4 = pB_scales
1471            "addi         s6, s5, 1024            \n\t"  // s6 = pB_qs
1472            "mv           s7, %[pB]               \n\t"  // s7 = pB_base
1473
1474            "vsetvli      t0, x0, e32, m1         \n\t"
1475            "vxor.vv      v31, v0, v0             \n\t"  // clear acc
1476            "vxor.vv      v30, v0, v0             \n\t"  // clear acc of K256
1477
1478            // ordinary vmadot: vle*10 vecIns*78 vmadot*16
1479            ".align 4                             \n\t"
1480            "BLK_LPST%=:                          \n\t"
1481            "K64_LPST%=:                          \n\t"
1482
1483            // K0-15
1484            // load B scales (32 bytes per K16, 16 times => 512B)
1485            "vsetvli      t0, x0, e8, m1          \n\t"
1486            "vle8.v       v2, (s4)                \n\t"
1487            "addi         s4, s4, 128             \n\t"
1488
1489            // load B qs chunk (128B per K16, 16 times => 2048B)
1490            "vle8.v       v4, (s6)                \n\t"
1491            "addi         s6, s6, 128             \n\t"
1492            "vle8.v       v5, (s6)                \n\t"
1493            "addi         s6, s6, 128             \n\t"
1494            "vle8.v       v6, (s6)                \n\t"
1495            "addi         s6, s6, 128             \n\t"
1496            "vle8.v       v7, (s6)                \n\t"
1497            "addi         s6, s6, 128             \n\t"
1498
1499            // load B hmask chunk (64B per K16, 16 times => 1024B)
1500            "vsetvli      t0, x0, e8, mf2         \n\t"
1501            "vle8.v       v0, (s5)                \n\t"
1502            "addi         s5, s5, 64              \n\t"
1503
1504            // load A data (16 bytes per K16, 16 times => 256B)
1505            "vsetvli      t0, x0, e8, mf2         \n\t"
1506            "vle8.v       v1, (s3)                \n\t"
1507            "addi         s3, s3, 64              \n\t"
1508
1509            // unpack 2-bit qs + hmask -> signed values
1510            "vsetvli      t0, x0, e8, m1          \n\t"
1511            "vnot.v       v0, v0                  \n\t"
1512            "vand.vi      v12, v4, 0x3            \n\t"
1513            "vand.vi      v13, v5, 0x3            \n\t"
1514            "vand.vi      v14, v6, 0x3            \n\t"
1515            "vand.vi      v15, v7, 0x3            \n\t"
1516
1517            "vsetvli      t0, x0, e8, m4          \n\t"
1518            "vadd.vi      v12, v12, -4, v0.t      \n\t"
1519
1520            "vsetvli      t0, x0, e32, m1         \n\t"
1521            "vxor.vv      v16, v16, v16           \n\t"
1522            "vxor.vv      v18, v16, v16           \n\t"
1523            "vxor.vv      v20, v16, v16           \n\t"
1524            "vxor.vv      v22, v16, v16           \n\t"
1525
1526            "vmadot       v16, v1, v12, i8        \n\t"
1527            "vmadot       v18, v1, v13, i8        \n\t"
1528            "vmadot       v20, v1, v14, i8        \n\t"
1529            "vmadot       v22, v1, v15, i8        \n\t"
1530
1531            "vsetvli      t0, x0, e16, m1         \n\t"
1532            "vpack.vv     v24, v16, v18, 2        \n\t"
1533            "vpack.vv     v26, v20, v22, 2        \n\t"
1534            "vpack.vv     v16, v24, v26, 3        \n\t"  // N0-N31 in v16
1535
1536            // apply B int8 scales (-32 bias has been applyed)
1537            "vsetvli      t0, x0, e8, mf4         \n\t"
1538            "vwadd.vx     v18, v2, x0             \n\t"  // int8 -> int16
1539
1540            "vsetvli      t0, x0, e16, mf2        \n\t"
1541            "vwadd.vx     v19, v18, x0            \n\t"  // int8 -> int16
1542
1543            // static_cast<int32_t>(qsum) * b_scale;
1544            "vsetvli      t0, x0, e32, m1         \n\t"
1545            "vmacc.vv     v30, v16, v19           \n\t"
1546
1547            //K16-31
1548            // load B scales (32 bytes per K16, 16 times => 512B)
1549            "vsetvli      t0, x0, e64, m1         \n\t"
1550            "vslidedown.vi  v2, v2, 4             \n\t"
1551
1552            // load B hmask chunk (64B per K16, 16 times => 1024B)
1553            "vsetvli      t0, x0, e8, mf2         \n\t"
1554            "vle8.v       v0, (s5)                \n\t"
1555            "addi         s5, s5, 64              \n\t"
1556
1557            // load A data (16 bytes per K16, 16 times => 256B)
1558            "vsetvli      t0, x0, e64, mf2        \n\t"
1559            "vslidedown.vi  v1, v1, 2             \n\t"
1560
1561            // unpack 2-bit qs + hmask -> signed values
1562            "vsetvli      t0, x0, e8, m1          \n\t"
1563            "vsll.vi      v8, v4, 4               \n\t"
1564            "vsll.vi      v9, v5, 4               \n\t"
1565            "vsll.vi      v10, v6, 4              \n\t"
1566            "vsll.vi      v11, v7, 4              \n\t"
1567            "vnot.v       v0, v0                  \n\t"
1568
1569            "vsrl.vi      v12, v8, 6              \n\t"
1570            "vsrl.vi      v13, v9, 6              \n\t"
1571            "vsrl.vi      v14, v10, 6             \n\t"
1572            "vsrl.vi      v15, v11, 6             \n\t"
1573
1574            "vsetvli      t0, x0, e8, m4          \n\t"
1575            "vadd.vi      v12, v12, -4, v0.t      \n\t"
1576
1577            "vsetvli      t0, x0, e32, m1         \n\t"
1578            "vxor.vv      v16, v16, v16           \n\t"
1579            "vxor.vv      v18, v16, v16           \n\t"
1580            "vxor.vv      v20, v16, v16           \n\t"
1581            "vxor.vv      v22, v16, v16           \n\t"
1582
1583            "vmadot       v16, v1, v12, i8        \n\t"
1584            "vmadot       v18, v1, v13, i8        \n\t"
1585            "vmadot       v20, v1, v14, i8        \n\t"
1586            "vmadot       v22, v1, v15, i8        \n\t"
1587
1588            "vsetvli      t0, x0, e16, m1         \n\t"
1589            "vpack.vv     v24, v16, v18, 2        \n\t"
1590            "vpack.vv     v26, v20, v22, 2        \n\t"
1591            "vpack.vv     v16, v24, v26, 3        \n\t"  // N0-N31 in v16
1592
1593            // apply B int8 scales (-32 bias has been applyed)
1594            "vsetvli      t0, x0, e8, mf4         \n\t"
1595            "vwadd.vx     v18, v2, x0              \n\t"  // int8 -> int16
1596
1597            "vsetvli      t0, x0, e16, mf2        \n\t"
1598            "vwadd.vx     v19, v18, x0              \n\t"  // int8 -> int16
1599
1600            // static_cast<int32_t>(qsum) * b_scale;
1601            "vsetvli      t0, x0, e32, m1         \n\t"
1602            "vmacc.vv     v30, v16, v19            \n\t"
1603
1604            //K32-47
1605            // load B scales (32 bytes per K16, 16 times => 512B)
1606            "vsetvli      t0, x0, e64, m1         \n\t"
1607            "vslidedown.vi  v2, v2, 4             \n\t"
1608
1609            // load B hmask chunk (64B per K16, 16 times => 1024B)
1610            "vsetvli      t0, x0, e8, mf2         \n\t"
1611            "vle8.v       v0, (s5)                \n\t"
1612            "addi         s5, s5, 64              \n\t"
1613
1614            // load A data (16 bytes per K16, 16 times => 256B)
1615            "vsetvli      t0, x0, e64, mf2        \n\t"
1616            "vslidedown.vi  v1, v1, 2             \n\t"
1617
1618            // unpack 2-bit qs + hmask -> signed values
1619            "vsetvli      t0, x0, e8, m1          \n\t"
1620            "vsll.vi      v8, v4, 2               \n\t"
1621            "vsll.vi      v9, v5, 2               \n\t"
1622            "vsll.vi      v10, v6, 2              \n\t"
1623            "vsll.vi      v11, v7, 2              \n\t"
1624            "vnot.v       v0, v0                  \n\t"
1625
1626            "vsrl.vi      v12, v8, 6              \n\t"
1627            "vsrl.vi      v13, v9, 6              \n\t"
1628            "vsrl.vi      v14, v10, 6             \n\t"
1629            "vsrl.vi      v15, v11, 6             \n\t"
1630
1631            "vsetvli      t0, x0, e8, m4          \n\t"
1632            "vadd.vi      v12, v12, -4, v0.t      \n\t"
1633
1634            "vsetvli      t0, x0, e32, m1         \n\t"
1635            "vxor.vv      v16, v16, v16           \n\t"
1636            "vxor.vv      v18, v16, v16           \n\t"
1637            "vxor.vv      v20, v16, v16           \n\t"
1638            "vxor.vv      v22, v16, v16           \n\t"
1639
1640            "vmadot       v16, v1, v12, i8        \n\t"
1641            "vmadot       v18, v1, v13, i8        \n\t"
1642            "vmadot       v20, v1, v14, i8        \n\t"
1643            "vmadot       v22, v1, v15, i8        \n\t"
1644
1645            "vsetvli      t0, x0, e16, m1         \n\t"
1646            "vpack.vv     v24, v16, v18, 2        \n\t"
1647            "vpack.vv     v26, v20, v22, 2        \n\t"
1648            "vpack.vv     v16, v24, v26, 3        \n\t"
1649
1650            // apply B int8 scales (-32 bias has been applyed)
1651            "vsetvli      t0, x0, e8, mf4         \n\t"
1652            "vwadd.vx     v18, v2, x0              \n\t"  // int8 -> int16
1653
1654            "vsetvli      t0, x0, e16, mf2        \n\t"
1655            "vwadd.vx     v19, v18, x0              \n\t"  // int8 -> int16
1656
1657            // static_cast<int32_t>(qsum) * b_scale;
1658            "vsetvli      t0, x0, e32, m1         \n\t"
1659            "vmacc.vv     v30, v16, v19            \n\t"
1660
1661            // K48-63
1662            // load B scales (32 bytes per K16, 16 times => 512B)
1663            "vsetvli      t0, x0, e64, m1         \n\t"
1664            "vslidedown.vi  v2, v2, 4             \n\t"
1665
1666            // load B hmask chunk (64B per K16, 16 times => 1024B)
1667            "vsetvli      t0, x0, e8, mf2         \n\t"
1668            "vle8.v       v0, (s5)                \n\t"
1669            "addi         s5, s5, 64              \n\t"
1670
1671            // load A data (16 bytes per K16, 16 times => 256B)
1672            "vsetvli      t0, x0, e64, mf2        \n\t"
1673            "vslidedown.vi  v1, v1, 2             \n\t"
1674
1675            "vsetvli      t0, x0, e8, m1          \n\t"
1676            "vnot.v       v0, v0                  \n\t"
1677            "vsrl.vi      v12, v4, 6              \n\t"
1678            "vsrl.vi      v13, v5, 6              \n\t"
1679            "vsrl.vi      v14, v6, 6              \n\t"
1680            "vsrl.vi      v15, v7, 6              \n\t"
1681
1682            "vsetvli      t0, x0, e8, m4          \n\t"
1683            "vadd.vi      v12, v12, -4, v0.t      \n\t"
1684
1685            "vsetvli      t0, x0, e32, m1         \n\t"
1686            "vxor.vv      v16, v16, v16           \n\t"
1687            "vxor.vv      v18, v16, v16           \n\t"
1688            "vxor.vv      v20, v16, v16           \n\t"
1689            "vxor.vv      v22, v16, v16           \n\t"
1690
1691            "vmadot       v16, v1, v12, i8        \n\t"
1692            "vmadot       v18, v1, v13, i8        \n\t"
1693            "vmadot       v20, v1, v14, i8        \n\t"
1694            "vmadot       v22, v1, v15, i8        \n\t"
1695
1696            "vsetvli      t0, x0, e16, m1         \n\t"
1697            "vpack.vv     v24, v16, v18, 2        \n\t"
1698            "vpack.vv     v26, v20, v22, 2        \n\t"
1699            "vpack.vv     v16, v24, v26, 3        \n\t"
1700
1701            // apply B int8 scales (-32 bias has been applyed)
1702            "vsetvli      t0, x0, e8, mf4         \n\t"
1703            "vwadd.vx     v18, v2, x0             \n\t"  // int8 -> int16
1704
1705            "vsetvli      t0, x0, e16, mf2        \n\t"
1706            "vwadd.vx     v19, v18, x0            \n\t"  // int8 -> int16
1707
1708            // static_cast<int32_t>(qsum) * b_scale;
1709            "vsetvli      t0, x0, e32, m1         \n\t"
1710            "vmacc.vv     v30, v16, v19           \n\t"
1711
1712            "addi         t3, t3, -1              \n\t"
1713            "bgtz         t3, K64_LPST%=          \n\t"
1714            "K64_LPND%=:                          \n\t"
1715
1716            // load A scale (fp32) and advance A to next superblock
1717            "flw          f0, (s2)                \n\t"
1718            "addi         s2, s2, 4+32+256        \n\t"
1719            "add          t4, s7, %[B_STR]        \n\t"  // t4 = next B blk base
1720            "addi         s3, s2, 4+32            \n\t"
1721
1722            // load B scales16[32] (fp16) at end of qs region
1723            "vsetvli      t0, x0, e16, mf2        \n\t"
1724            "vle16.v      v2, (s6)                \n\t"
1725
1726            // pointer modify
1727            "addi         s5, t4, 32*16           \n\t"
1728            "mv           s4, t4                  \n\t"
1729            "addi         s6, s5, 32*32           \n\t"
1730            "addi         s7, t4, 0               \n\t"
1731
1732            // b_scale fp16 -> fp32
1733            "vsetvli      t0, x0, e16, mf2        \n\t"
1734            "vfwcvt.f.f.v v24, v2                 \n\t"
1735
1736            // a_scale * b_scale;
1737            "vsetvli      t0, x0, e32, m1         \n\t"
1738            "vfcvt.f.x.v v26, v30                 \n\t"
1739            "vfmul.vf     v1, v24, f0             \n\t"
1740            "vsetvli      t0, x0, e32, m1         \n\t"
1741            // static_cast<float>(qsum) * a_scale * b_scale;
1742            "vfmacc.vv    v31, v1, v26            \n\t"
1743
1744            // next K-superblock
1745            "addi         t2, t2, -1              \n\t"
1746            "vxor.vv      v30, v0, v0             \n\t"  // clear acc of K256
1747            "li           t3, 4                   \n\t"
1748            "bgtz         t2, BLK_LPST%=          \n\t"
1749
1750            "BLK_LPND%=:                          \n\t"
1751            "vsetvli      t0, %[NBLKS], e32, m1   \n\t"
1752            "vse32.v      v31, (%[pC])            \n\t"
1753            "FUNC_END%=:                          \n\t"
1754
1755            :
1756            : [KBLKS] "r"(k_blks), [NBLKS] "r"(nb_real), [pA] "r"(quant_a_ptr), [pB] "r"(quant_b_blk_data),
1757              [pC] "r"(c_ptr), [B_STR] "r"(b_ncol_block_stride)
1758            : "cc", "memory", "t0", "t2", "t3", "t4", "t5", "f0", "s2", "s3", "s4", "s5", "s6", "s7");
1759#else
1760
1761        __asm__ volatile(
1762            // =========================
1763            // Kernel overview (M1 x N32)
1764            // =========================
1765            // Process one output row (M=1) and 32 columns (N=32) per call.
1766            //
1767            // Loop structure:
1768            //   - Outer loop: K superblocks of size K=256 (k_blks times)
1769            //   - Each K256 superblock is broken into 4 x K64
1770            //   - Each K64 is processed as 4 x K16 "sub-blocks" (via unpack+dot)
1771            //
1772            // Data layout (high level):
1773            //   A (q8k K=256, per superblock):
1774            //     [ fp32 a_scale ][ int16 a_sum[16] ][ int8 a_qs[256] ]
1775            //   B (nrow_block_q3_k<32>, per superblock):
1776            //     [ int8  scales[32*16] ][ hmask[1024] ][ qs[2048] ][ fp16 scales16[32] ]
1777            //
1778            // Registers/pointers:
1779            //   s2: pA (points at A superblock header; used to load fp32 a_scale)
1780            //   s3: pA_qs (points at A int8 data within the current superblock)
1781            //   s4: pB_scales (points at B int8 per-K16 scales)
1782            //   s5: pB_hmask (points at B sign mask area)
1783            //   s6: pB_qs (points at B 2-bit packed qs area)
1784            //   s8: pB_scales16 (points at B fp16 scales16[32] at the end of block)
1785            //   s7: pB_base (base pointer to current B block; used for block-to-block stride)
1786
1787            // t2 = number of K256 superblocks
1788            "mv           t2, %[KBLKS]            \n\t"
1789            // t3 = number of K64 chunks per K256 superblock (256 / 64)
1790            "li           t3, 4                   \n\t"
1791
1792            // A pointers
1793            "mv           s2, %[pA]               \n\t"  // s2 = pA_superblock (a_scale at +0)
1794            "addi         s3, %[pA], 4+32         \n\t"  // s3 = pA_qs (skip a_scale + a_sum[16])
1795
1796            // B pointers for nrow_block_q3_k<32>
1797            "addi         s5, %[pB], 32*16        \n\t"  // s5 = pB_hmask  (skip scales[32*16])
1798            "mv           s4, %[pB]               \n\t"  // s4 = pB_scales
1799            "addi         s6, s5, 1024            \n\t"  // s6 = pB_qs     (skip hmask)
1800            // scales16 is at the end of the block: qs(2048) after hmask
1801            "addi         s8, s6, 1024            \n\t"
1802            "addi         s8, s8, 1024            \n\t"  // s8 = pB_scales16 (fp16 scales16[32])
1803            "mv           s7, %[pB]               \n\t"  // s7 = pB_base (for next-block address calc)
1804
1805            // v31: final FP32 accumulator for N=32
1806            "vsetvli      t0, x0, e32, m1         \n\t"
1807            "vxor.vv      v31, v0, v0             \n\t"
1808
1809            // ---- Preload B scales16[32] and build FP16 scale vector used by vmadot.hp ----
1810            "vsetvli      t0, x0, e16, mf2        \n\t"
1811            "vle16.v      v1, (s8)                \n\t"  // load fp16 scales16[32]
1812            "vsetvli      t0, x0, e16, m1         \n\t"
1813            "vpack.vv     v26, v1, v1, 3          \n\t"  // broadcast/pack to match lanes
1814            "vmv.v.v      v17, v26                \n\t"
1815            "vsetvli      t0, x0, e16, m1         \n\t"
1816            "vfmul.vf     v30, v17, %[q3_step]    \n\t"  // v30 = scales16 * (1/16)
1817
1818            // v24-v27: fp16 partial accumulators for a K64 chunk (vmadot.hp outputs)
1819            "vsetvli      t0, x0, e32, m1         \n\t"
1820            "vxor.vv      v24, v16, v16           \n\t"
1821            "vxor.vv      v25, v16, v16           \n\t"
1822            "vxor.vv      v26, v16, v16           \n\t"
1823            "vxor.vv      v27, v16, v16           \n\t"
1824
1825            // HP vmadot: vle*10 vecIns*38 vmadot.hp*16
1826            ".align 4                             \n\t"
1827            "BLK_LPST%=:                          \n\t"  // loop over K256 superblocks
1828            "K64_LPST%=:                          \n\t"  // loop over 4 x K64 chunks
1829
1830            // ------------------------------------------------------------
1831            // K0-15: load B scales + {hmask, qs} + A data; unpack and dot
1832            // ------------------------------------------------------------
1833            "vsetvli      t0, x0, e8, m1          \n\t"
1834            "vle8.v       v2, (s4)                \n\t"  // B int8 scales for this K16
1835            "addi         s4, s4, 128             \n\t"
1836
1837            "vle8.v       v4, (s6)                \n\t"
1838            "addi         s6, s6, 128             \n\t"
1839            "vle8.v       v5, (s6)                \n\t"
1840            "addi         s6, s6, 128             \n\t"
1841            "vle8.v       v6, (s6)                \n\t"
1842            "addi         s6, s6, 128             \n\t"
1843            "vle8.v       v7, (s6)                \n\t"
1844            "addi         s6, s6, 128             \n\t"
1845
1846            "vsetvli      t0, x0, e8, mf2         \n\t"
1847            "vle8.v       v0, (s5)                \n\t"  // B hmask for this K16
1848            "addi         s5, s5, 64              \n\t"
1849
1850            "vsetvli      t0, x0, e8, mf2         \n\t"
1851            "vle8.v       v3, (s3)                \n\t"  // A int8 data for this K16
1852            "addi         s3, s3, 64              \n\t"
1853
1854            // Convert B int8 scales to FP16 and apply scales16*(1/16)
1855            "vsetvli      t0, x0, e8, m1          \n\t"
1856            "vfwcvt.f.x.v v28, v2                 \n\t"  // int8 -> fp16
1857            "vsetvli      t0, x0, e16, m1         \n\t"
1858            "vfmul.vv     v1, v28, v30            \n\t"  // v1: FP16 scale vector for vmadot.hp
1859            "vfmul.vv     v29, v29, v30           \n\t"
1860
1861            // Unpack B 2-bit qs + hmask -> signed int8 in v12..v15
1862            "vsetvli      t0, x0, e8, m1          \n\t"
1863            "vnot.v       v0, v0                  \n\t"
1864            "vand.vi      v12, v4, 0x3            \n\t"
1865            "vand.vi      v13, v5, 0x3            \n\t"
1866            "vand.vi      v14, v6, 0x3            \n\t"
1867            "vand.vi      v15, v7, 0x3            \n\t"
1868            "vsetvli      t0, x0, e8, m4          \n\t"
1869            "vadd.vi      v12, v12, -4, v0.t      \n\t"
1870
1871            // (Next K16 unpack path uses a fresh hmask load)
1872            "vsetvli      t0, x0, e8, mf2         \n\t"
1873            "vle8.v       v0, (s5)                \n\t"
1874            "addi         s5, s5, 64              \n\t"
1875
1876            // Prepare another group from packed qs (bit shifts) + apply sign from hmask
1877            "vsetvli      t0, x0, e8, m1          \n\t"
1878            "vsll.vi      v8, v4, 4               \n\t"
1879            "vsll.vi      v9, v5, 4               \n\t"
1880            "vsll.vi      v10, v6, 4              \n\t"
1881            "vsll.vi      v11, v7, 4              \n\t"
1882            "vsrl.vi      v16, v8, 6              \n\t"
1883            "vsrl.vi      v17, v9, 6              \n\t"
1884            "vnot.v       v0, v0                  \n\t"
1885            "vsrl.vi      v18, v10, 6             \n\t"
1886            "vsrl.vi      v19, v11, 6             \n\t"
1887            "vsetvli      t0, x0, e8, m4          \n\t"
1888            "vadd.vi      v16, v16, -4, v0.t      \n\t"
1889
1890            // A shift for the second dot within this K64
1891            "vsetvli      t0, x0, e64, mf2        \n\t"
1892            "vslidedown.vi  v2, v3, 2             \n\t"
1893
1894            // Dot products with FP16 scaling (accumulate into v24..v27)
1895            "vsetvli      t0, x0, e32, m1         \n\t"
1896            "vmadot.hp    v24, v3, v12, v1, 0, i8 \n\t"
1897            "vmadot.hp    v25, v3, v13, v1, 1, i8 \n\t"
1898            "vmadot.hp    v26, v3, v14, v1, 2, i8 \n\t"
1899            "vmadot.hp    v27, v3, v15, v1, 3, i8 \n\t"
1900            "vmadot.hp    v24, v2, v16, v1, 4, i8 \n\t"
1901            "vmadot.hp    v25, v2, v17, v1, 5, i8 \n\t"
1902            "vmadot.hp    v26, v2, v18, v1, 6, i8 \n\t"
1903            "vmadot.hp    v27, v2, v19, v1, 7, i8 \n\t"
1904
1905            // (K32-47 / K48-63 blocks continue unchanged...)
1906            // load B scales (32 bytes per K16, 16 times => 512B)
1907            "vsetvli      t0, x0, e64, m1         \n\t"
1908            "vmv.v.v      v1, v29                 \n\t"
1909
1910            // load B hmask chunk (64B per K16, 16 times => 1024B)
1911            "vsetvli      t0, x0, e8, mf2         \n\t"
1912            "vle8.v       v0, (s5)                \n\t"
1913            "addi         s5, s5, 64              \n\t"
1914
1915            // load A data (16 bytes per K16, 16 times => 256B)
1916            "vsetvli      t0, x0, e64, mf2        \n\t"
1917            "vslidedown.vi  v3, v3, 4             \n\t"
1918
1919            // unpack 2-bit qs + hmask -> signed values
1920            "vsetvli      t0, x0, e8, m1          \n\t"
1921            "vsll.vi      v8, v4, 2               \n\t"
1922            "vsll.vi      v9, v5, 2               \n\t"
1923            "vsll.vi      v10, v6, 2              \n\t"
1924            "vsll.vi      v11, v7, 2              \n\t"
1925
1926            "vsrl.vi      v20, v8, 6              \n\t"
1927            "vsrl.vi      v21, v9, 6              \n\t"
1928            "vnot.v       v0, v0                  \n\t"
1929            "vsrl.vi      v22, v10, 6             \n\t"
1930            "vsrl.vi      v23, v11, 6             \n\t"
1931
1932            "vsetvli      t0, x0, e8, m4          \n\t"
1933            "vadd.vi      v20, v20, -4, v0.t      \n\t"
1934
1935            // K48-63
1936            "vsetvli      t0, x0, e8, mf2         \n\t"
1937            "vle8.v       v0, (s5)                \n\t"
1938            "addi         s5, s5, 64              \n\t"
1939
1940            "vsetvli      t0, x0, e8, m1          \n\t"
1941            "vsrl.vi      v8, v4, 6               \n\t"
1942            "vsrl.vi      v9, v5, 6               \n\t"
1943            "vnot.v       v0, v0                  \n\t"
1944            "vsrl.vi      v10, v6, 6              \n\t"
1945            "vsrl.vi      v11, v7, 6              \n\t"
1946
1947            "vsetvli      t0, x0, e8, m4          \n\t"
1948            "vadd.vi      v8, v8, -4, v0.t        \n\t"
1949
1950            // load A data (16 bytes per K16, 16 times => 256B)
1951            "vsetvli      t0, x0, e64, mf2        \n\t"
1952            "vslidedown.vi  v2, v3, 2             \n\t"
1953
1954            "vsetvli      t0, x0, e32, m1         \n\t"
1955            "vmadot.hp    v24, v3, v20, v1, 0, i8 \n\t"
1956            "vmadot.hp    v25, v3, v21, v1, 1, i8 \n\t"
1957            "vmadot.hp    v26, v3, v22, v1, 2, i8 \n\t"
1958            "vmadot.hp    v27, v3, v23, v1, 3, i8 \n\t"
1959            "vmadot.hp    v24, v2, v8, v1, 4, i8  \n\t"
1960            "vmadot.hp    v25, v2, v9, v1, 5, i8  \n\t"
1961            "vmadot.hp    v26, v2, v10, v1, 6, i8 \n\t"
1962            "vmadot.hp    v27, v2, v11, v1, 7, i8 \n\t"
1963
1964            "addi         t3, t3, -1              \n\t"
1965            "bgtz         t3, K64_LPST%=          \n\t"
1966            "K64_LPND%=:                          \n\t"
1967
1968            // ---- End of K64 chunk: reduce fp16 accumulators -> fp32 and scale by A ----
1969            "vsetvli      t0, x0, e16, m1         \n\t"
1970            "vpack.vv     v12, v24, v25, 1        \n\t"
1971            "vpack.vv     v14, v26, v27, 1        \n\t"
1972            "vpack.vv     v16, v12, v14, 2        \n\t"
1973            "vsetvli      t0, x0, e16, mf2        \n\t"
1974            "vfwcvt.f.f.v v26, v16                \n\t"  // fp16 -> fp32 vector (qsum * b_scales)
1975
1976            // Load A scale and advance A pointer to next K256 superblock
1977            "flw          f0, (s2)                \n\t"
1978            "addi         s2, s2, 4+32+256        \n\t"
1979            "add          t4, s7, %[B_STR]        \n\t"  // next B block base
1980            "addi         s3, s2, 4+32            \n\t"  // reset A data pointer for next block
1981
1982            // Advance B pointers to next K256 superblock
1983            "addi         s5, t4, 32*16           \n\t"
1984            "mv           s4, t4                  \n\t"
1985            "addi         s6, s5, 32*32           \n\t"
1986            "addi         s8, s6, 1024            \n\t"
1987            "addi         s8, s8, 1024            \n\t"
1988            "addi         s7, t4, 0               \n\t"
1989            "addi         t2, t2, -1              \n\t"
1990
1991            // Final per-block scaling: a_scale * 16.0f
1992            "fmul.s       f0, f0, %[a_post_mul]   \n\t"
1993            // acc += (qsum * b_scales) * (a_scale*16)
1994            "vsetvli      t0, x0, e32, m1         \n\t"
1995            "vfmacc.vf    v31, f0, v26            \n\t"
1996
1997            "beqz         t2, BLK_LPND%=          \n\t"
1998
1999            // Preload next block's scales16 and rebuild v30 for vmadot.hp
2000            "vsetvli      t0, x0, e16, mf2        \n\t"
2001            "vle16.v      v1, (s8)                \n\t"
2002            "vsetvli      t0, x0, e16, m1         \n\t"
2003            "vpack.vv     v26, v1, v1, 3          \n\t"
2004            "vmv.v.v      v17, v26                \n\t"
2005            "vsetvli      t0, x0, e16, m1         \n\t"
2006            "vfmul.vf     v30, v17, %[q3_step]    \n\t"
2007
2008            // Reset fp16 partial accumulators for next K64 loop(s)
2009            "vsetvli      t0, x0, e32, m1         \n\t"
2010            "vxor.vv      v24, v16, v16           \n\t"
2011            "vxor.vv      v25, v16, v16           \n\t"
2012            "vxor.vv      v26, v16, v16           \n\t"
2013            "vxor.vv      v27, v16, v16           \n\t"
2014
2015            "li           t3, 4                   \n\t"
2016            "bgtz         t2, BLK_LPST%=          \n\t"
2017
2018            "BLK_LPND%=:                          \n\t"
2019            "vsetvli      t0, %[NBLKS], e32, m1   \n\t"
2020            "vse32.v      v31, (%[pC])            \n\t"
2021
2022            :
2023            : [KBLKS] "r"(k_blks), [NBLKS] "r"(nb_real), [pA] "r"(quant_a_ptr), [pB] "r"(quant_b_blk_data),
2024              [pC] "r"(c_ptr), [B_STR] "r"(b_ncol_block_stride), [q3_step] "f"(k_q3k_scale_step),
2025              [a_post_mul] "f"(k_a_scale_post_mul)
2026            : "cc", "memory", "t0", "t2", "t3", "t4", "t5", "f0", "f1", "s2", "s3", "s4", "s5", "s6", "s7", "s8");
2027#endif
2028    }
2029}
2030
2031void gemm_kernel_i8i3k_m4(size_t          blk_len,
2032                          const uint8_t * quant_a_ptr,
2033                          const uint8_t * quant_b_data,
2034                          float *         c_ptr,
2035                          size_t          count_m,
2036                          size_t          count_n,
2037                          size_t          k_blks,
2038                          size_t          ldc) {
2039    using blk_type           = nrow_block_q3_k<32>;
2040    constexpr size_t NB_COLS = 32;  //only support 32 in ASM
2041
2042    const blk_type * b_base = reinterpret_cast<const blk_type *>(quant_b_data);
2043
2044    int64_t a_blk_stride        = q8k_blk_size(256);
2045    int64_t a_nrow_block_stride = a_blk_stride * 4;
2046    int64_t b_ncol_block_stride = sizeof(blk_type);
2047
2048    for (size_t ni = 0; ni < count_n; ni += NB_COLS, c_ptr += NB_COLS) {
2049        size_t           nb_real          = std::min<size_t>(NB_COLS, count_n - ni);
2050        const blk_type * quant_b_blk_data = b_base + (ni / NB_COLS) * k_blks;
2051
2052        //------------------------------------------------------------------------------
2053        // A format
2054        // Ascale   fp32 * 1* 4row    128bit
2055        // Asum     int16 * 16 4row  1024bit
2056        // A M1K256 int8 4row        8192bit
2057        //------------------------------------------------------------------------------
2058        // B format
2059        // B_scl    uint8*N32*16    4096bit
2060        // B_Hmask  N32K16*16 1bit  8192bit
2061        // B_Qs     N32K16*16 2bit  16384bit
2062        // B scl16  fp16 * N32      512bit;
2063        //------------------------------------------------------------------------------
2064        //bias always be nullptr
2065        __asm__ volatile(
2066            // t2 = k_blks (each is K256 superblock)
2067            "mv           t2, %[KBLKS]            \n\t"
2068            // t3 = 256/64 = 4 (K64 iterations per superblock)
2069            "li           t3, 4                   \n\t"
2070            "mv           s2, %[pA]               \n\t"  // s2 = pASCL
2071            "addi         s3, %[pA], 16+128       \n\t"  // s3 = pAData, (pA+AScl+ASum)
2072
2073            // B block layout for nrow_block_q3_k<32>:
2074            // scales: 512B, hmask: 1024B, qs: 2048B, scales16: 64B
2075            "addi         s5, %[pB], 32*16        \n\t"  // s5 = pB_hmask (skip scales)
2076            "mv           s4, %[pB]               \n\t"  // s4 = pB_scales
2077            "addi         s6, s5, 1024            \n\t"  // s6 = pB_qs (skip hmask)
2078            "mv           s7, %[pB]               \n\t"  // s7 = pB_base
2079
2080            "vsetvli      t0, x0, e32, m1         \n\t"
2081            "vxor.vv      v24, v0, v0             \n\t"  // v24-v27: K256 temp accumulator
2082            "vxor.vv      v25, v0, v0             \n\t"
2083            "vxor.vv      v26, v0, v0             \n\t"
2084            "vxor.vv      v27, v0, v0             \n\t"
2085            "vxor.vv      v28, v0, v0             \n\t"  // v28-v31: final accumulator
2086            "vxor.vv      v29, v0, v0             \n\t"
2087            "vxor.vv      v30, v0, v0             \n\t"
2088            "vxor.vv      v31, v0, v0             \n\t"
2089
2090            // ordinary vmadot: vle*13 vecIns*96 vmadot*16
2091            ".align 4                             \n\t"
2092            "BLK_LPST%=:                          \n\t"
2093            "K64_LPST%=:                          \n\t"
2094
2095            // ========== K0-15: First K16 sub-block ==========
2096            // Load B INT8 scale factors (32 cols × 16 K16 blocks)
2097            "vsetvli      t0, x0, e8, m1          \n\t"
2098            "vle8.v       v8, (s4)                \n\t"
2099            "addi         s4, s4, 128             \n\t"
2100
2101            // Load B quantized data (32 cols × 16 elements × 2bit, stored in 4 groups)
2102            "vle8.v       v4, (s6)                \n\t"
2103            "addi         s6, s6, 128             \n\t"
2104            "vle8.v       v5, (s6)                \n\t"
2105            "addi         s6, s6, 128             \n\t"
2106            "vle8.v       v6, (s6)                \n\t"
2107            "addi         s6, s6, 128             \n\t"
2108            "vle8.v       v7, (s6)                \n\t"
2109            "addi         s6, s6, 128             \n\t"
2110
2111            // Load B hmask (32 cols × 16bit sign mask)
2112            "vsetvli      t0, x0, e8, mf2         \n\t"
2113            "vle8.v       v0, (s5)                \n\t"
2114            "addi         s5, s5, 64              \n\t"
2115
2116            // Load A data (4 rows × 16 elements × INT8)
2117            "vsetvli      t0, x0, e8, mf2         \n\t"
2118            "vle8.v       v12, (s3)               \n\t"
2119            "addi         s3, s3, 256             \n\t"  // Jump to next row
2120            "vle8.v       v13, (s3)               \n\t"
2121            "addi         s3, s3, 256             \n\t"
2122            "vle8.v       v14, (s3)               \n\t"
2123            "addi         s3, s3, 256             \n\t"
2124            "vle8.v       v15, (s3)               \n\t"
2125            "addi         s3, s3, -768+64         \n\t"  // Back to first row, advance 16 elements
2126
2127            // Pack A data: merge 4 rows into 2 vectors
2128            "vsetvli      t0, x0, e8, m1          \n\t"
2129            "vpack.vv     v16, v12, v13, 1        \n\t"
2130            "vpack.vv     v18, v14, v15, 1        \n\t"
2131            "vpack.vv     v2, v16, v18, 2         \n\t"
2132
2133            // unpack 2-bit qs + hmask -> signed values
2134            "vsetvli      t0, x0, e8, m1          \n\t"
2135            "vnot.v       v0, v0                  \n\t"
2136            "vand.vi      v12, v4, 0x3            \n\t"
2137            "vand.vi      v13, v5, 0x3            \n\t"
2138            "vand.vi      v14, v6, 0x3            \n\t"
2139            "vand.vi      v15, v7, 0x3            \n\t"
2140
2141            "vsetvli      t0, x0, e8, m4          \n\t"
2142            "vadd.vi      v12, v12, -4, v0.t      \n\t"
2143
2144            "vsetvli      t0, x0, e32, m1         \n\t"
2145            "vxor.vv      v16, v16, v16           \n\t"
2146            "vxor.vv      v18, v16, v16           \n\t"
2147            "vxor.vv      v20, v16, v16           \n\t"
2148            "vxor.vv      v22, v16, v16           \n\t"
2149
2150            "vmadot       v16, v2, v12, i8        \n\t"  // 4 rows × cols 0-7
2151            "vmadot       v18, v2, v13, i8        \n\t"  // 4 rows × cols 8-15
2152            "vmadot       v20, v2, v14, i8        \n\t"  // 4 rows × cols 16-23
2153            "vmadot       v22, v2, v15, i8        \n\t"  // 4 rows × cols 24-31
2154
2155            "vsetvli      t0, x0, e16, m1         \n\t"
2156            "vpack.vv     v12, v16, v18, 2        \n\t"  // Merge cols 0-15
2157            "vpack.vv     v14, v20, v22, 2        \n\t"  // Merge cols 16-31
2158            "vpack.vv     v16, v12, v14, 3        \n\t"  // Inter-row results (INT16)
2159            "vpack.vv     v18, v13, v15, 3        \n\t"
2160
2161            // apply B int8 scales (-32 bias has been applyed)
2162            "vsetvli      t0, x0, e8, mf4         \n\t"
2163            "vwadd.vx     v21, v8, x0             \n\t"  // INT8 → INT16
2164
2165            "vsetvli      t0, x0, e16, mf2        \n\t"
2166            "vwadd.vx     v23, v21, x0            \n\t"  // INT16 → INT32
2167
2168            // Accumulate to K256 accumulator: qsum * b_scale
2169            "vsetvli      t0, x0, e32, m1         \n\t"
2170            "vmacc.vv     v24, v16, v23           \n\t"  // Row 0
2171            "vmacc.vv     v25, v17, v23           \n\t"  // Row 1
2172            "vmacc.vv     v26, v18, v23           \n\t"  // Row 2
2173            "vmacc.vv     v27, v19, v23           \n\t"
2174
2175            // ========== K16-31, K32-47, K48-63: Similar processing ==========
2176            // load B scales (32 bytes per K16, 16 times => 512B)
2177            "vsetvli      t0, x0, e64, m1         \n\t"
2178            "vslidedown.vi  v8, v8, 4             \n\t"
2179
2180            // load B hmask chunk (64B per K16, 16 times => 1024B)
2181            "vsetvli      t0, x0, e8, mf2         \n\t"
2182            "vle8.v       v0, (s5)                \n\t"
2183            "addi         s5, s5, 64              \n\t"
2184
2185            // load A data (16 bytes per K16, 16 times => 256B)
2186            "vsetvli      t0, x0, e64, m1         \n\t"
2187            "vslidedown.vi  v2, v2, 8             \n\t"
2188
2189            // unpack 2-bit qs + hmask -> signed values
2190            "vsetvli      t0, x0, e8, m1          \n\t"
2191            "vsll.vi      v12, v4, 4              \n\t"
2192            "vsll.vi      v13, v5, 4              \n\t"
2193            "vsll.vi      v14, v6, 4              \n\t"
2194            "vsll.vi      v15, v7, 4              \n\t"
2195            "vnot.v       v0, v0                  \n\t"
2196
2197            "vsrl.vi      v12, v12, 6             \n\t"
2198            "vsrl.vi      v13, v13, 6             \n\t"
2199            "vsrl.vi      v14, v14, 6             \n\t"
2200            "vsrl.vi      v15, v15, 6             \n\t"
2201
2202            "vsetvli      t0, x0, e8, m4          \n\t"
2203            "vadd.vi      v12, v12, -4, v0.t      \n\t"
2204
2205            "vsetvli      t0, x0, e32, m1         \n\t"
2206            "vxor.vv      v16, v16, v16           \n\t"
2207            "vxor.vv      v18, v16, v16           \n\t"
2208            "vxor.vv      v20, v16, v16           \n\t"
2209            "vxor.vv      v22, v16, v16           \n\t"
2210
2211            "vmadot       v16, v2, v12, i8        \n\t"
2212            "vmadot       v18, v2, v13, i8        \n\t"
2213            "vmadot       v20, v2, v14, i8        \n\t"
2214            "vmadot       v22, v2, v15, i8        \n\t"
2215
2216            "vsetvli      t0, x0, e16, m1         \n\t"
2217            "vpack.vv     v12, v16, v18, 2        \n\t"
2218            "vpack.vv     v14, v20, v22, 2        \n\t"
2219            "vpack.vv     v16, v12, v14, 3        \n\t"  // N0-N31 in v16
2220            "vpack.vv     v18, v13, v15, 3        \n\t"
2221
2222            // apply B int8 scales (-32 bias has been applyed)
2223            "vsetvli      t0, x0, e8, mf4         \n\t"
2224            "vwadd.vx     v21, v8, x0             \n\t"  // int8 -> int16
2225
2226            "vsetvli      t0, x0, e16, mf2        \n\t"
2227            "vwadd.vx     v23, v21, x0            \n\t"  // int8 -> int16
2228
2229            // static_cast<int32_t>(qsum) * b_scale;
2230            "vsetvli      t0, x0, e32, m1         \n\t"
2231            "vmacc.vv     v24, v16, v23           \n\t"
2232            "vmacc.vv     v25, v17, v23           \n\t"
2233            "vmacc.vv     v26, v18, v23           \n\t"
2234            "vmacc.vv     v27, v19, v23           \n\t"
2235
2236            //K32-47
2237            // load B scales (32 bytes per K16, 16 times => 512B)
2238            "vsetvli      t0, x0, e64, m1         \n\t"
2239            "vslidedown.vi  v8, v8, 4             \n\t"
2240
2241            // load B hmask chunk (64B per K16, 16 times => 1024B)
2242            "vsetvli      t0, x0, e8, mf2         \n\t"
2243            "vle8.v       v0, (s5)                \n\t"
2244            "addi         s5, s5, 64              \n\t"
2245
2246            // load A data (16 bytes per K16, 16 times => 256B)
2247
2248            // unpack 2-bit qs + hmask -> signed values
2249            "vsetvli      t0, x0, e8, m1          \n\t"
2250            "vsll.vi      v12, v4, 2              \n\t"
2251            "vsll.vi      v13, v5, 2              \n\t"
2252            "vsll.vi      v14, v6, 2              \n\t"
2253            "vsll.vi      v15, v7, 2              \n\t"
2254            "vnot.v       v0, v0                  \n\t"
2255
2256            "vsrl.vi      v12, v12, 6             \n\t"
2257            "vsrl.vi      v13, v13, 6             \n\t"
2258            "vsrl.vi      v14, v14, 6             \n\t"
2259            "vsrl.vi      v15, v15, 6             \n\t"
2260
2261            "vsetvli      t0, x0, e8, m4          \n\t"
2262            "vadd.vi      v12, v12, -4, v0.t      \n\t"
2263
2264            "vsetvli      t0, x0, e32, m1         \n\t"
2265            "vxor.vv      v16, v16, v16           \n\t"
2266            "vxor.vv      v18, v16, v16           \n\t"
2267            "vxor.vv      v20, v16, v16           \n\t"
2268            "vxor.vv      v22, v16, v16           \n\t"
2269
2270            "vmadot       v16, v3, v12, i8        \n\t"
2271            "vmadot       v18, v3, v13, i8        \n\t"
2272            "vmadot       v20, v3, v14, i8        \n\t"
2273            "vmadot       v22, v3, v15, i8        \n\t"
2274
2275            "vsetvli      t0, x0, e16, m1         \n\t"
2276            "vpack.vv     v12, v16, v18, 2        \n\t"
2277            "vpack.vv     v14, v20, v22, 2        \n\t"
2278            "vpack.vv     v16, v12, v14, 3        \n\t"  // N0-N31 in v16
2279            "vpack.vv     v18, v13, v15, 3        \n\t"
2280
2281            // apply B int8 scales (-32 bias has been applyed)
2282            "vsetvli      t0, x0, e8, mf4         \n\t"
2283            "vwadd.vx     v21, v8, x0             \n\t"  // int8 -> int16
2284
2285            "vsetvli      t0, x0, e16, mf2        \n\t"
2286            "vwadd.vx     v23, v21, x0            \n\t"  // int8 -> int16
2287
2288            // static_cast<int32_t>(qsum) * b_scale;
2289            "vsetvli      t0, x0, e32, m1         \n\t"
2290            "vmacc.vv     v24, v16, v23           \n\t"
2291            "vmacc.vv     v25, v17, v23           \n\t"
2292            "vmacc.vv     v26, v18, v23           \n\t"
2293            "vmacc.vv     v27, v19, v23           \n\t"
2294
2295            // K48-63
2296            // load B scales (32 bytes per K16, 16 times => 512B)
2297            "vsetvli      t0, x0, e64, m1         \n\t"
2298            "vslidedown.vi  v8, v8, 4             \n\t"
2299
2300            // load B hmask chunk (64B per K16, 16 times => 1024B)
2301            "vsetvli      t0, x0, e8, mf2         \n\t"
2302            "vle8.v       v0, (s5)                \n\t"
2303            "addi         s5, s5, 64              \n\t"
2304
2305            // load A data (16 bytes per K16, 16 times => 256B)
2306            "vsetvli      t0, x0, e64, m1         \n\t"
2307            "vslidedown.vi  v3, v3, 8             \n\t"
2308
2309            "vsetvli      t0, x0, e8, m1          \n\t"
2310            "vnot.v       v0, v0                  \n\t"
2311            "vsrl.vi      v12, v4, 6              \n\t"
2312            "vsrl.vi      v13, v5, 6              \n\t"
2313            "vsrl.vi      v14, v6, 6              \n\t"
2314            "vsrl.vi      v15, v7, 6              \n\t"
2315
2316            "vsetvli      t0, x0, e8, m4          \n\t"
2317            "vadd.vi      v12, v12, -4, v0.t      \n\t"
2318
2319            "vsetvli      t0, x0, e32, m1         \n\t"
2320            "vxor.vv      v16, v16, v16           \n\t"
2321            "vxor.vv      v18, v16, v16           \n\t"
2322            "vxor.vv      v20, v16, v16           \n\t"
2323            "vxor.vv      v22, v16, v16           \n\t"
2324
2325            "vmadot       v16, v3, v12, i8        \n\t"
2326            "vmadot       v18, v3, v13, i8        \n\t"
2327            "vmadot       v20, v3, v14, i8        \n\t"
2328            "vmadot       v22, v3, v15, i8        \n\t"
2329
2330            "vsetvli      t0, x0, e16, m1         \n\t"
2331            "vpack.vv     v12, v16, v18, 2        \n\t"
2332            "vpack.vv     v14, v20, v22, 2        \n\t"
2333            "vpack.vv     v16, v12, v14, 3        \n\t"  // N0-N31 in v16
2334            "vpack.vv     v18, v13, v15, 3        \n\t"
2335
2336            // apply B int8 scales (-32 bias has been applyed)
2337            "vsetvli      t0, x0, e8, mf4         \n\t"
2338            "vwadd.vx     v21, v8, x0             \n\t"  // int8 -> int16
2339
2340            "vsetvli      t0, x0, e16, mf2        \n\t"
2341            "vwadd.vx     v23, v21, x0            \n\t"  // int8 -> int16
2342
2343            // static_cast<int32_t>(qsum) * b_scale;
2344            "vsetvli      t0, x0, e32, m1         \n\t"
2345            "vmacc.vv     v24, v16, v23           \n\t"
2346            "vmacc.vv     v25, v17, v23           \n\t"
2347            "vmacc.vv     v26, v18, v23           \n\t"
2348            "vmacc.vv     v27, v19, v23           \n\t"
2349
2350            "addi         t3, t3, -1              \n\t"
2351            "bgtz         t3, K64_LPST%=          \n\t"
2352            "K64_LPND%=:                          \n\t"
2353
2354            // ========== K256 superblock complete, apply scale factors ==========
2355            // Load A's 4 row scale factors (FP32)
2356            "flw          f0, (s2)                \n\t"
2357            "flw          f1, 4(s2)               \n\t"
2358            "flw          f2, 8(s2)               \n\t"
2359            "flw          f3, 12(s2)              \n\t"
2360            "add          s2, s2, %[A_STR]        \n\t"  // Advance to next superblock
2361            "add          t4, s7, %[B_STR]        \n\t"  // t4 = next B block address
2362            "addi         s3, s2, (4+32)*4        \n\t"
2363
2364            // Load B FP16 global scale factors (32 cols)
2365            "vsetvli      t0, x0, e16, mf2        \n\t"
2366            "vle16.v      v8, (s6)                \n\t"
2367
2368            // Update B pointers to next block
2369            "addi         s5, t4, 32*16           \n\t"
2370            "mv           s4, t4                  \n\t"
2371            "addi         s6, s5, 32*32           \n\t"
2372            "addi         s7, t4, 0               \n\t"
2373
2374            // ========== Type conversion and final scaling ==========
2375            // FP16 → FP32
2376            "vsetvli      t0, x0, e16, mf2        \n\t"
2377            "vfwcvt.f.f.v v9, v8                 \n\t"
2378
2379            // INT32 → FP32
2380            "vsetvli      t0, x0, e32, m1         \n\t"
2381            "vfcvt.f.x.v  v24, v24                \n\t"
2382            "vfcvt.f.x.v  v25, v25                \n\t"
2383            "vfcvt.f.x.v  v26, v26                \n\t"
2384            "vfcvt.f.x.v  v27, v27                \n\t"
2385
2386            // Compute a_scale * b_scale (4 rows)
2387            "vfmul.vf     v12, v9, f0             \n\t"
2388            "vfmul.vf     v13, v9, f1             \n\t"
2389            "vfmul.vf     v14, v9, f2             \n\t"
2390            "vfmul.vf     v15, v9, f3             \n\t"
2391
2392            // Final accumulation: result += qsum * a_scale * b_scale
2393            "vsetvli      t0, x0, e32, m1         \n\t"
2394            "vfmacc.vv    v28, v12, v24           \n\t"
2395            "vfmacc.vv    v29, v13, v25           \n\t"
2396            "vfmacc.vv    v30, v14, v26           \n\t"
2397            "vfmacc.vv    v31, v15, v27           \n\t"
2398
2399            // Prepare for next K superblock
2400            "addi         t2, t2, -1              \n\t"
2401            "vxor.vv      v24, v0, v0             \n\t"  // Clear K256 accumulator
2402            "vxor.vv      v25, v0, v0             \n\t"
2403            "vxor.vv      v26, v0, v0             \n\t"
2404            "vxor.vv      v27, v0, v0             \n\t"
2405            "li           t3, 4                   \n\t"
2406            "bgtz         t2, BLK_LPST%=          \n\t"
2407
2408            "BLK_LPND%=:                          \n\t"
2409
2410            // ========== Store results (4 rows × 32 cols) ==========
2411            "mv           t5, %[pC]               \n\t"
2412            "vsetvli      t0, %[NBLKS], e32, m1   \n\t"
2413            "vse32.v      v28, (%[pC])            \n\t"
2414            "add          t5, t5, %[LDC]          \n\t"
2415            "vse32.v      v29, (t5)               \n\t"
2416            "add          t5, t5, %[LDC]          \n\t"
2417            "vse32.v      v30, (t5)               \n\t"
2418            "add          t5, t5, %[LDC]          \n\t"
2419            "vse32.v      v31, (t5)               \n\t"
2420            "add          t5, t5, %[LDC]          \n\t"
2421            "FUNC_END%=:                          \n\t"
2422
2423            :
2424            : [KBLKS] "r"(k_blks), [NBLKS] "r"(nb_real), [pA] "r"(quant_a_ptr), [pB] "r"(quant_b_blk_data),
2425              [pC] "r"(c_ptr), [B_STR] "r"(b_ncol_block_stride), [A_STR] "r"(a_nrow_block_stride), [LDC] "r"(ldc * 4)
2426            : "cc", "memory", "t0", "t2", "t3", "t4", "t5", "f0", "f1", "f2", "f3", "s2", "s3", "s4", "s5", "s6", "s7");
2427    }
2428}
2429
2430void gemm_kernel_i8i4_m1(size_t          blk_len,
2431                         const uint8_t * quant_a_ptr,
2432                         const uint8_t * quant_b_data,
2433                         const uint8_t * quant_b_zp,
2434                         float *         c_ptr,
2435                         size_t          count_m,
2436                         size_t          count_n,
2437                         size_t          k_blks,
2438                         size_t          ldc) {
2439    if (quant_b_zp == NULL) {
2440        for (size_t n = 0; n < count_n; n += 32) {
2441            size_t    nblks         = (count_n - n) > 32 ? 32 : count_n - n;
2442            uint8_t * QuantBDataPtr = (uint8_t *) quant_b_data +      //
2443                                      n * k_blks * blk_len / 2 +      // b data
2444                                      n * k_blks * sizeof(_Float16);  // scale
2445            float * CPtr = c_ptr + n;
2446            size_t  cnt  = k_blks;
2447
2448            // A format Version_1 (FP32 SCALE FOR Normal VMADOTins of IME2)
2449            // A M1K32 int8    256bit
2450            // Ascale fp32 * 1  32bit
2451            // || scl*1(fp32) | Asum(int16) | blk0 || scl*1(fp32) | Asum(int16) | blk0 || ...
2452            // || Element                          || Element                          || ...
2453            // B format
2454            // B N8K32 int4    1024bit
2455            //   4VRF, N32K32, 4096bit
2456            // Bscale fp16 * N32 512bit;
2457            // || scl*32..(fp16) | blk0 blk1 ... blk31 || scl*32..(fp16) | blk0 blk1 ... blk31 || ...
2458            // || Element                              || Element                              || ...
2459#if 0
2460            //bias always be nullptr
2461            __asm__ volatile(
2462
2463                // t3 = k/32
2464                "mv           t3, %[BCK]              \n\t"
2465                "mv           t4, %[NBLKS]            \n\t"
2466                "mv           s2, %[pA]               \n\t"  // s2 = pASCL
2467                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pAData, (pA+AScl+ASum)
2468                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
2469                "addi         s5, %[pB], 32*2         \n\t"  // s5 = pBdata;
2470                "mv           s6, %[pC]               \n\t"
2471
2472                "vsetvli      t0, x0, e32, m1         \n\t"
2473                "vxor.vv      v2, v0, v0              \n\t"  // clear acc
2474
2475                // ordinary vmadot: vle*6 flw*1 vecIns*21 vmadot*8
2476                ".align 4                             \n\t"
2477                "_K_LPST%=:                           \n\t"
2478
2479                "vsetvli      t0, x0, e8, m1          \n\t"
2480                "vl4r.v       v4, (s5)                \n\t"  // B Data 4VRF * 8Row * 32
2481                "addi         s5, s5, 128*4+64        \n\t"  // 1024bit
2482
2483                "vsetvli      t0, x0, e8, mf2         \n\t"
2484                "vle8.v       v0, (s4)                \n\t"  // B Scale 4VRF*8Row*FP16 = 512bit
2485                "addi         s4, s4, 64+128*4        \n\t"
2486
2487                "vsetvli      t0, x0, e8, mf4         \n\t"
2488                "vle8.v       v3, (s3)                \n\t"  // A Data M1*K32*int8 = 256bit
2489                "addi         s3, s3, 32+6            \n\t"
2490
2491                "flw          f0, (s2)                \n\t"  // A Scale fp32
2492                "lh           t2, 4(s2)               \n\t"  // A sum of int16
2493                "addi         s2, s2, 6+32            \n\t"
2494
2495                "vsetvli      t0, zero, e8, m1        \n\t"
2496                "vsrl.vi      v24, v3, 4              \n\t"
2497
2498                "vnpack4.vv   v8, v3, v3, 3           \n\t"  // lo4 of A
2499                "vnpack4.vv   v10, v24, v24, 3        \n\t"  // hi4 of A
2500
2501                "vsetvli      t0, x0, e32, m1         \n\t"
2502                "vxor.vv      v16, v16, v16           \n\t"
2503                "vxor.vv      v18, v16, v16           \n\t"
2504                "vxor.vv      v20, v16, v16           \n\t"
2505                "vxor.vv      v22, v16, v16           \n\t"
2506
2507                "vmadotsu     v16, v10, v4, i4        \n\t"  // M0 N0 - N7 INT32(256bit)
2508                "vmadotsu     v18, v10, v5, i4        \n\t"  // M0 N8 - N15
2509                "vmadotsu     v20, v10, v6, i4        \n\t"  // M0 N16 - N23
2510                "vmadotsu     v22, v10, v7, i4        \n\t"  // M0 N24 - N31
2511
2512                "vsll.vi      v16, v16, 4             \n\t"
2513                "vsll.vi      v18, v18, 4             \n\t"
2514                "vsll.vi      v20, v20, 4             \n\t"
2515                "vsll.vi      v22, v22, 4             \n\t"
2516
2517                "vmadotu      v16, v8, v4, i4         \n\t"
2518                "vmadotu      v18, v8, v5, i4         \n\t"
2519                "vmadotu      v20, v8, v6, i4         \n\t"
2520                "vmadotu      v22, v8, v7, i4         \n\t"
2521
2522                "vsetvli      t0, x0, e16, m1         \n\t"
2523                "vmv.v.i      v28, 8                  \n\t"
2524                "vpack.vv     v24, v16, v18, 2        \n\t"
2525                "vpack.vv     v26, v20, v22, 2        \n\t"
2526                "vpack.vv     v16, v24, v26, 3        \n\t"
2527
2528                "vwmul.vx     v24, v28, t2            \n\t"
2529                "vsetvli      t0, x0, e32, m1         \n\t"
2530                "vadd.vv      v16, v16, v24           \n\t"
2531
2532                // b_scale fp16 -> fp32
2533                "vsetvli      t0, x0, e16, mf2        \n\t"
2534                "vfwcvt.f.f.v v24, v0                 \n\t"
2535                // mac result i32 -> fp32
2536                "vsetvli      t0, x0, e32, m1         \n\t"
2537                "vfcvt.f.x.v  v26, v16                \n\t"
2538                // a_scale * b_scale;
2539                "vfmul.vf     v1, v24, f0             \n\t"
2540                // static_cast<float>(qsum) * a_scale * b_scale;
2541                "vfmacc.vv    v2, v1, v26             \n\t"
2542
2543                "addi         t3, t3, -1              \n\t"
2544                "bgtz         t3, _K_LPST%=           \n\t"
2545                "_K_LPND%=:                           \n\t"
2546
2547                //-----------------------------------------
2548                // STORE Equal 32N-------------------------
2549                "_ST32%=:                             \n\t"
2550                "vsetvli      t0, t4, e32, m1         \n\t"
2551                "vse32.v      v2, (s6)                \n\t"  // M0 [N0 : N32]; FP32(1024bit)
2552
2553                "_FUNC_END%=:                         \n\t"
2554
2555                :
2556                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
2557                : "cc", "t0", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6");
2558#else
2559            __asm__ volatile(
2560
2561                // t3 = k/32
2562                "mv           t3, %[BCK]              \n\t"
2563                "mv           t4, %[NBLKS]            \n\t"
2564                "vsetvli      t0, x0, e16, m1         \n\t"
2565                "vmv.v.i      v0, 1                   \n\t"  // init the scale
2566                "mv           s2, %[pA]               \n\t"  // s2 = pASCL
2567                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pAData, (pA+AScl+ASum)
2568                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
2569                "addi         s5, %[pB], 32*2         \n\t"  // s5 = pBdata;
2570                "mv           s6, %[pC]               \n\t"
2571
2572                "vsll.vi      v1, v0, 4               \n\t"
2573                "vxor.vv      v2, v0, v0              \n\t"  // clear acc
2574                "vfcvt.f.x.v  v0, v0                  \n\t"
2575                "vfcvt.f.x.v  v1, v1                  \n\t"
2576
2577                // vmadot hp: vle*7 flw*1 vecIns*14 vmadot*8
2578                ".align 4                             \n\t"
2579                "_K_LPST%=:                           \n\t"
2580
2581                "vsetvli      t0, x0, e8, m1          \n\t"
2582                "vl4r.v       v4, (s5)                \n\t"  // B Data 4VRF * 8Row * 32
2583                "addi         s5, s5, 128*4+64        \n\t"  // 1024bit
2584
2585                "vsetvli      t0, x0, e8, mf2         \n\t"
2586                "vle8.v       v30, (s4)               \n\t"  // B Scale 4VRF*8Row*FP16 = 512bit
2587                "addi         s4, s4, 64+128*4        \n\t"
2588
2589                "vsetvli      t0, x0, e8, mf4         \n\t"
2590                "vle8.v       v3, (s3)                \n\t"  // A Data M1*K32*int8 = 256bit
2591                "addi         s3, s3, 32+6            \n\t"
2592
2593                "flw          f0, (s2)                \n\t"  // A Scale fp32
2594                "lh           t2, 4(s2)               \n\t"  // A sum of int16
2595                "addi         s2, s2, 6+32            \n\t"
2596
2597                "vsetvli      t0, x0, e16, m1         \n\t"
2598                "vmv.v.i      v28, 8                  \n\t"  // Bzp u8 -> u16
2599                "vsetvli      t0, x0, e8, m1          \n\t"
2600                "vsrl.vi      v24, v3, 4              \n\t"
2601
2602                "vsetvli      t0, x0, e16, m1         \n\t"
2603                "vmul.vx      v26, v28, t2            \n\t"  // asum*zp i16*i16
2604                "vnpack4.vv   v8, v3, v3, 3           \n\t"  // lo4 of A
2605                "vnpack4.vv   v10, v24, v24, 3        \n\t"  // hi4 of A
2606
2607                "vfcvt.f.x.v  v16, v26                \n\t"  // zp i16 -> fp16
2608                "vadd.vi      v18, v16, 0             \n\t"
2609                "vadd.vi      v20, v16, 0             \n\t"
2610                "vadd.vi      v22, v16, 0             \n\t"
2611
2612                "vmadotsu.hp  v16, v10, v4, v1, 0, i4 \n\t"  // high 4
2613                "vmadotsu.hp  v18, v10, v5, v1, 0, i4 \n\t"
2614                "vmadotsu.hp  v20, v10, v6, v1, 0, i4 \n\t"
2615                "vmadotsu.hp  v22, v10, v7, v1, 0, i4 \n\t"
2616                "vmadotu.hp   v16, v8, v4, v0, 0, i4  \n\t"  // low 4
2617                "vmadotu.hp   v18, v8, v5, v0, 0, i4  \n\t"
2618                "vmadotu.hp   v20, v8, v6, v0, 0, i4  \n\t"
2619                "vmadotu.hp   v22, v8, v7, v0, 0, i4  \n\t"
2620
2621                "vpack.vv     v24, v16, v18, 1        \n\t"
2622                "vpack.vv     v26, v20, v22, 1        \n\t"
2623                "vpack.vv     v16, v24, v26, 2        \n\t"
2624
2625                "vsetvli      t0, x0, e16, mf2        \n\t"
2626                // mac result * b_scale; f16*f16->f32
2627                "vfwmul.vv     v31, v30, v16          \n\t"
2628
2629                "vsetvli      t0, x0, e32, m1         \n\t"
2630                // static_cast<float>(qsum * b_scale) * a_scale;
2631                "vfmacc.vf    v2, f0, v31             \n\t"
2632
2633                "addi         t3, t3, -1              \n\t"
2634                "bgtz         t3, _K_LPST%=           \n\t"
2635                "_K_LPND%=:                           \n\t"
2636
2637                //-----------------------------------------
2638                // STORE Equal 32N-------------------------
2639                "_ST32%=:                             \n\t"
2640                "vsetvli      t0, t4, e32, m1         \n\t"
2641                "vse32.v      v2, (s6)                \n\t"  // M0 [N0 : N32]; FP32(1024bit)
2642
2643                "_FUNC_END%=:                         \n\t"
2644
2645                :
2646                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
2647                : "cc", "t0", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6");
2648
2649#endif
2650        }
2651    } else {
2652        for (size_t n = 0; n < count_n; n += 32) {
2653            size_t    nblks         = (count_n - n) > 32 ? 32 : count_n - n;
2654            uint8_t * QuantBDataPtr = (uint8_t *) quant_b_data +      //
2655                                      n * k_blks * blk_len / 2 +      // b data
2656                                      n * k_blks * sizeof(uint8_t) +  // b zp
2657                                      n * k_blks * sizeof(_Float16);  // scale
2658            float * CPtr = c_ptr + n;
2659            size_t  cnt  = k_blks;
2660
2661            // A format Version_1 (FP32 SCALE FOR Normal VMADOTins of IME2)
2662            // A M1K32 int8    256bit
2663            // Ascale fp32 * 1  32bit
2664            // || scl*1(fp32) | Asum(int16) | blk0 || scl*1(fp32) | Asum(int16) | blk0 || ...
2665            // || Element                          || Element                          || ...
2666            // B format
2667            // B N8K32 int4    1024bit
2668            //   4VRF, N32K32, 4096bit
2669            // Bscale fp16 * N32 512bit;
2670            // Bzp uint8_t * N32 256bit;
2671            // || scl*32..(fp16) | zp*32(uint8) | blk0 blk1 ... blk31 || scl*32..(fp16)  ...
2672            // || Element                                             || Element         ...
2673
2674            //bias always be nullptr
2675#if 0
2676            __asm__ volatile(
2677
2678                // t3 = k/32
2679                "mv           t3, %[BCK]              \n\t"
2680                "mv           t4, %[NBLKS]            \n\t"
2681                "mv           s2, %[pA]               \n\t"  // s2 = pASCL
2682                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pAData, (pA+AScl+ASum)
2683                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
2684                "addi         s5, %[pB], 32*3         \n\t"  // s5 = pBdata, (pB+BScl+Bzp)
2685                "mv           s6, %[pC]               \n\t"
2686
2687                "vsetvli      t0, x0, e32, m1         \n\t"
2688                "vxor.vv      v2, v0, v0              \n\t"  // clear acc
2689
2690                // ordinary vmadot: vle*6 flw*1 vecIns*21 vmadot*8
2691                ".align 4                             \n\t"
2692                "_K_LPST%=:                           \n\t"
2693
2694                "vsetvli      t0, x0, e8, m1          \n\t"
2695                "vl4r.v       v4, (s5)                \n\t"  // B Data 4VRF * 8Row * 32
2696                "addi         s5, s5, 128*4+96        \n\t"  // 1024bit
2697
2698                "vsetvli      t0, x0, e8, mf2         \n\t"
2699                "vle8.v       v0, (s4)                \n\t"  // B Scale 4VRF*8Row*FP16 = 512bit
2700                "addi         s4, s4, 64              \n\t"
2701
2702                "vsetvli      t0, x0, e8, mf4         \n\t"
2703                "vle8.v       v3, (s3)                \n\t"  // A Data M1*K32*int8 = 256bit
2704                "addi         s3, s3, 32+6            \n\t"
2705
2706                "flw          f0, (s2)                \n\t"  // A Scale fp32
2707                "lh           t2, 4(s2)               \n\t"  // A sum of int16
2708                "addi         s2, s2, 6+32            \n\t"
2709
2710                "vsetvli      t0, zero, e8, m1        \n\t"
2711                "vsrl.vi      v24, v3, 4              \n\t"
2712
2713                "vnpack4.vv   v8, v3, v3, 3           \n\t"  // lo4 of A
2714                "vnpack4.vv   v10, v24, v24, 3        \n\t"  // hi4 of A
2715
2716                "vsetvli      t0, x0, e32, m1         \n\t"
2717                "vxor.vv      v16, v16, v16           \n\t"
2718                "vxor.vv      v18, v16, v16           \n\t"
2719                "vxor.vv      v20, v16, v16           \n\t"
2720                "vxor.vv      v22, v16, v16           \n\t"
2721
2722                "vmadotsu     v16, v10, v4, i4        \n\t"  // M0 N0 - N7 INT32(256bit)
2723                "vmadotsu     v18, v10, v5, i4        \n\t"  // M0 N8 - N15
2724                "vmadotsu     v20, v10, v6, i4        \n\t"  // M0 N16 - N23
2725                "vmadotsu     v22, v10, v7, i4        \n\t"  // M0 N24 - N31
2726
2727                "vsll.vi      v16, v16, 4             \n\t"
2728                "vsll.vi      v18, v18, 4             \n\t"
2729                "vsll.vi      v20, v20, 4             \n\t"
2730                "vsll.vi      v22, v22, 4             \n\t"
2731
2732                "vsetvli      t0, x0, e8, m1          \n\t"
2733                "vle8.v       v1, (s4)                \n\t"  // Bzp
2734                "addi         s4, s4, 32+128*4        \n\t"
2735
2736                "vmadotu      v16, v8, v4, i4         \n\t"
2737                "vmadotu      v18, v8, v5, i4         \n\t"
2738                "vmadotu      v20, v8, v6, i4         \n\t"
2739                "vmadotu      v22, v8, v7, i4         \n\t"
2740
2741                "vwaddu.vx    v28, v1, x0             \n\t"  // uint8 -> uint16
2742                "vpack.vv     v24, v16, v18, 2        \n\t"
2743                "vpack.vv     v26, v20, v22, 2        \n\t"
2744                "vpack.vv     v16, v24, v26, 3        \n\t"
2745
2746                "vsetvli      t0, x0, e16, m1         \n\t"
2747                "vwmul.vx     v24, v28, t2            \n\t"
2748                "vsetvli      t0, x0, e32, m1         \n\t"
2749                "vadd.vv      v16, v16, v24           \n\t"
2750
2751                // b_scale fp16 -> fp32
2752                "vsetvli      t0, x0, e16, mf2        \n\t"
2753                "vfwcvt.f.f.v v24, v0                 \n\t"
2754                // mac result i32 -> fp32
2755                "vsetvli      t0, x0, e32, m1         \n\t"
2756                "vfcvt.f.x.v  v26, v16                \n\t"
2757                // a_scale * b_scale;
2758                "vfmul.vf     v1, v24, f0             \n\t"
2759                // static_cast<float>(qsum) * a_scale * b_scale;
2760                "vfmacc.vv    v2, v1, v26             \n\t"
2761
2762                "addi         t3, t3, -1              \n\t"
2763                "bgtz         t3, _K_LPST%=           \n\t"
2764                "_K_LPND%=:                           \n\t"
2765
2766                //-----------------------------------------
2767                // STORE Equal 32N-------------------------
2768                "_ST32%=:                             \n\t"
2769                "vsetvli      t0, t4, e32, m1         \n\t"
2770                "vse32.v      v2, (s6)                \n\t"  // M0 [N0 : N32]; FP32(1024bit)
2771
2772                "_FUNC_END%=:                         \n\t"
2773
2774                :
2775                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
2776                : "cc", "t0", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6");
2777#else
2778            __asm__ volatile(
2779
2780                // t3 = k/32
2781                "mv           t3, %[BCK]              \n\t"
2782                "mv           t4, %[NBLKS]            \n\t"
2783                "vsetvli      t0, x0, e16, m1         \n\t"
2784                "vmv.v.i      v0, 1                   \n\t"  // init the scale
2785                "mv           s2, %[pA]               \n\t"  // s2 = pASCL
2786                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pAData, (pA+AScl+ASum)
2787                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
2788                "addi         s5, %[pB], 32*3         \n\t"  // s5 = pBdata, (pB+BScl+Bzp)
2789                "mv           s6, %[pC]               \n\t"
2790
2791                "vsll.vi      v1, v0, 4               \n\t"
2792                "vxor.vv      v2, v0, v0              \n\t"  // clear acc
2793                "vfcvt.f.x.v  v0, v0                  \n\t"
2794                "vfcvt.f.x.v  v1, v1                  \n\t"
2795
2796                // vmadot hp: vle*6 flw*1 vecIns*14 vmadot*8
2797                ".align 4                             \n\t"
2798                "_K_LPST%=:                           \n\t"
2799
2800                "vsetvli      t0, x0, e8, m1          \n\t"
2801                "vl4r.v       v4, (s5)                \n\t"  // B Data 4VRF * 8Row * 32
2802                "addi         s5, s5, 128*4+96        \n\t"  // 1024bit
2803
2804                "vsetvli      t0, x0, e8, mf2         \n\t"
2805                "vle8.v       v30, (s4)               \n\t"  // B Scale 4VRF*8Row*FP16 = 512bit
2806                "addi         s4, s4, 64              \n\t"
2807
2808                "vsetvli      t0, x0, e8, mf4         \n\t"
2809                "vle8.v       v31, (s4)               \n\t"  // B zp 32Row*uint8 = 256bit
2810                "addi         s4, s4, 32+128*4        \n\t"
2811
2812                "vle8.v       v3, (s3)                \n\t"  // A Data M1*K32*int8 = 256bit
2813                "addi         s3, s3, 32+6            \n\t"
2814
2815                "flw          f0, (s2)                \n\t"  // A Scale fp32
2816                "lh           t2, 4(s2)               \n\t"  // A sum of int16
2817                "addi         s2, s2, 6+32            \n\t"
2818
2819                "vsetvli      t0, x0, e8, m1          \n\t"
2820                "vsrl.vi      v24, v3, 4              \n\t"
2821
2822                "vsetvli      t0, x0, e16, m1         \n\t"
2823                "vnpack4.vv   v8, v3, v3, 3           \n\t"  // lo4 of A
2824                "vnpack4.vv   v10, v24, v24, 3        \n\t"  // hi4 of A
2825
2826                "vxor.vv      v16, v16, v16           \n\t"
2827                "vxor.vv      v18, v16, v16           \n\t"
2828                "vxor.vv      v20, v16, v16           \n\t"
2829                "vxor.vv      v22, v16, v16           \n\t"
2830
2831                "vmadotsu.hp  v16, v10, v4, v1, 0, i4 \n\t"  // high 4
2832                "vmadotsu.hp  v18, v10, v5, v1, 0, i4 \n\t"
2833                "vmadotsu.hp  v20, v10, v6, v1, 0, i4 \n\t"
2834                "vmadotsu.hp  v22, v10, v7, v1, 0, i4 \n\t"
2835                "vmadotu.hp   v16, v8, v4, v0, 0, i4  \n\t"  // low 4
2836                "vmadotu.hp   v18, v8, v5, v0, 0, i4  \n\t"
2837                "vmadotu.hp   v20, v8, v6, v0, 0, i4  \n\t"
2838                "vmadotu.hp   v22, v8, v7, v0, 0, i4  \n\t"
2839
2840                "vsetvli      t0, x0, e8, mf4         \n\t"
2841                "vwaddu.vx    v28, v31, x0            \n\t"  // Bzp u8 -> u16
2842
2843                "vsetvli      t0, x0, e8, m1          \n\t"
2844                "vpack.vv     v24, v16, v18, 1        \n\t"
2845                "vpack.vv     v26, v20, v22, 1        \n\t"
2846                "vpack.vv     v16, v24, v26, 2        \n\t"
2847
2848                "vsetvli      t0, x0, e16, mf2        \n\t"
2849                "vmul.vx      v26, v28, t2            \n\t"  // asum*zp i16*i16
2850                "vfwcvt.f.f.v v22, v30                \n\t"  // b_scale fp16 -> fp32
2851                "vfcvt.f.x.v  v18, v26                \n\t"  // zp i16 -> fp16
2852                "vsetvli      t0, x0, e16, m1         \n\t"
2853                "vfwadd.vv    v20, v18, v16           \n\t"
2854
2855                "vsetvli      t0, x0, e32, m1         \n\t"
2856                // mac result * b_scale; f32*f32->f32
2857                "vfmul.vv     v31, v22, v20           \n\t"
2858
2859                "vsetvli      t0, x0, e32, m1         \n\t"
2860                // static_cast<float>(qsum * b_scale) * a_scale;
2861                "vfmacc.vf    v2, f0, v31             \n\t"
2862
2863                "addi         t3, t3, -1              \n\t"
2864                "bgtz         t3, _K_LPST%=           \n\t"
2865                "_K_LPND%=:                           \n\t"
2866
2867                //-----------------------------------------
2868                // STORE Equal 32N-------------------------
2869                "_ST32%=:                             \n\t"
2870                "vsetvli      t0, t4, e32, m1         \n\t"
2871                "vse32.v      v2, (s6)                \n\t"  // M0 [N0 : N32]; FP32(1024bit)
2872
2873                "_FUNC_END%=:                         \n\t"
2874
2875                :
2876                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
2877                : "cc", "t0", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6");
2878#endif
2879        }
2880    }
2881}
2882
2883void gemm_kernel_i8i4_hp_m1(size_t          blk_len,
2884                            const uint8_t * quant_a_ptr,
2885                            const uint8_t * quant_b_data,
2886                            const uint8_t * quant_b_zp,
2887                            float *         c_ptr,
2888                            size_t          count_m,
2889                            size_t          count_n,
2890                            size_t          k_blks,
2891                            size_t          ldc) {
2892    constexpr size_t NB_COLS                = 32;
2893    constexpr size_t k_subblks_per_superblk = 8;
2894
2895    struct block_q4_0x32_layout {
2896        _Float16 d[NB_COLS];
2897        uint8_t  qs[16 * NB_COLS];
2898    };
2899
2900    GGML_ASSERT(blk_len == 256);
2901
2902    const size_t b_superblk_stride = sizeof(block_q4_0x32_layout) * k_subblks_per_superblk +
2903                                     (quant_b_zp ? NB_COLS * k_subblks_per_superblk * sizeof(uint8_t) : 0);
2904    const size_t b_tile_stride = k_blks * b_superblk_stride;
2905
2906    if (quant_b_zp == NULL) {
2907        for (size_t ni = 0; ni < count_n; ni += 32) {
2908            uint8_t * b_data = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_tile_stride;
2909            int8_t *  a_data = (int8_t *) quant_a_ptr;
2910            float *   dst_c  = c_ptr + ni;
2911
2912            asm volatile(
2913                "vsetvli        t0, x0, e16, m1         \n\t"
2914                "vxor.vv        v31, v31, v31           \n\t"  // init acc to zero
2915                "mv             t4, %[BK]               \n\t"
2916                "li             t0, 0x4c00              \n\t"  // 16 in fp16
2917                "fmv.h.x        fa0, t0                 \n\t"
2918
2919                ".align 4                               \n\t"
2920                "BLK_LOOP%=:                            \n\t"
2921                "li             t5, 8                   \n\t"
2922                "addi           t6, %[A], 288           \n\t"  // point to blk scale
2923                "flh            ft1, (t6)               \n\t"
2924                "addi           t6, %[A], 272           \n\t"  // point to asum
2925
2926                // init the acc fp16
2927                "vsetvli        t0, x0, e16, m1         \n\t"
2928                "vxor.vv        v16, v18, v18           \n\t"
2929                "vxor.vv        v17, v18, v18           \n\t"
2930                "vxor.vv        v18, v18, v18           \n\t"
2931                "vxor.vv        v19, v18, v18           \n\t"
2932
2933                "INNER_BLK_LOOP%=:                      \n\t"
2934                // load a sum and scale
2935                "flh            fa1, (t6)               \n\t"
2936                "addi           t6, t6, 2               \n\t"
2937                "flh            ft0, (%[A])             \n\t"
2938                "addi           %[A], %[A], 2           \n\t"
2939                // load A
2940                "vsetvli        t0, x0, e8, mf4         \n\t"
2941                "vle8.v         v3, (%[A])              \n\t"  // 1x32@i8
2942                "addi           %[A], %[A], 32          \n\t"
2943
2944                // load scale B and B
2945                "vsetvli        t0, x0, e16, mf2        \n\t"
2946                "vle16.v        v8, (%[B])              \n\t"  // b_scale fp16
2947                "addi           %[B], %[B], 64          \n\t"
2948                "vl4r.v         v4, (%[B])              \n\t"  // 32*32@i4
2949                "addi           %[B], %[B], 512         \n\t"
2950                "vfmul.vf       v8, v8, ft0             \n\t"  // scale b * scale a
2951                "vfmul.vf       v9, v8, fa0             \n\t"
2952                "vfmul.vf       v10, v8, fa1            \n\t"  // scale b * scale a * asm
2953                "vfwmacc.vf     v31, ft1, v10           \n\t"  // asum * scale a * scale b * blk scale
2954
2955                "vsetvli        t0, x0, e8, m1          \n\t"
2956                "vpack.vv       v0, v8, v9, 3           \n\t"
2957                "vsrl.vi        v28, v3, 4              \n\t"
2958
2959                "vsetvli        t0, x0, e16, m1         \n\t"
2960                "vnpack4.vv     v2, v3, v3, 3           \n\t"  // lo4 of A
2961                "vnpack4.vv     v3, v28, v28, 3         \n\t"  // hi4 of A
2962
2963                // i4 * i4 vmadot
2964                "vsetvli        t0, x0, e16, m1         \n\t"
2965                "vmadotsu.hp    v16, v3, v4, v0, 4, i4  \n\t"  // high 4
2966                "vmadotsu.hp    v17, v3, v5, v0, 5, i4  \n\t"
2967                "vmadotsu.hp    v18, v3, v6, v0, 6, i4  \n\t"
2968                "vmadotsu.hp    v19, v3, v7, v0, 7, i4  \n\t"
2969                "vmadotu.hp     v16, v2, v4, v0, 0, i4  \n\t"  // low 4
2970                "vmadotu.hp     v17, v2, v5, v0, 1, i4  \n\t"
2971                "vmadotu.hp     v18, v2, v6, v0, 2, i4  \n\t"
2972                "vmadotu.hp     v19, v2, v7, v0, 3, i4  \n\t"
2973
2974                "addi           t5, t5, -1              \n\t"
2975                "bgtz           t5, INNER_BLK_LOOP%=    \n\t"
2976
2977                "vpack.vv       v8, v16, v17, 1         \n\t"
2978                "vpack.vv       v12, v18, v19, 1        \n\t"
2979                "vpack.vv       v20, v8, v12, 2         \n\t"
2980
2981                "vsetvli        t0, x0, e16, mf2        \n\t"
2982                "addi           t4, t4, -1              \n\t"
2983                "vfwmacc.vf     v31, ft1, v20           \n\t"
2984                //"vsetvli        t0, x0, e32, m1         \n\t"
2985                //"vfmul.vf       v31, v31, ft1           \n\t"  // blk scale
2986
2987                // update A ptr
2988                "addi           %[A], t6, 2             \n\t"
2989
2990                "bgtz           t4, BLK_LOOP%=          \n\t"
2991
2992                // save
2993                "vsetvli        t0, x0, e32, m1         \n\t"
2994                "vse32.v        v31, (%[DST])           \n\t"
2995                : [A] "+r"(a_data), [B] "+r"(b_data)
2996                : [DST] "r"(dst_c), [BK] "r"(k_blks)
2997                : "t0", "t1", "t2", "t3", "t4", "t5", "t6", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9",
2998                  "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23",
2999                  "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "ft0", "ft1");
3000        }
3001    } else {
3002        // TODO: support quant_b_zp for i8i4 hp kernel
3003        GGML_ABORT("gemm_kernel_i8i4_hp_m1 with quant_b_zp is not supported yet");
3004    }
3005}
3006
3007void gemm_kernel_i8i4_m4(size_t          blk_len,
3008                         const uint8_t * quant_a_ptr,
3009                         const uint8_t * quant_b_data,
3010                         const uint8_t * quant_b_zp,
3011                         float *         c_ptr,
3012                         size_t          count_m,
3013                         size_t          count_n,
3014                         size_t          k_blks,
3015                         size_t          ldc) {
3016    int64_t b_data_stride =
3017        k_blks * (sizeof(ggml_fp16_t) + 16 * sizeof(int8_t) + (quant_b_zp != NULL ? sizeof(int8_t) : 0));
3018    if (quant_b_zp == NULL) {
3019        for (size_t ni = 0; ni < count_n; ni += 32) {
3020            uint8_t * b_data = (uint8_t *) quant_b_data + ni * b_data_stride;
3021            int8_t *  a_data = (int8_t *) quant_a_ptr;
3022            float *   dst_c  = c_ptr + ni;
3023#if 0
3024            asm volatile(
3025                "li             t1,  8              \n\t"
3026                "vsetvli        t0, x0, e32, m1     \n\t"
3027                "vxor.vv        v28, v28, v28       \n\t"
3028                "vxor.vv        v29, v29, v29       \n\t"
3029                "vxor.vv        v30, v30, v30       \n\t"
3030                "vxor.vv        v31, v31, v31       \n\t"
3031                "mv             t4, %[BK]           \n\t"
3032
3033                ".align 4                           \n\t"
3034                "BLK_LOOP%=:                        \n\t"
3035                // load scale A
3036                "flw            fa0, (%[A])         \n\t"
3037                "flw            fa1, 4(%[A])        \n\t"
3038                "flw            fa2, 8(%[A])        \n\t"
3039                "flw            fa3, 12(%[A])       \n\t"
3040                "addi           %[A], %[A], 16      \n\t"
3041
3042                // load scale B
3043                "vsetvli        t0, x0, e16, mf2    \n\t"
3044                "vle16.v        v12, (%[B])         \n\t"
3045                "addi           %[B], %[B], 64      \n\t"
3046                "vfwcvt.f.f.v   v14, v12            \n\t"
3047
3048                "vsetivli       t0, 4, e16, mf2     \n\t"
3049                "vle16.v        v8, (%[A])          \n\t"  // asum
3050                "addi           %[A], %[A], 8       \n\t"
3051                "vwmul.vx       v10, v8, t1         \n\t"  // 8*asum
3052
3053                "vsetvli        t0, x0, e8, m1      \n\t"
3054                "vl1r.v         v0, (%[A])          \n\t"
3055                "addi           %[A], %[A], 128     \n\t"  // 4*32@i8
3056                "vl4r.v         v4, (%[B])          \n\t"  // 32*32@i4
3057                "addi           %[B], %[B], 512     \n\t"
3058                "vsrl.vi        v1, v0, 4           \n\t"
3059                "vnpack4.vv     v12, v0, v1, 3      \n\t"  // A low  u4
3060                "vupack.vv      v2, v12, v12, 2     \n\t"
3061
3062                // init the accumu to asum * zp
3063                "vsetvli        t0, x0, e32, m1     \n\t"
3064                "vxor.vv        v16, v16, v16       \n\t"
3065                "vxor.vv        v18, v16, v16       \n\t"
3066                "vxor.vv        v20, v16, v16       \n\t"
3067                "vxor.vv        v22, v16, v16       \n\t"
3068
3069                // i4 * i4 vmadot
3070                "vsetvli        t0, x0, e32, m1     \n\t"
3071                "vmadotsu       v16, v3, v4, i4     \n\t"   // high 4
3072                "vmadotsu       v18, v3, v5, i4     \n\t"
3073                "vmadotsu       v20, v3, v6, i4     \n\t"
3074                "vmadotsu       v22, v3, v7, i4     \n\t"
3075                "vsll.vi        v16, v16, 4         \n\t"
3076                "vsll.vi        v18, v18, 4         \n\t"
3077                "vsll.vi        v20, v20, 4         \n\t"
3078                "vsll.vi        v22, v22, 4         \n\t"
3079                "vmadotu        v16, v2, v4, i4     \n\t"   // low 4
3080                "vmadotu        v18, v2, v5, i4     \n\t"
3081                "vmadotu        v20, v2, v6, i4     \n\t"
3082                "vmadotu        v22, v2, v7, i4     \n\t"
3083
3084                "vpack.vv       v0, v16, v18, 2     \n\t"
3085                "vpack.vv       v2, v20, v22, 2     \n\t"
3086                "vpack.vv       v16, v0, v2, 3      \n\t"
3087                "vpack.vv       v18, v1, v3, 3      \n\t"
3088
3089                "vrgather.vi    v0, v10, 0          \n\t"
3090                "vrgather.vi    v1, v10, 1          \n\t"
3091                "vrgather.vi    v2, v10, 2          \n\t"
3092                "vrgather.vi    v3, v10, 3          \n\t"
3093
3094                "vadd.vv        v16, v16, v0        \n\t"
3095                "vadd.vv        v17, v17, v1        \n\t"
3096                "vadd.vv        v18, v18, v2        \n\t"
3097                "vadd.vv        v19, v19, v3        \n\t"
3098
3099                "vfcvt.f.x.v    v16, v16            \n\t"
3100                "vfcvt.f.x.v    v17, v17            \n\t"
3101                "vfcvt.f.x.v    v18, v18            \n\t"
3102                "vfcvt.f.x.v    v19, v19            \n\t"
3103
3104                // mul scale
3105                "vfmul.vv       v16, v16, v14       \n\t"
3106                "vfmul.vv       v17, v17, v14       \n\t"
3107                "vfmul.vv       v18, v18, v14       \n\t"
3108                "vfmul.vv       v19, v19, v14       \n\t"
3109
3110                "addi           t4, t4, -1          \n\t"
3111                "vfmacc.vf      v28, fa0, v16       \n\t"
3112                "vfmacc.vf      v29, fa1, v17       \n\t"
3113                "vfmacc.vf      v30, fa2, v18       \n\t"
3114                "vfmacc.vf      v31, fa3, v19       \n\t"
3115
3116                "bgtz           t4, BLK_LOOP%=      \n\t"
3117
3118                // save
3119                "vsetvli        t0, x0, e32, m1     \n\t"
3120                "add            t2, %[LDC], %[DST]  \n\t"
3121                "vse32.v        v28, (%[DST])       \n\t"
3122                "add            t3, %[LDC], t2      \n\t"
3123                "vse32.v        v29, (t2)           \n\t"
3124                "add            t2, %[LDC], t3      \n\t"
3125                "vse32.v        v30, (t3)           \n\t"
3126                "vse32.v        v31, (t2)           \n\t"
3127                : [A] "+r"(a_data), [B] "+r"(b_data)
3128                : [DST] "r"(dst_c), [LDC] "r"(ldc*4), [BK] "r"(k_blks)
3129                : "t0", "t1", "t2", "t3", "t4", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11",
3130                  "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25",
3131                  "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "fa3");
3132#else
3133            asm volatile(
3134                "vsetvli        t0, x0, e16, m1         \n\t"
3135                "vxor.vv        v28, v28, v28           \n\t"
3136                "vxor.vv        v29, v29, v29           \n\t"
3137                "vxor.vv        v30, v30, v30           \n\t"
3138                "vxor.vv        v31, v31, v31           \n\t"
3139                "vmv.v.i        v0, 1                   \n\t"  // init the scale
3140                "vsll.vi        v1, v0, 4               \n\t"
3141                "vfcvt.f.x.v    v0, v0                  \n\t"
3142                "vfcvt.f.x.v    v1, v1                  \n\t"
3143                "mv             t4, %[BK]               \n\t"
3144
3145                ".align 4                               \n\t"
3146                "BLK_LOOP%=:                            \n\t"
3147                // load scale A
3148                "flw            fa0, (%[A])             \n\t"
3149                "flw            fa1, 4(%[A])            \n\t"
3150                "flw            fa2, 8(%[A])            \n\t"
3151                "flw            fa3, 12(%[A])           \n\t"
3152                "addi           %[A], %[A], 16          \n\t"
3153
3154                // load scale B
3155                "vsetvli        t0, x0, e16, mf2        \n\t"
3156                "vle16.v        v12, (%[B])             \n\t"
3157                "addi           %[B], %[B], 64          \n\t"
3158                "vsetvli        t0, x0, e16, m1         \n\t"
3159                "vpack.vv       v14, v12, v12, 3        \n\t"
3160
3161                "vsetivli       t0, 4, e16, mf2         \n\t"
3162                "vle16.v        v8, (%[A])              \n\t"  // asum
3163                "addi           %[A], %[A], 8           \n\t"
3164                "vsll.vi        v8, v8, 3               \n\t"  // asum * 8
3165                "vfcvt.f.x.v    v9, v8                  \n\t"
3166                "vsetvli        t0, x0, e64, m1         \n\t"
3167                "vrgather.vi    v10, v9, 0              \n\t"
3168
3169                "vsetvli        t0, x0, e8, m1          \n\t"
3170                "vl1r.v         v16, (%[A])             \n\t"
3171                "addi           %[A], %[A], 128         \n\t"  // 4*32@i8
3172                "vl4r.v         v4, (%[B])              \n\t"  // 32*32@i4
3173                "addi           %[B], %[B], 512         \n\t"
3174                "vsrl.vi        v17, v16, 4             \n\t"
3175                "vnpack4.vv     v12, v16, v17, 3        \n\t"  // A low  u4
3176                "vupack.vv      v2, v12, v12, 2         \n\t"
3177
3178                // init the accumu to asum * zp
3179                "vsetvli        t0, x0, e16, m1         \n\t"
3180                "vpack.vv       v16, v10, v10,0         \n\t"
3181                "vsetvli        t0, x0, e32, m1         \n\t"
3182                "vpack.vv       v20, v16, v16,0         \n\t"
3183                "vsetvli        t0, x0, e64, m1         \n\t"
3184                "vpack.vv       v18, v20, v20, 0        \n\t"
3185                "vor.vv         v20, v18, v18           \n\t"
3186                "vor.vv         v21, v18, v18           \n\t"
3187
3188                // i4 * i4 vmadot
3189                "vsetvli        t0, x0, e16, m1         \n\t"
3190                "vmadotsu.hp    v18, v3, v4, v1, 0, i4  \n\t"  // high 4
3191                "vmadotsu.hp    v19, v3, v5, v1, 0, i4  \n\t"
3192                "vmadotsu.hp    v20, v3, v6, v1, 0, i4  \n\t"
3193                "vmadotsu.hp    v21, v3, v7, v1, 0, i4  \n\t"
3194                "vmadotu.hp     v18, v2, v4, v0, 0, i4  \n\t"  // low 4
3195                "vmadotu.hp     v19, v2, v5, v0, 0, i4  \n\t"
3196                "vmadotu.hp     v20, v2, v6, v0, 0, i4  \n\t"
3197                "vmadotu.hp     v21, v2, v7, v0, 0, i4  \n\t"
3198
3199                "vpack.vv       v8, v18, v19, 1         \n\t"
3200                "vpack.vv       v12, v20, v21, 1        \n\t"
3201                "vpack.vv       v20, v8, v12, 2         \n\t"
3202
3203                "vfwmul.vv      v16, v20, v14           \n\t"
3204                "vfwmul.vv      v18, v21, v14           \n\t"
3205
3206                "vsetvli        t0, x0, e32, m1         \n\t"
3207
3208                "addi           t4, t4, -1              \n\t"
3209                "vfmacc.vf      v28, fa0, v16           \n\t"
3210                "vfmacc.vf      v29, fa1, v17           \n\t"
3211                "vfmacc.vf      v30, fa2, v18           \n\t"
3212                "vfmacc.vf      v31, fa3, v19           \n\t"
3213
3214                "bgtz           t4, BLK_LOOP%=          \n\t"
3215
3216                // save
3217                "vsetvli        t0, x0, e32, m1         \n\t"
3218                "add            t2, %[LDC], %[DST]      \n\t"
3219                "vse32.v        v28, (%[DST])           \n\t"
3220                "add            t3, %[LDC], t2          \n\t"
3221                "vse32.v        v29, (t2)               \n\t"
3222                "add            t2, %[LDC], t3          \n\t"
3223                "vse32.v        v30, (t3)               \n\t"
3224                "vse32.v        v31, (t2)               \n\t"
3225                : [A] "+r"(a_data), [B] "+r"(b_data)
3226                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [BK] "r"(k_blks)
3227                : "t0", "t1", "t2", "t3", "t4", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10",
3228                  "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24",
3229                  "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "fa3");
3230#endif
3231        }
3232    } else {
3233        for (size_t ni = 0; ni < count_n; ni += 32) {
3234            uint8_t * b_data = (uint8_t *) quant_b_data + ni * b_data_stride;
3235            int8_t *  a_data = (int8_t *) quant_a_ptr;
3236            float *   dst_c  = c_ptr + ni;
3237
3238            asm volatile(
3239                "li             t1,  8          \n\t"
3240                "vsetvli        t0, x0, e32, m1 \n\t"
3241                "vxor.vv        v28, v28, v28   \n\t"
3242                "vxor.vv        v29, v29, v29   \n\t"
3243                "vxor.vv        v30, v30, v30   \n\t"
3244                "vxor.vv        v31, v31, v31   \n\t"
3245                "mv             t4, %[BK]       \n\t"
3246
3247                ".align 4                        \n\t"
3248                "BLK_LOOP%=:                     \n\t"
3249                // load scale A
3250                "flw            fa0, (%[A])     \n\t"
3251                "flw            fa1, 4(%[A])    \n\t"
3252                "flw            fa2, 8(%[A])    \n\t"
3253                "flw            fa3, 12(%[A])   \n\t"
3254                "addi           %[A], %[A], 16  \n\t"
3255
3256                // load scale B
3257                "vsetvli        t0, x0, e16, mf2\n\t"
3258                "vle16.v        v12, (%[B])     \n\t"
3259                "addi           %[B], %[B], 64  \n\t"
3260                "vfwcvt.f.f.v   v14, v12        \n\t"
3261
3262                // load zp
3263                "vsetvli        t0, x0, e8, mf4 \n\t"
3264                "vle8.v         v8, (%[B])      \n\t"
3265                "addi           %[B], %[B], 32  \n\t"
3266                "vwaddu.vx      v10, v8, x0     \n\t"
3267
3268                // load a sum
3269                "lh             s1, (%[A])      \n\t"
3270                "lh             s2, 2(%[A])     \n\t"
3271                "lh             s3, 4(%[A])     \n\t"
3272                "lh             s4, 6(%[A])     \n\t"
3273                "addi           %[A], %[A], 8   \n\t"
3274
3275                "vsetvli        t0, x0, e8, m1  \n\t"
3276                "vl1r.v         v0, (%[A])      \n\t"
3277                "addi           %[A], %[A], 128 \n\t"  // 4*32@i8
3278                "vl4r.v         v4, (%[B])      \n\t"  // 32*32@i4
3279                "addi           %[B], %[B], 512 \n\t"
3280                "vsrl.vi        v1, v0, 4       \n\t"
3281                "vnpack4.vv     v12, v0, v1, 3  \n\t"  // A low  u4
3282                "vupack.vv      v2, v12, v12, 2 \n\t"
3283
3284                // init the accumu to asum * zp
3285                "vsetvli        t0, x0, e32, m1 \n\t"
3286                "vxor.vv        v16, v16, v16   \n\t"
3287                "vxor.vv        v18, v16, v16   \n\t"
3288                "vxor.vv        v20, v16, v16   \n\t"
3289                "vxor.vv        v22, v16, v16   \n\t"
3290
3291                // i4 * i4 vmadot
3292                "vsetvli        t0, x0, e32, m1 \n\t"
3293                "vmadotsu       v16, v3, v4, i4 \n\t"  // high 4
3294                "vmadotsu       v18, v3, v5, i4 \n\t"
3295                "vmadotsu       v20, v3, v6, i4 \n\t"
3296                "vmadotsu       v22, v3, v7, i4 \n\t"
3297                "vsll.vi        v16, v16, 4     \n\t"
3298                "vsll.vi        v18, v18, 4     \n\t"
3299                "vsll.vi        v20, v20, 4     \n\t"
3300                "vsll.vi        v22, v22, 4     \n\t"
3301                "vmadotu        v16, v2, v4, i4 \n\t"  // low 4
3302                "vmadotu        v18, v2, v5, i4 \n\t"
3303                "vmadotu        v20, v2, v6, i4 \n\t"
3304                "vmadotu        v22, v2, v7, i4 \n\t"
3305
3306                "vpack.vv       v0, v16, v18, 2 \n\t"
3307                "vpack.vv       v2, v20, v22, 2 \n\t"
3308                "vpack.vv       v16, v0, v2, 3  \n\t"
3309                "vpack.vv       v18, v1, v3, 3  \n\t"
3310
3311                "vsetvli        t0, x0, e16, m1 \n\t"
3312                "vwmul.vx       v0, v10, s1     \n\t"
3313                "vwmul.vx       v2, v10, s2     \n\t"
3314                "vwmul.vx       v4, v10, s3     \n\t"
3315                "vwmul.vx       v6, v10, s4     \n\t"
3316
3317                "vsetvli        t0, x0, e32, m1 \n\t"
3318                "vadd.vv        v16, v16, v0    \n\t"
3319                "vadd.vv        v17, v17, v2    \n\t"
3320                "vadd.vv        v18, v18, v4    \n\t"
3321                "vadd.vv        v19, v19, v6    \n\t"
3322
3323                "vfcvt.f.x.v    v16, v16        \n\t"
3324                "vfcvt.f.x.v    v17, v17        \n\t"
3325                "vfcvt.f.x.v    v18, v18        \n\t"
3326                "vfcvt.f.x.v    v19, v19        \n\t"
3327
3328                // mul scale
3329                "vfmul.vv       v16, v16, v14   \n\t"
3330                "vfmul.vv       v17, v17, v14   \n\t"
3331                "vfmul.vv       v18, v18, v14   \n\t"
3332                "vfmul.vv       v19, v19, v14   \n\t"
3333
3334                "addi           t4, t4, -1      \n\t"
3335                "vfmacc.vf      v28, fa0, v16   \n\t"
3336                "vfmacc.vf      v29, fa1, v17   \n\t"
3337                "vfmacc.vf      v30, fa2, v18   \n\t"
3338                "vfmacc.vf      v31, fa3, v19   \n\t"
3339
3340                "bgtz           t4, BLK_LOOP%=  \n\t"
3341
3342                // save
3343                "vsetvli        t0, x0, e32, m1 \n\t"
3344                "add            t2, %[LDC], %[DST]\n\t"
3345                "vse32.v        v28, (%[DST])   \n\t"
3346                "add            t3, %[LDC], t2  \n\t"
3347                "vse32.v        v29, (t2)       \n\t"
3348                "add            t2, %[LDC], t3  \n\t"
3349                "vse32.v        v30, (t3)       \n\t"
3350                "vse32.v        v31, (t2)       \n\t"
3351                : [A] "+r"(a_data), [B] "+r"(b_data)
3352                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [BK] "r"(k_blks)
3353                : "t0", "t1", "t2", "t3", "t4", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10",
3354                  "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24",
3355                  "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "fa3", "s1", "s2", "s3", "s4");
3356        }
3357    }
3358}
3359
3360void gemm_kernel_i8i4_hp_m4(size_t          blk_len,
3361                            const uint8_t * quant_a_ptr,
3362                            const uint8_t * quant_b_data,
3363                            const uint8_t * quant_b_zp,
3364                            float *         c_ptr,
3365                            size_t          count_m,
3366                            size_t          count_n,
3367                            size_t          k_blks,
3368                            size_t          ldc) {
3369    constexpr size_t NB_COLS                = 32;
3370    constexpr size_t K_SUBBLKS_PER_SUPERBLK = 8;
3371    constexpr size_t K_SUBBLK_LEN           = 32;
3372
3373    struct block_q4_0x32_layout {
3374        _Float16 d[NB_COLS];
3375        uint8_t  qs[16 * NB_COLS];
3376    };
3377
3378    GGML_ASSERT(blk_len == 256);
3379    GGML_ASSERT(count_m >= 4);
3380
3381    // Contract:
3382    // - computes a 4-row x 32-col tile per inner invocation
3383    // - A is q8 HP packed in m4 layout, one logical K256 block at a time
3384    // - B is q4 HP packed in N32 tiles, optionally with a separate zp area
3385    // - tail-N is currently not handled here; the caller must provide full N32 tiles
3386
3387    const size_t b_superblk_stride = sizeof(block_q4_0x32_layout) * K_SUBBLKS_PER_SUPERBLK +
3388                                     (quant_b_zp ? NB_COLS * K_SUBBLKS_PER_SUPERBLK * sizeof(uint8_t) : 0);
3389    const size_t b_tile_stride       = k_blks * b_superblk_stride;
3390    const size_t a_nrow_block_stride = q8_hp_blk_size(blk_len, true, true) * 4;
3391    const size_t a_subblk_stride     = q8_hp_blk_size(K_SUBBLK_LEN, false, false) * 4;
3392
3393    if (quant_b_zp != nullptr) {
3394        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
3395            const size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
3396            if (nb_real != NB_COLS) {
3397                break;
3398            }
3399
3400            uint8_t * b_tile_base = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_tile_stride;
3401            uint8_t * a_block     = (uint8_t *) quant_a_ptr;
3402            float *   dst_c       = c_ptr + ni;
3403
3404            // Data layout summary for the with-zp path.
3405            //
3406            // A: M4 x K256 q8 HP block
3407            //   - split into 8 x K32 subblocks
3408            //   - each K32 subblock is 136B:
3409            //       8B   = 4 x fp16 row scales
3410            //       128B = 4 x int8[32] row payloads
3411            //   - trailer after 8 subblocks is 72B:
3412            //       4 rows x fp16[8] a_sum values, indexed as [row][ksi]
3413            //       4 rows x fp16 scale_avg tail
3414            //
3415            // B: N32 x K256 q4 HP block with explicit zp area
3416            //   - each K32 subblock is 576B:
3417            //       64B  = fp16 scale[32]
3418            //       512B = packed q4 payload for 32 columns x 32 k-elements
3419            //   - zp is stored separately, not interleaved with the 576B payload block
3420            //   - one K256 superblock is laid out as:
3421            //       8 x (scale + qs) blocks = 4608B
3422            //       8 x zp[32]              =  256B
3423            //
3424            // C: 4 rows x 32 fp32 outputs
3425            //
3426            // ASM pointer convention:
3427            //   - t6: current A K32 subblock base
3428            //   - t2: current A a_sum base for this ksi
3429            //         row1/row2/row3 are at +16/+32/+48 bytes
3430            //   - s5: current B (scale + qs) K32 subblock base
3431            //   - s6: current B zp[32] base for this ksi
3432            //
3433            // Loop progression:
3434            //   - per ksi: A += 136, a_sum += 2, B_data += 576, B_zp += 32
3435            //   - per ki : skip the 72B A trailer and advance B to the next 4864B superblock
3436
3437            const _Float16 hp_scale_16   = (_Float16) 16.0f;
3438            const _Float16 hp_scale_1    = (_Float16) 1.0f;
3439            const _Float16 hp_scale_0125 = (_Float16) 0.125f;
3440
3441            // VPR grouping used below:
3442            // - v4-v7   : B q4 payload for N32 split as 4 x N8 groups
3443            // - v8/v10  : zp u8 / widened fp16
3444            // - v12     : B fp16 scale[32]
3445            // - v14-v15 : packed (Bscale * Ascale) for rows [0,1] / [2,3]
3446            // - v16-v19 : temporary per-row scaled B scales
3447            // - v28-v31 : final fp32 accumulators for rows 0..3
3448
3449            asm volatile(
3450                "mv             t5, %[BK]                 \n\t"
3451                "mv             t6, %[A]                  \n\t"
3452                "mv             s5, %[B]                  \n\t"
3453                "vsetvli        t0, x0, e32, m1           \n\t"
3454                "vxor.vv        v28, v28, v28             \n\t"
3455                "vxor.vv        v29, v29, v29             \n\t"
3456                "vxor.vv        v30, v30, v30             \n\t"
3457                "vxor.vv        v31, v31, v31             \n\t"
3458                "li             t4, 8                     \n\t"
3459                "li             t1, 4608                  \n\t"
3460                "addi           t2, t6, 1088              \n\t"  // 8 * 136B A K32 subblocks, a_sum trailer starts here
3461                "add            s6, s5, t1                \n\t"  // 8 * 576B B(scale+qs), zp area starts here
3462
3463                ".align 4                                 \n\t"
3464                "_BLK_LPST%=:                             \n\t"
3465                "flh            fa1, 64(t2)               \n\t"  // a_scale_avg_row[0]
3466                "vsetvli        t0, x0, e32, m1           \n\t"
3467                "vxor.vv        v18, v30, v30             \n\t"
3468                "vxor.vv        v19, v31, v31             \n\t"
3469                "vxor.vv        v20, v30, v30             \n\t"
3470                "vxor.vv        v21, v31, v31             \n\t"
3471                "_KsubBLK_LPST%=:                         \n\t"
3472                // load first subblock scales for 4 rows
3473                "flh            fa0,   0(t6)              \n\t"  // ascale_fp16
3474
3475                // load B fp16 scales[32]
3476                "vsetvli        t0, x0, e16, mf2          \n\t"
3477                "vle16.v        v12, (s5)                 \n\t"
3478
3479                // load Bzp[32] for the current ksi from the dedicated zp area
3480                "vsetvli        t0, x0, e8, mf4           \n\t"
3481                "vle8.v         v8, (s6)                  \n\t"
3482
3483                "fmul.h         fa2, fa0, %[HP16]         \n\t"
3484                "vfwcvt.f.xu.v  v10, v8                   \n\t"  // uint8 -> fp16
3485
3486                "vsetvli        t0, x0, e16, mf2          \n\t"
3487                "vfmul.vf       v16, v12, fa0             \n\t"  // row0: Bscale * Ascale
3488                "vfmul.vf       v17, v12, fa2             \n\t"
3489
3490                // load a_sum[row][ksi] from the trailer; t2 points to row0[ksi]
3491                "flh            ft1, 0(t2)                \n\t"
3492                "flh            ft2, 16(t2)               \n\t"
3493                "flh            ft3, 32(t2)               \n\t"
3494                "flh            ft4, 48(t2)               \n\t"
3495
3496                "fmul.h         ft1, ft1, %[HP0125]       \n\t"
3497                "fmul.h         ft2, ft2, %[HP0125]       \n\t"
3498                "fmul.h         ft3, ft3, %[HP0125]       \n\t"
3499                "fmul.h         ft4, ft4, %[HP0125]       \n\t"
3500
3501                // load A payload from current K32 subblock and B q4 payload from current 576B block
3502                "addi           t3, t6, 8                 \n\t"
3503                "vsetvli        t0, x0, e8, m1            \n\t"
3504                "vl1r.v         v0, (t3)                  \n\t"  //A
3505                "addi           t3, s5, 64                \n\t"
3506                "vl4r.v         v4, (t3)                  \n\t"  //B
3507
3508                "vsetvli        t0, x0, e8, m1            \n\t"
3509                "vsrl.vi        v1, v0, 4                 \n\t"
3510                "vnpack4.vv     v12, v0, v1, 3            \n\t"
3511                "vpack.vv       v0, v17, v16, 3           \n\t"
3512                "vupack.vv      v2, v12, v12, 2           \n\t"
3513
3514                "vsetvli        t0, x0, e16, mf2          \n\t"  // mf2 -> mf2
3515                "vfmul.vv       v10, v10, v16             \n\t"  // zp * ascale * bscale; fp16*fp16
3516
3517                "vsetvli        t0, x0, e16, mf2          \n\t"  // mf2 -> m1
3518                "vfmul.vf       v12, v10, ft1             \n\t"  // zp(1:n)* abscale * asum_m0; fp16*fp16
3519                "vfmul.vf       v13, v10, ft2             \n\t"  // zp(1:n)* abscale * asum_m1; fp16*fp16
3520                "vfmul.vf       v24, v10, ft3             \n\t"  // zp(1:n)* abscale * asum_m2; fp16*fp16
3521                "vfmul.vf       v25, v10, ft4             \n\t"  // zp(1:n)* abscale * asum_m3; fp16*fp16
3522
3523                "vsetvli        t0, x0, e16, mf2           \n\t"
3524                "vfwmacc.vf     v28, fa1, v12             \n\t"  // row0/1 accum += dot * packed scale
3525                "vfwmacc.vf     v29, fa1, v13             \n\t"
3526                "vfwmacc.vf     v30, fa1, v24             \n\t"
3527                "vfwmacc.vf     v31, fa1, v25             \n\t"
3528
3529                "vsetvli        t0, x0, e32, m1           \n\t"
3530                "vmadotsu.hp    v18, v3, v4, v0, 0, i4    \n\t"  //lo4;n0n7
3531                "vmadotsu.hp    v19, v3, v5, v0, 1, i4    \n\t"  //lo4;n8n15
3532                "vmadotsu.hp    v20, v3, v6, v0, 2, i4    \n\t"  //lo4;n16n23
3533                "vmadotsu.hp    v21, v3, v7, v0, 3, i4    \n\t"  //lo4;n24n31
3534                "vmadotu.hp     v18, v2, v4, v0, 4, i4    \n\t"  //hi4;n0n7
3535                "vmadotu.hp     v19, v2, v5, v0, 5, i4    \n\t"  //hi4;n8n15
3536                "vmadotu.hp     v20, v2, v6, v0, 6, i4    \n\t"  //hi4;n16n23
3537                "vmadotu.hp     v21, v2, v7, v0, 7, i4    \n\t"  //hi4;n24n31
3538
3539                "addi           t4, t4, -1                \n\t"
3540                "addi           t6, t6, 8+128             \n\t"  // next A K32 subblock
3541                "addi           t2, t2, 2                 \n\t"  // next ksi entry in each a_sum row
3542                "addi           s5, s5, 64+512            \n\t"  // next B (scale + qs) K32 block
3543                "addi           s6, s6, 32                \n\t"  // next zp[32]
3544                "bgtz           t4, _KsubBLK_LPST%=       \n\t"
3545
3546                "vsetvli        t0, x0, e16, m1           \n\t"
3547                "vpack.vv       v8, v18, v19, 1           \n\t"  // 128(16*8)->256(16*16)
3548                "vpack.vv       v12, v20, v21, 1          \n\t"
3549                "vpack.vv       v26, v8, v12, 2           \n\t"  // 256(16*16)->512(16*32)
3550
3551                "vsetvli        t0, x0, e16, m1           \n\t"
3552                "vfwmacc.vf     v28, fa1, v26             \n\t"  // row0/1 accum += dot * packed scale
3553                "vfwmacc.vf     v30, fa1, v27             \n\t"
3554
3555                "li             t4, 8                     \n\t"
3556                "addi           t5, t5, -1                \n\t"
3557                "addi           t6, t6, 72                \n\t"  // skip A trailer after 8 subblocks and scale_avg tail
3558                "mv             s5, s6                    \n\t"  // s6 already points to next B superblock base
3559                "addi           t2, t6, 1088              \n\t"  // 8 * 136B A K32 subblocks, a_sum trailer starts here
3560                "add            s6, s5, t1                \n\t"  // 8 * 576B B(scale+qs), zp area starts here
3561                "bgtz           t5, _BLK_LPST%=           \n\t"
3562
3563                "_BLK_LPND%=:                             \n\t"
3564                "vsetvli        t0, x0, e32, m1           \n\t"
3565                "add            t2, %[LDC], %[DST]        \n\t"
3566                "vse32.v        v28, (%[DST])             \n\t"
3567                "add            t3, %[LDC], t2            \n\t"
3568                "vse32.v        v29, (t2)                 \n\t"
3569                "add            t2, %[LDC], t3            \n\t"
3570                "vse32.v        v30, (t3)                 \n\t"
3571                "vse32.v        v31, (t2)                 \n\t"
3572                : [A] "+r"(a_block), [B] "+r"(b_tile_base)
3573                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [BK] "r"(k_blks), [HP16] "f"(hp_scale_16),
3574                  [HP1] "f"(hp_scale_1), [HP0125] "f"(hp_scale_0125)
3575                : "t0", "t1", "t2", "t3", "t4", "t5", "t6", "s5", "s6", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
3576                  "v8", "v10", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v24",
3577                  "v25", "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "ft1", "ft2", "ft3", "ft4",
3578                  "memory");
3579        }
3580        return;
3581    } else {
3582        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
3583            const size_t nb_real = std::min<size_t>(NB_COLS, count_n - ni);
3584            if (nb_real != NB_COLS) {
3585                break;
3586            }
3587
3588            uint8_t * b_tile_base = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_tile_stride;
3589            uint8_t * a_block     = (uint8_t *) quant_a_ptr;
3590            float *   dst_c       = c_ptr + ni;
3591
3592            // Data layout summary for the no-zp path.
3593            //
3594            // A layout is identical to the with-zp branch.
3595            //
3596            // B: N32 x K256 q4 HP block without explicit zp storage
3597            //   - each K32 subblock is still 576B:
3598            //       64B  = fp16 scale[32]
3599            //       512B = packed q4 payload
3600            //   - zp is implicit and treated as a constant value 8 in the kernel
3601            //   - one K256 superblock therefore contains only:
3602            //       8 x (scale + qs) blocks = 4608B
3603            //
3604            // C: 4 rows x 32 fp32 outputs
3605            //
3606            // ASM pointer convention:
3607            //   - t6: current A K32 subblock base
3608            //   - t2: current A a_sum base for this ksi
3609            //   - s5: current B (scale + qs) K32 subblock base
3610            //
3611            // Loop progression:
3612            //   - per ksi: A += 136, a_sum += 2, B_data += 576
3613            //   - per ki : skip the 72B A trailer and advance B to the next 4608B superblock
3614
3615            const _Float16 hp_scale_16 = (_Float16) 16.0f;
3616            const _Float16 hp_scale_1  = (_Float16) 1.0f;
3617
3618            // VPR grouping used below matches the with-zp path:
3619            // - v4-v7   : B q4 payload for N32 split as 4 x N8 groups
3620            // - v8/v10  : implicit zp lane / widened fp16
3621            // - v12     : B fp16 scale[32]
3622            // - v14-v15 : packed (Bscale * Ascale) for rows [0,1] / [2,3]
3623            // - v16-v19 : temporary per-row scaled B scales
3624            // - v28-v31 : final fp32 accumulators for rows 0..3
3625
3626            asm volatile(
3627                "mv             t5, %[BK]                 \n\t"
3628                "mv             t6, %[A]                  \n\t"
3629                "mv             s5, %[B]                  \n\t"
3630                "vsetvli        t0, x0, e32, m1           \n\t"
3631                "vxor.vv        v28, v28, v28             \n\t"
3632                "vxor.vv        v29, v29, v29             \n\t"
3633                "vxor.vv        v30, v30, v30             \n\t"
3634                "vxor.vv        v31, v31, v31             \n\t"
3635                "li             t4, 8                     \n\t"
3636                "addi           t2, t6, 1088              \n\t"  // 8 * 136B A K32 subblocks, a_sum trailer starts here
3637
3638                ".align 4                                 \n\t"
3639                "_BLK_LPST%=:                             \n\t"
3640                "flh            fa1, 64(t2)               \n\t"  // a_scale_avg_row[0]
3641                "vsetvli        t0, x0, e32, m1           \n\t"
3642                "vxor.vv        v18, v30, v30             \n\t"
3643                "vxor.vv        v19, v31, v31             \n\t"
3644                "vxor.vv        v20, v30, v30             \n\t"
3645                "vxor.vv        v21, v31, v31             \n\t"
3646                "_KsubBLK_LPST%=:                         \n\t"
3647                // load first subblock scales for 4 rows
3648                "flh            fa0,   0(t6)              \n\t"  // ascale_fp16
3649
3650                // load B fp16 scales[32]
3651                "vsetvli        t0, x0, e16, mf2          \n\t"
3652                "vle16.v        v12, (s5)                 \n\t"
3653
3654                "fmul.h         fa2, fa0, %[HP16]         \n\t"
3655
3656                "vsetvli        t0, x0, e16, mf2          \n\t"
3657                "vfmul.vf       v16, v12, fa0             \n\t"  // row0: Bscale * Ascale
3658                "vfmul.vf       v17, v12, fa2             \n\t"
3659
3660                // load a_sum[row][ksi] from the trailer; t2 points to row0[ksi]
3661                "flh            ft1, 0(t2)                \n\t"
3662                "flh            ft2, 16(t2)               \n\t"
3663                "flh            ft3, 32(t2)               \n\t"
3664                "flh            ft4, 48(t2)               \n\t"
3665
3666                // load A payload from current K32 subblock and B q4 payload from current 576B block
3667                "addi           t3, t6, 8                 \n\t"
3668                "vsetvli        t0, x0, e8, m1            \n\t"
3669                "vl1r.v         v0, (t3)                  \n\t"  //A
3670                "addi           t3, s5, 64                \n\t"
3671                "vl4r.v         v4, (t3)                  \n\t"  //B
3672
3673                "vsetvli        t0, x0, e8, m1            \n\t"
3674                "vsrl.vi        v1, v0, 4                 \n\t"
3675                "vnpack4.vv     v12, v0, v1, 3            \n\t"
3676                "vpack.vv       v0, v17, v16, 3           \n\t"
3677                "vupack.vv      v2, v12, v12, 2           \n\t"
3678
3679                "vsetvli        t0, x0, e16, mf2          \n\t"  // mf2 -> m1
3680                "vfmul.vf       v12, v16, ft1             \n\t"  // zp(1:n)* abscale * asum_m0; fp16*fp16
3681                "vfmul.vf       v13, v16, ft2             \n\t"  // zp(1:n)* abscale * asum_m1; fp16*fp16
3682                "vfmul.vf       v24, v16, ft3             \n\t"  // zp(1:n)* abscale * asum_m2; fp16*fp16
3683                "vfmul.vf       v25, v16, ft4             \n\t"  // zp(1:n)* abscale * asum_m3; fp16*fp16
3684
3685                "vsetvli        t0, x0, e16, mf2          \n\t"
3686                "vfwmacc.vf     v28, fa1, v12             \n\t"
3687                "vfwmacc.vf     v29, fa1, v13             \n\t"
3688                "vfwmacc.vf     v30, fa1, v24             \n\t"
3689                "vfwmacc.vf     v31, fa1, v25             \n\t"
3690
3691                "vsetvli        t0, x0, e32, m1           \n\t"
3692                "vmadotsu.hp    v18, v3, v4, v0, 0, i4    \n\t"  //lo4;n0n7
3693                "vmadotsu.hp    v19, v3, v5, v0, 1, i4    \n\t"  //lo4;n8n15
3694                "vmadotsu.hp    v20, v3, v6, v0, 2, i4    \n\t"  //lo4;n16n23
3695                "vmadotsu.hp    v21, v3, v7, v0, 3, i4    \n\t"  //lo4;n24n31
3696                "vmadotu.hp     v18, v2, v4, v0, 4, i4    \n\t"  //hi4;n0n7
3697                "vmadotu.hp     v19, v2, v5, v0, 5, i4    \n\t"  //hi4;n8n15
3698                "vmadotu.hp     v20, v2, v6, v0, 6, i4    \n\t"  //hi4;n16n23
3699                "vmadotu.hp     v21, v2, v7, v0, 7, i4    \n\t"  //hi4;n24n31
3700
3701                "addi           t4, t4, -1                \n\t"
3702
3703                "addi           t6, t6, 8+128             \n\t"  // next A K32 subblock
3704                "addi           t2, t2, 2                 \n\t"  // next ksi entry in each a_sum row
3705                "addi           s5, s5, 64+512            \n\t"  // next B (scale + qs) K32 block
3706                "bgtz           t4, _KsubBLK_LPST%=       \n\t"
3707
3708                "vsetvli        t0, x0, e16, m1           \n\t"  //N32in1register
3709                "vpack.vv       v8, v18, v19, 1           \n\t"  // 128(16*8)->256(16*16)
3710                "vpack.vv       v12, v20, v21, 1          \n\t"
3711                "vpack.vv       v26, v8, v12, 2           \n\t"  // 256(16*16)->512(16*32)
3712
3713                "vsetvli        t0, x0, e16, m1           \n\t"
3714                "vfwmacc.vf     v28, fa1, v26             \n\t"  // row0/1 accum += dot * packed scale
3715                "vfwmacc.vf     v30, fa1, v27             \n\t"
3716
3717                "li             t4, 8                     \n\t"
3718                "addi           t5, t5, -1                \n\t"
3719                "addi           t6, t6, 72                \n\t"  // skip A trailer after 8 subblocks and scale_avg tail
3720                // s5 already points to next B superblock base
3721                "addi           t2, t6, 1088              \n\t"  // 8 * 136B A K32 subblocks, a_sum trailer starts here
3722                "bgtz           t5, _BLK_LPST%=           \n\t"
3723
3724                "_BLK_LPND%=:                             \n\t"
3725                "vsetvli        t0, x0, e32, m1           \n\t"
3726                "add            t2, %[LDC], %[DST]        \n\t"
3727                "vse32.v        v28, (%[DST])             \n\t"
3728                "add            t3, %[LDC], t2            \n\t"
3729                "vse32.v        v29, (t2)                 \n\t"
3730                "add            t2, %[LDC], t3            \n\t"
3731                "vse32.v        v30, (t3)                 \n\t"
3732                "vse32.v        v31, (t2)                 \n\t"
3733                : [A] "+r"(a_block), [B] "+r"(b_tile_base)
3734                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [BK] "r"(k_blks), [HP16] "f"(hp_scale_16), [HP1] "f"(hp_scale_1)
3735                : "t0", "t2", "t3", "t4", "t5", "t6", "s5", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v10",
3736                  "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v24", "v25", "v26",
3737                  "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "ft1", "ft2", "ft3", "ft4", "memory");
3738        }
3739        return;
3740    }
3741}
3742
3743void gemm_kernel_i8mxfp4_m1(size_t          blk_len,
3744                            const uint8_t * quant_a_ptr,
3745                            const uint8_t * quant_b_data,
3746                            const uint8_t * quant_b_zp,
3747                            float *         c_ptr,
3748                            size_t          count_m,
3749                            size_t          count_n,
3750                            size_t          k_blks,
3751                            size_t          ldc) {
3752    constexpr size_t NB_COLS = 32;
3753    constexpr size_t K_TILE  = 32;
3754    using blk_type           = nrow_block_mxfp4<NB_COLS>;
3755
3756    GGML_ASSERT(blk_len == K_TILE);
3757    GGML_ASSERT(count_m == 1);
3758    GGML_UNUSED(quant_b_zp);
3759
3760    const size_t a_blk_stride  = q8_blk_size(blk_len, true);
3761    const size_t b_blk_stride  = sizeof(blk_type);
3762    const size_t b_tile_stride = k_blks * b_blk_stride;
3763
3764    if (quant_b_zp == NULL) {
3765        for (size_t n = 0; n < count_n; n += 32) {
3766            size_t    nblks         = (count_n - n) > 32 ? 32 : count_n - n;
3767            // MXFP4 no-zp: per column per k-block stride = scale_e8m0(1B) + qs(16B) + qh(4B) = 21B
3768            uint8_t * QuantBDataPtr = (uint8_t *) quant_b_data +     //
3769                                      n * k_blks * (blk_len / 8) +   // qh sign/high-bit mask: n×k_blks×4
3770                                      n * k_blks * blk_len / 2 +     // qs packed 4-bit magnitudes: n×k_blks×16
3771                                      n * k_blks * sizeof(uint8_t);  // scale: n×k_blks×1
3772            float * CPtr = c_ptr + n;
3773            size_t  cnt  = k_blks;
3774
3775            // A format (q8 block with per-block scale and stored sum field):
3776            //   || scl(fp32,4B) | asum(int16,2B) | data(int8,32B) || × k_blks
3777            //
3778            // Register map:
3779            //   t3 = k_blks loop counter   t4 = nblks (tail)
3780            //   f0 = A scale (fp32)
3781            //   s2 = pA (scale/asum)       s3 = pA data
3782            //   s4 = pB scales (u8×32)
3783            //   s5 = pB qh (sign/high-bit mask, 128B)
3784            //   s6 = pB qs (packed 4-bit magnitudes, 512B)
3785            //   s7 = pC
3786            //   v3  = fp32 accumulator (N32)
3787            //   v2  = B scales u8 (loaded as bytes; later widened)
3788            //   v0  = qh mask bytes (also used as v0.t mask after load)
3789            //   v1  = A int8 (K32)
3790            //   v8..v15 / v16..v23 = qs unpack/pack temporaries (build signed vmadot lanes)
3791            //   v24/v26/v28/v30    = int32 dot accumulators & packing temps
3792
3793            __asm__ volatile(
3794                "mv           t3, %[BCK]              \n\t"  // t3 = k_blks
3795                "mv           t4, %[NBLKS]            \n\t"  // t4 = nblks (tail guard)
3796
3797                // ---- pre-loop: init fp16 constants in e16 m1 context ----
3798                "vsetvli      t0, x0, e16, m1         \n\t"
3799                "vmv.v.i      v0, 1                   \n\t"  // v0 = int16(1)
3800                "vfcvt.f.x.v  v0, v0                  \n\t"  // v0  = 1.0_fp16
3801                "vxor.vv      v3, v16, v16            \n\t"
3802
3803                // ---- pointer setup ----
3804                "mv           s2, %[pA]               \n\t"  // s2 = pA (scale, fp32)
3805                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pA data (skip scale+asum)
3806                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
3807                "addi         s5, %[pB], 32           \n\t"  // s5 = pBh  (pB + 32B scale)
3808                "addi         s6, %[pB], 32+128       \n\t"  // s6 = pBs  (pB + 32 + 128 = pB+192)
3809                "mv           s7, %[pC]               \n\t"  // s7 = pC
3810
3811                // =====================================================================
3812                // K-block loop: each iteration processes one N32×K32 block
3813                // Stride per k-block = 672B = 32(scl) + 512(Bs) + 128(Bh)
3814                // =====================================================================
3815                ".align 4                             \n\t"
3816                "BLK_LPST%=:                          \n\t"
3817
3818                // ---- load qs (512B = 4 VRF) from s6, advance s6 by 672 ----
3819                "vsetvli      t0, x0, e8, m1          \n\t"
3820                "vl4r.v       v8, (s6)                \n\t"  // v8..v11 = qs N32K32 packed 4-bit magnitudes
3821                "addi         s6, s6, 128*4+128+32    \n\t"  // s6 += 672 (512+128+32)
3822
3823                // ---- load B scale (32B = 32×u8) from s4, advance s4 by 672 ----
3824                "vsetvli      t0, x0, e8, mf2         \n\t"
3825                "vle8.v       v2, (s4)                \n\t"  // v2 = scale_u8 × 32
3826                "addi         s4, s4, 32+128*4+128    \n\t"  // s4 += 672 (32+512+128)
3827
3828                // ---- load qh (128B = 1 VRF) from s5, advance s5 by 672 ----
3829                "vsetvli      t0, x0, e8, m1          \n\t"
3830                "vle8.v       v0, (s5)                \n\t"  // v0 = qh N32K32 sign/high-bit packed
3831                "addi         s5, s5, 128+32+128*4    \n\t"  // s5 += 672 (128+32+512)
3832
3833                // ---- load A data (32B = K32 int8) from s3 ----
3834                "vsetvli      t0, x0, e8, mf4         \n\t"
3835                "vle8.v       v1, (s3)                \n\t"  // v1 = A M1K32 int8
3836                "addi         s3, s3, 32+6            \n\t"  // s3 += 38 (data + scl + asum)
3837
3838                // ---- load A scale (fp32) and asum (int16) from s2 ----
3839                "flw          f0, (s2)                \n\t"  // f0 = A scale (fp32)
3840                "addi         s2, s2, 6+32            \n\t"  // s2 += 38
3841
3842                // ---- Decode packed MXFP4 payload into a vmadot-friendly signed-lane layout ----
3843                "vsetvli      t0, x0, e8, m1          \n\t"
3844                "vand.vi      v12, v8, 0xF            \n\t"  //8bit(lo4) //[8*32]
3845                "vand.vi      v13, v9, 0xF            \n\t"
3846                "vand.vi      v14, v10, 0xF           \n\t"
3847                "vand.vi      v15, v11, 0xF           \n\t"
3848                "vsrl.vi      v8, v8, 4               \n\t"  //8bit(hi4)
3849                "vsrl.vi      v9, v9, 4               \n\t"
3850                "vsrl.vi      v10, v10, 4             \n\t"
3851                "vsrl.vi      v11, v11, 4             \n\t"
3852
3853                // [4*32]*2
3854                "vsetvli      t0, x0, e8, m1          \n\t"
3855                "vpack.vv     v16, v12, v8, 0         \n\t"
3856                "vpack.vv     v18, v13, v9, 0         \n\t"
3857                "vpack.vv     v20, v14, v10, 0        \n\t"
3858                "vpack.vv     v22, v15, v11, 0        \n\t"
3859
3860                "vsetvli      t0, x0, e8, m8          \n\t"
3861                "vrsub.vi     v16, v16, 0, v0.t       \n\t"
3862
3863                // [4*32]*2 -> [8*16]
3864                "vsetvli      t0, x0, e8, m1          \n\t"
3865                "vupack.vv    v8, v16, v17, 1         \n\t"
3866                "vupack.vv    v10, v18, v19, 1        \n\t"
3867                "vupack.vv    v12, v20, v21, 1        \n\t"
3868                "vupack.vv    v14, v22, v23, 1        \n\t"
3869
3870                "vsetvli      t0, x0, e64, m1         \n\t"
3871                "vslidedown.vi  v16, v1, 2            \n\t"
3872
3873                // init the accumu to 0
3874                "vsetvli      t0, x0, e32, m1         \n\t"
3875                "vxor.vv      v24, v16, v16           \n\t"
3876                "vxor.vv      v26, v16, v16           \n\t"
3877                "vxor.vv      v28, v16, v16           \n\t"
3878                "vxor.vv      v30, v16, v16           \n\t"
3879
3880                // ---- int8 dot products over the decoded MXFP4 lane groups ----
3881                "vmadot       v24, v1, v8, i8         \n\t"  // N0..7
3882                "vmadot       v26, v1, v10, i8        \n\t"  // N8..15
3883                "vmadot       v28, v1, v12, i8        \n\t"  // N16..23
3884                "vmadot       v30, v1, v14, i8        \n\t"  // N24..31
3885                "vmadot       v24, v16, v9, i8        \n\t"  // N0..7
3886                "vmadot       v26, v16, v11, i8       \n\t"  // N8..15
3887                "vmadot       v28, v16, v13, i8       \n\t"  // N16..23
3888                "vmadot       v30, v16, v15, i8       \n\t"  // N24..31
3889
3890                "vsetvli      t0, x0, e32, m1         \n\t"
3891                "vpack.vv     v16, v24, v26, 2        \n\t"  // v16 = N0..15
3892                "vpack.vv     v18, v28, v30, 2        \n\t"  // v18 = N16..31
3893                "vpack.vv     v24, v16, v18, 3        \n\t"  // v24 = N0..31
3894
3895                "lui          t1, 0x00200             \n\t"
3896                "vmv.v.x      v30, t1                 \n\t"
3897                // b_scale e8m0 -> fp32
3898                "vsetvli      t0, x0, e8, mf4         \n\t"
3899                "vwaddu.vx    v28, v2, x0             \n\t"
3900                "vsetvli      t0, x0, e16, mf2        \n\t"
3901                "vwadd.vx     v2, v28, x0             \n\t"
3902                "vsetvli      t0, x0, e32, m1         \n\t"
3903                "vmsle.vi     v0, v2, 1               \n\t"
3904                "vadd.vi      v28, v2, -1             \n\t"
3905                "vsll.vi      v28, v28, 23            \n\t"
3906                "vsll.vv      v28, v30, v2, v0.t      \n\t"
3907
3908                // a_scale * b_scale;
3909                "vsetvli      t0, x0, e32, m1         \n\t"
3910                "vfcvt.f.x.v  v26, v24                \n\t"
3911                "vfmul.vf     v30, v28, f0            \n\t"
3912                "vsetvli      t0, x0, e32, m1         \n\t"
3913                // static_cast<float>(qsum) * a_scale * b_scale;
3914                "vfmacc.vv    v3, v30, v26            \n\t"
3915
3916                "addi         t3, t3, -1              \n\t"
3917                "bgtz         t3, BLK_LPST%=          \n\t"
3918                "BLK_LPND%=:                          \n\t"
3919                "vsetvli      t0, %[NBLKS], e32, m1   \n\t"
3920                "vse32.v      v3, (%[pC])             \n\t"
3921                "FUNC_END%=:                          \n\t"
3922
3923                :
3924                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
3925                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6", "s7", "v0", "v1",
3926                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v16", "v17", "v18", "v19",
3927                  "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31");
3928        }
3929    }
3930}
3931
3932void gemm_kernel_i8mxfp4_m4(size_t          blk_len,
3933                            const uint8_t * quant_a_ptr,
3934                            const uint8_t * quant_b_data,
3935                            const uint8_t * quant_b_zp,
3936                            float *         c_ptr,
3937                            size_t          count_m,
3938                            size_t          count_n,
3939                            size_t          k_blks,
3940                            size_t          ldc) {
3941    constexpr size_t NB_COLS = 32;
3942    constexpr size_t K_TILE  = 32;
3943    using blk_type           = nrow_block_mxfp4<NB_COLS>;
3944
3945    GGML_ASSERT(blk_len == K_TILE);
3946    GGML_ASSERT(count_m == 4);
3947    GGML_UNUSED(quant_b_zp);
3948
3949    const size_t a_blk_stride  = q8_blk_size(blk_len, true);
3950    const size_t b_blk_stride  = sizeof(blk_type);
3951    const size_t b_tile_stride = k_blks * b_blk_stride;
3952
3953    if (quant_b_zp == NULL) {
3954        // MXFP4 block layout per K32/N32 tile:
3955        //   [scale_e8m0 x 32][qh sign/high-bit mask x 128B][qs packed 4-bit magnitudes x 512B]
3956        // There is no explicit zp stream; qh is combined with qs to reconstruct signed MXFP4 values.
3957        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
3958            size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
3959            uint8_t * b_data  = (uint8_t *) quant_b_data + (ni / NB_COLS) * b_tile_stride;
3960            uint8_t * a_data  = (uint8_t *) quant_a_ptr;
3961            float *   dst_c   = c_ptr + ni;
3962            size_t    cnt     = k_blks;
3963
3964            asm volatile(
3965                // v4-v7 are the fp32 accumulators for rows 0..3 of the current N32 tile.
3966                "vsetvli        t0, x0, e32, m1         \n\t"
3967                "vxor.vv        v4, v4, v4              \n\t"
3968                "vxor.vv        v5, v5, v5              \n\t"
3969                "vxor.vv        v6, v6, v6              \n\t"
3970                "vxor.vv        v7, v7, v7              \n\t"
3971
3972                ".align 4                               \n\t"
3973                "BLK_LOOP%=:                            \n\t"
3974                // Load the 4 A-row scales for this K32 block and build row data pointers.
3975                "flw            fa0, 0(%[A])            \n\t"
3976                "flw            fa1, 4(%[A])            \n\t"
3977                "flw            fa2, 8(%[A])            \n\t"
3978                "flw            fa3, 12(%[A])           \n\t"
3979                "addi           t3, %[A], 24            \n\t"
3980                "addi           t4, t3, 32              \n\t"
3981                "addi           t5, t3, 64              \n\t"
3982                "addi           t6, t3, 96              \n\t"
3983                "addi           %[A], %[A], 152         \n\t"
3984
3985                // B-side pointers:
3986                //   t1 -> qh bitmask stream, t2 -> qs low-nibble stream.
3987                "addi           t1, %[B], 32            \n\t"
3988                "addi           t2, %[B], 160           \n\t"
3989                "vsetvli        t0, x0, e8, mf2         \n\t"
3990                "vle8.v         v2, (%[B])              \n\t"
3991                "addi           %[B], %[B], 672         \n\t"
3992                "vsetvli        t0, x0, e8, m1          \n\t"
3993                "vle8.v         v0, (t1)                \n\t"
3994                "vl4r.v         v8, (t2)                \n\t"
3995
3996                // Decode the packed MXFP4 payload once for the whole tile and expand it
3997                // into a vmadot-friendly layout.
3998                "vand.vi        v12, v8, 0xF            \n\t"
3999                "vand.vi        v13, v9, 0xF            \n\t"
4000                "vand.vi        v14, v10, 0xF           \n\t"
4001                "vand.vi        v15, v11, 0xF           \n\t"
4002                "vsrl.vi        v8, v8, 4               \n\t"
4003                "vsrl.vi        v9, v9, 4               \n\t"
4004                "vsrl.vi        v10, v10, 4             \n\t"
4005                "vsrl.vi        v11, v11, 4             \n\t"
4006
4007                "vpack.vv       v16, v12, v8, 0         \n\t"
4008                "vpack.vv       v18, v13, v9, 0         \n\t"
4009                "vpack.vv       v20, v14, v10, 0        \n\t"
4010                "vpack.vv       v22, v15, v11, 0        \n\t"
4011
4012                "vsetvli        t0, x0, e8, m8          \n\t"
4013                "vrsub.vi       v16, v16, 0, v0.t       \n\t"
4014
4015                "vsetvli        t0, x0, e8, m1          \n\t"
4016                "vupack.vv      v8, v16, v17, 1         \n\t"
4017                "vupack.vv      v10, v18, v19, 1        \n\t"
4018                "vupack.vv      v12, v20, v21, 1        \n\t"
4019                "vupack.vv      v14, v22, v23, 1        \n\t"
4020
4021                "lui            t1, 0x00200             \n\t"
4022                "vmv.v.x        v30, t1                 \n\t"
4023                // b_scale e8m0 -> fp32
4024                "vsetvli        t0, x0, e8, mf4         \n\t"
4025                "vwaddu.vx      v28, v2, x0             \n\t"
4026                "vsetvli        t0, x0, e16, mf2        \n\t"
4027                "vwadd.vx       v26, v28, x0            \n\t"
4028                "vsetvli        t0, x0, e32, m1         \n\t"
4029                "vmsle.vi       v0, v26, 1              \n\t"
4030                "vadd.vi        v24, v26, -1            \n\t"
4031                "vsll.vi        v18, v24, 23            \n\t"
4032                "vsll.vv        v18, v30, v26, v0.t     \n\t"
4033
4034                // Row 0: dot(A0, decoded MXFP4 lane groups), accumulate in int32 and
4035                // then apply A/B scaling.
4036                "vsetvli        t0, x0, e8, m1          \n\t"
4037                "vle8.v         v1, (t3)                \n\t"
4038                "vsetvli        t0, x0, e64, m1         \n\t"
4039                "vupack.vv      v16, v1, v2, 1          \n\t"
4040                "vsetvli        t0, x0, e32, m1         \n\t"
4041                "vxor.vv        v24, v24, v24           \n\t"
4042                "vxor.vv        v26, v26, v26           \n\t"
4043                "vxor.vv        v28, v28, v28           \n\t"
4044                "vxor.vv        v30, v30, v30           \n\t"
4045                "vmadot         v24, v16, v8, i8        \n\t"
4046                "vmadot         v26, v16, v10, i8       \n\t"
4047                "vmadot         v28, v16, v12, i8       \n\t"
4048                "vmadot         v30, v16, v14, i8       \n\t"
4049                "vmadot         v24, v17, v9, i8        \n\t"
4050                "vmadot         v26, v17, v11, i8       \n\t"
4051                "vmadot         v28, v17, v13, i8       \n\t"
4052                "vmadot         v30, v17, v15, i8       \n\t"
4053                "vpack.vv       v16, v24, v26, 2        \n\t"
4054                "vpack.vv       v20, v28, v30, 2        \n\t"
4055                "vpack.vv       v24, v16, v20, 3        \n\t"
4056                "vpack.vv       v26, v17, v21, 3        \n\t"
4057                "vfcvt.f.x.v    v24, v24                \n\t"
4058                "vfcvt.f.x.v    v25, v25                \n\t"
4059                "vfcvt.f.x.v    v26, v26                \n\t"
4060                "vfcvt.f.x.v    v27, v27                \n\t"
4061                "vfmul.vv       v24, v24, v18           \n\t"
4062                "vfmul.vv       v25, v25, v18           \n\t"
4063                "vfmul.vv       v26, v26, v18           \n\t"
4064                "vfmul.vv       v27, v27, v18           \n\t"
4065                "vfmacc.vf      v4, fa0, v24            \n\t"
4066                "vfmacc.vf      v5, fa1, v25            \n\t"
4067                "vfmacc.vf      v6, fa2, v26            \n\t"
4068                "vfmacc.vf      v7, fa3, v27            \n\t"
4069
4070                "addi           %[BK], %[BK], -1        \n\t"
4071                "bgtz           %[BK], BLK_LOOP%=       \n\t"
4072
4073                // Tail-aware store for the final N tile (`nb_real` may be < 32).
4074                "vsetvli        t0, %[NBLKS], e32, m1   \n\t"
4075                "add            t1, %[LDC], %[DST]      \n\t"
4076                "vse32.v        v4, (%[DST])            \n\t"
4077                "vse32.v        v5, (t1)                \n\t"
4078                "add            t2, t1, %[LDC]          \n\t"
4079                "vse32.v        v6, (t2)                \n\t"
4080                "add            t3, t2, %[LDC]          \n\t"
4081                "vse32.v        v7, (t3)                \n\t"
4082                : [A] "+r"(a_data), [B] "+r"(b_data), [BK] "+r"(cnt)
4083                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [NBLKS] "r"(nb_real)
4084                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "t5", "t6", "s1", "s2", "s3", "s4", "v0", "v1", "v2",
4085                  "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
4086                  "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31",
4087                  "fa0", "fa1", "fa2", "fa3");
4088        }
4089    }
4090}
4091
4092void gemm_kernel_i8i5_m1(size_t          blk_len,
4093                         const uint8_t * quant_a_ptr,
4094                         const uint8_t * quant_b_data,
4095                         const uint8_t * quant_b_zp,
4096                         float *         c_ptr,
4097                         size_t          count_m,
4098                         size_t          count_n,
4099                         size_t          k_blks,
4100                         size_t          ldc) {
4101    // =========================================================================
4102    // i8i5: 8-bit activation × 5-bit weight (4-bit low + 1-bit high mask)
4103    //
4104    // B layout per N32K32 k-block (no-zp):
4105    //   [0  .. 63 ] : scale_fp16 × 32              (64B)
4106    //   [64 .. 191] : Bh i1-high-bit  × 32N × 32K  (128B = 1 VRF)
4107    //   [192.. 703] : Bs i4-low-nibble × 32N × 32K (512B = 4 VRF)
4108    //   Total: 704B per k-block stride
4109    //
4110    // B layout per N32K32 k-block (with-zp):
4111    //   [0  .. 63 ] : scale_fp16 × 32              (64B)
4112    //   [64 .. 95 ] : zp_uint8 × 32                (32B)
4113    //   [96 .. 223] : Bh i1-high-bit  × 32N × 32K  (128B = 1 VRF)
4114    //   [224.. 735] : Bs i4-low-nibble × 32N × 32K (512B = 4 VRF)
4115    //   Total: 736B per k-block stride
4116    //
4117    // Bh format per N8K32 sub-block (32B):
4118    //   K rows × N cols × 1bit packed as bytes (8 cols per byte, K groups of 4B)
4119    //   Byte k gives 8 mask bits for columns N7..N0 at k-th K-element.
4120    //
4121    // Computation:
4122    //   B5bit_signed = (Bs | (Bh << 4)) - zp
4123    //   dot(A, B5) = dot(A, Bs_u4) + 16*dot(A, Bh_u1) - zp*asum
4124    //   No-zp: implicit zp = 16 (unsigned [0..31] centered at 16)
4125    //   With-zp: explicit zp from data
4126    //
4127    // =========================================================================
4128
4129    if (quant_b_zp == NULL) {
4130        for (size_t n = 0; n < count_n; n += 32) {
4131            size_t    nblks         = (count_n - n) > 32 ? 32 : count_n - n;
4132            // i8i5 no-zp: per column per k-block stride = fp16(2B) + i4(16B) + i1(4B) = 22B
4133            uint8_t * QuantBDataPtr = (uint8_t *) quant_b_data +      //
4134                                      n * k_blks * (blk_len / 8) +    // Bh i1 mask: n×k_blks×4
4135                                      n * k_blks * blk_len / 2 +      // Bs i4 data: n×k_blks×16
4136                                      n * k_blks * sizeof(_Float16);  // scale: n×k_blks×2
4137            float * CPtr = c_ptr + n;
4138            size_t  cnt  = k_blks;
4139
4140            // A format (same as i8i4):
4141            //   || scl(fp32,4B) | asum(int16,2B) | data(int8,32B) || × k_blks
4142            //
4143            // Register map:
4144            //   t3 = k_blks loop counter   t4 = nblks (tail)
4145            //   t2 = A asum (int16) << 4   f0 = A scale (fp32)
4146            //   s2 = pA (scale/asum)       s3 = pA data
4147            //   s4 = pB scales (fp16×32)
4148            //   s5 = pB Bh (i1 mask, 128B)
4149            //   s6 = pB Bs (i4 packed, 512B)
4150            //   s7 = pC
4151            //   v3  = fp32 accumulator (N32)
4152            //   v2  = B scales fp16 (loaded as bytes; later widened)
4153            //   v0  = Bh mask bytes (also used as v0.t mask after load)
4154            //   v1  = A int8 (K32)
4155            //   v8..v15 / v16..v23 = Bs unpack/pack temporaries (build b5bit bytes)
4156            //   v24/v26/v28/v30    = int32 dot accumulators & packing temps
4157
4158            __asm__ volatile(
4159                "mv           t3, %[BCK]              \n\t"  // t3 = k_blks
4160                "mv           t4, %[NBLKS]            \n\t"  // t4 = nblks (tail guard)
4161
4162                // ---- pre-loop: init fp16 constants in e16 m1 context ----
4163                "vsetvli      t0, x0, e16, m1         \n\t"
4164                "vmv.v.i      v0, 1                   \n\t"  // v0 = int16(1)
4165                "vfcvt.f.x.v  v0, v0                  \n\t"  // v0  = 1.0_fp16
4166                "vxor.vv      v3, v16, v16            \n\t"
4167
4168                // ---- pointer setup ----
4169                "mv           s2, %[pA]               \n\t"  // s2 = pA (scale, fp32)
4170                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pA data (skip scale+asum)
4171                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
4172                "addi         s5, %[pB], 32*2         \n\t"  // s5 = pBh  (pB + 64B scale)
4173                "addi         s6, %[pB], 32*2+128     \n\t"  // s6 = pBs  (pB + 64 + 128 = pB+192)
4174                "mv           s7, %[pC]               \n\t"  // s7 = pC
4175
4176                // =====================================================================
4177                // K-block loop: each iteration processes one N32×K32 block
4178                // Stride per k-block = 704B = 64(scl) + 512(Bs) + 128(Bh)
4179                // =====================================================================
4180                ".align 4                             \n\t"
4181                "BLK_LPST%=:                          \n\t"
4182
4183                // ---- load Bs (512B = 4 VRF) from s6, advance s6 by 704 ----
4184                "vsetvli      t0, x0, e8, m1          \n\t"
4185                "vl4r.v       v8, (s6)                \n\t"  // v8..v11 = Bs N32K32 i4
4186                "addi         s6, s6, 128*4+128+64    \n\t"  // s6 += 704 (512+128+64)
4187
4188                // ---- load B scale (64B = 32×fp16) from s4, advance s4 by 704 ----
4189                "vsetvli      t0, x0, e8, mf2         \n\t"
4190                "vle8.v       v2, (s4)                \n\t"  // v2 = scale_fp16 × 32
4191                "addi         s4, s4, 64+128*4+128    \n\t"  // s4 += 704 (64+512+128)
4192
4193                // ---- load Bh (128B = 1 VRF) from s5, advance s5 by 704 ----
4194                "vsetvli      t0, x0, e8, m1          \n\t"
4195                "vle8.v       v0, (s5)                \n\t"  // v0 = Bh N32K32 1-bit packed
4196                "addi         s5, s5, 128+64+128*4    \n\t"  // s5 += 704 (128+64+512)
4197
4198                // ---- load A data (32B = K32 int8) from s3 ----
4199                "vsetvli      t0, x0, e8, mf4         \n\t"
4200                "vle8.v       v1, (s3)                \n\t"  // v1 = A M1K32 int8
4201                "addi         s3, s3, 32+6            \n\t"  // s3 += 38 (data + scl + asum)
4202
4203                // ---- load A scale (fp32) and asum (int16) from s2 ----
4204                "flw          f0, (s2)                \n\t"  // f0 = A scale (fp32)
4205                "lh           t2, 4(s2)               \n\t"  // t2 = A asum (int16)
4206                "addi         s2, s2, 6+32            \n\t"  // s2 += 38
4207
4208                //// ---- A nibble unpacking ----
4209                "vsetvli      t0, x0, e8, m1          \n\t"
4210                "vand.vi      v12, v8, 0xF            \n\t"  //8bit(lo4) //[8*32]
4211                "vand.vi      v13, v9, 0xF            \n\t"
4212                "vand.vi      v14, v10, 0xF           \n\t"
4213                "vand.vi      v15, v11, 0xF           \n\t"
4214                "vsrl.vi      v8, v8, 4               \n\t"  //8bit(hi4)
4215                "vsrl.vi      v9, v9, 4               \n\t"
4216                "vsrl.vi      v10, v10, 4             \n\t"
4217                "vsrl.vi      v11, v11, 4             \n\t"
4218
4219                "slli         t2, t2, 4               \n\t"  // a_sum * 16;
4220                // [4*32]*2
4221                "vsetvli      t0, x0, e8, m1          \n\t"
4222                "vpack.vv     v16, v12, v8, 0         \n\t"
4223                "vpack.vv     v18, v13, v9, 0         \n\t"
4224                "vpack.vv     v20, v14, v10, 0        \n\t"
4225                "vpack.vv     v22, v15, v11, 0        \n\t"
4226
4227                "li           t1, 16                  \n\t"
4228                "vsetvli      t0, x0, e8, m8          \n\t"
4229                "vadd.vx      v16, v16, t1, v0.t      \n\t"
4230
4231                // [4*32]*2 -> [8*16]
4232                "vsetvli      t0, x0, e8, m1          \n\t"
4233                "vupack.vv    v8, v16, v17, 1         \n\t"
4234                "vupack.vv    v10, v18, v19, 1        \n\t"
4235                "vupack.vv    v12, v20, v21, 1        \n\t"
4236                "vupack.vv    v14, v22, v23, 1        \n\t"
4237
4238                "vsetvli      t0, x0, e64, m1         \n\t"
4239                "vslidedown.vi  v16, v1, 2            \n\t"
4240
4241                // init the accumu to asum * zp
4242                "vsetvli        t0, x0, e32, m1 \n\t"
4243                "vxor.vv        v24, v16, v16   \n\t"
4244                "vxor.vv        v26, v16, v16   \n\t"
4245                "vxor.vv        v28, v16, v16   \n\t"
4246                "vxor.vv        v30, v16, v16   \n\t"
4247
4248                // ---- i8 main dot products ----
4249                // vmadot: A × unsigned Bh × 16 → fp16 accumulate
4250                "vmadot       v24, v1, v8, i8         \n\t"  // N0..7
4251                "vmadot       v26, v1, v10, i8        \n\t"  // N8..15
4252                "vmadot       v28, v1, v12, i8        \n\t"  // N16..23
4253                "vmadot       v30, v1, v14, i8        \n\t"  // N24..31
4254                //// vmadot: A × unsigned Bh × 1 → fp16 accumulate
4255                "vmadot       v24, v16, v9, i8        \n\t"  // N0..7
4256                "vmadot       v26, v16, v11, i8       \n\t"  // N8..15
4257                "vmadot       v28, v16, v13, i8       \n\t"  // N16..23
4258                "vmadot       v30, v16, v15, i8       \n\t"  // N24..31
4259
4260                "vsetvli      t0, x0, e32, m1         \n\t"
4261                "vpack.vv     v16, v24, v26, 2        \n\t"  // v16 = N0..15
4262                "vpack.vv     v18, v28, v30, 2        \n\t"  // v18 = N16..31
4263                "vpack.vv     v24, v16, v18, 3        \n\t"  // v24 = N0..31
4264
4265                "vadd.vx      v24, v24, t2            \n\t"
4266                // b_scale fp16 -> fp32
4267                "vsetvli      t0, x0, e16, mf2        \n\t"
4268                "vfwcvt.f.f.v v28, v2                 \n\t"
4269
4270                // a_scale * b_scale;
4271                "vsetvli      t0, x0, e32, m1         \n\t"
4272                "vfcvt.f.x.v  v26, v24                \n\t"
4273                "vfmul.vf     v30, v28, f0            \n\t"
4274                "vsetvli      t0, x0, e32, m1         \n\t"
4275                // static_cast<float>(qsum) * a_scale * b_scale;
4276                "vfmacc.vv    v3, v30, v26            \n\t"
4277
4278                "addi         t3, t3, -1              \n\t"
4279                "bgtz         t3, BLK_LPST%=          \n\t"
4280                "BLK_LPND%=:                          \n\t"
4281                "vsetvli      t0, %[NBLKS], e32, m1   \n\t"
4282                "vse32.v      v3, (%[pC])             \n\t"
4283                "FUNC_END%=:                          \n\t"
4284
4285                :
4286                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
4287                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6", "s7", "v0", "v1",
4288                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v16", "v17", "v18", "v19",
4289                  "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31");
4290        }
4291    } else {
4292        for (size_t n = 0; n < count_n; n += 32) {
4293            size_t    nblks         = (count_n - n) > 32 ? 32 : count_n - n;
4294            // i8i5 with-zp: per column per k-block stride = fp16(2B)+zp(1B)+i4(16B)+i1(4B)=23B
4295            uint8_t * QuantBDataPtr = (uint8_t *) quant_b_data +      //
4296                                      n * k_blks * blk_len / 2 +      // Bs i4: n×k_blks×16
4297                                      n * k_blks * (blk_len / 8) +    // Bh i1: n×k_blks×4
4298                                      n * k_blks * sizeof(uint8_t) +  // zp: n×k_blks×1
4299                                      n * k_blks * sizeof(_Float16);  // scale: n×k_blks×2
4300            float * CPtr = c_ptr + n;
4301            size_t  cnt  = k_blks;
4302
4303            // A format (same as i8i4):
4304            //   || scl(fp32,4B) | asum(int16,2B) | data(int8,32B) || × k_blks
4305            //
4306            // Register map:
4307            //   t3 = k_blks loop counter   t4 = nblks (tail)
4308            //   t2 = A asum (int16) << 4   f0 = A scale (fp32)
4309            //   s2 = pA (scale/asum)       s3 = pA data
4310            //   s4 = pB scales (fp16×32); 每个 k-block 先 +64 指向 zp,再 +672 到下一个 block
4311            //   s5 = pB Bh (i1 mask, 128B) (offset +96)
4312            //   s6 = pB Bs (i4 packed, 512B) (offset +224)
4313            //   s7 = pC
4314            //   v3  = fp32 accumulator (N32)
4315            //   v2  = B scales fp16 (loaded as bytes; later widened)
4316            //   v0  = Bh mask bytes (also used as v0.t mask after load)
4317            //   v1  = A int8 (K32) / later reused to hold Bzp bytes
4318            //   v8..v15 / v16..v23 = Bs unpack/pack temporaries (build b5bit bytes)
4319            //   v24/v26/v28/v30    = int32 dot accumulators & packing temps
4320
4321            __asm__ volatile(
4322                "mv           t3, %[BCK]              \n\t"  // t3 = k_blks
4323                "mv           t4, %[NBLKS]            \n\t"  // t4 = nblks (tail guard)
4324
4325                // ---- pre-loop: init fp16 constants in e16 m1 context ----
4326                "vsetvli      t0, x0, e16, m1         \n\t"
4327                "vmv.v.i      v0, 1                   \n\t"  // v0 = int16(1)
4328                "vfcvt.f.x.v  v0, v0                  \n\t"  // v0  = 1.0_fp16
4329                "vxor.vv      v3, v16, v16            \n\t"
4330
4331                // ---- pointer setup ----
4332                "mv           s2, %[pA]               \n\t"  // s2 = pA (scale, fp32)
4333                "addi         s3, %[pA], 4+2          \n\t"  // s3 = pA data (skip scale+asum)
4334                "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
4335                "addi         s5, %[pB], 32*3         \n\t"  // s5 = pBh  (pB + 64B scale + 32B zp = pB+96)
4336                "addi         s6, %[pB], 32*3+128     \n\t"  // s6 = pBs  (pB + 96 + 128 = pB+224)
4337                "mv           s7, %[pC]               \n\t"  // s7 = pC
4338
4339                // =====================================================================
4340                // K-block loop: each iteration processes one N32×K32 block
4341                // Stride per k-block = 736B = 64(scale) + 32(zp) + 128(Bh) + 512(Bs)
4342                // =====================================================================
4343                ".align 4                             \n\t"
4344                "BLK_LPST%=:                          \n\t"
4345
4346                // ---- load Bs (512B = 4 VRF) from s6, advance s6 by 736 ----
4347                "vsetvli      t0, x0, e8, m1          \n\t"
4348                "vl4r.v       v8, (s6)                \n\t"  // v8..v11 = Bs N32K32 i4
4349                "addi         s6, s6, 128*4+128+96    \n\t"  // s6 += 736 (512+128+96)
4350
4351                // ---- load B scale (64B = 32×fp16) from s4; then s4 points to zp[32] ----
4352                "vsetvli      t0, x0, e8, mf2         \n\t"
4353                "vle8.v       v2, (s4)                \n\t"  // v2 = scale_fp16 × 32
4354                "addi         s4, s4, 64              \n\t"  // s4 += 64 (now points to zp)
4355
4356                // ---- load Bh (128B = 1 VRF) from s5, advance s5 by 736 ----
4357                "vsetvli      t0, x0, e8, m1          \n\t"
4358                "vle8.v       v0, (s5)                \n\t"  // v0 = Bh N32K32 1-bit packed
4359                "addi         s5, s5, 128+96+128*4    \n\t"  // s5 += 736 (128+96+512)
4360
4361                // ---- load A data (32B = K32 int8) from s3 ----
4362                "vsetvli      t0, x0, e8, mf4         \n\t"
4363                "vle8.v       v1, (s3)                \n\t"  // v1 = A M1K32 int8
4364                "addi         s3, s3, 32+6            \n\t"  // s3 += 38 (data + scl + asum)
4365
4366                // ---- load A scale (fp32) and asum (int16) from s2 ----
4367                "flw          f0, (s2)                \n\t"  // f0 = A scale (fp32)
4368                "lh           t2, 4(s2)               \n\t"  // t2 = A asum (int16)
4369                "addi         s2, s2, 6+32            \n\t"  // s2 += 38
4370
4371                //// ---- A nibble unpacking ----
4372                "vsetvli      t0, x0, e8, m1          \n\t"
4373                "vand.vi      v12, v8, 0xF            \n\t"  //8bit(lo4) //[8*32]
4374                "vand.vi      v13, v9, 0xF            \n\t"
4375                "vand.vi      v14, v10, 0xF           \n\t"
4376                "vand.vi      v15, v11, 0xF           \n\t"
4377                "vsrl.vi      v8, v8, 4               \n\t"  //8bit(hi4)
4378                "vsrl.vi      v9, v9, 4               \n\t"
4379                "vsrl.vi      v10, v10, 4             \n\t"
4380                "vsrl.vi      v11, v11, 4             \n\t"
4381
4382                // [4*32]*2
4383                "vsetvli      t0, x0, e8, m1          \n\t"
4384                "vpack.vv     v16, v12, v8, 0         \n\t"
4385                "vpack.vv     v18, v13, v9, 0         \n\t"
4386                "vpack.vv     v20, v14, v10, 0        \n\t"
4387                "vpack.vv     v22, v15, v11, 0        \n\t"
4388
4389                "li           t1, 16                  \n\t"
4390                "vsetvli      t0, x0, e8, m8          \n\t"
4391                "vadd.vx      v16, v16, t1, v0.t      \n\t"
4392
4393                // [4*32]*2 -> [8*16]
4394                "vsetvli      t0, x0, e8, m1          \n\t"
4395                "vupack.vv    v8, v16, v17, 1         \n\t"
4396                "vupack.vv    v10, v18, v19, 1        \n\t"
4397                "vupack.vv    v12, v20, v21, 1        \n\t"
4398                "vupack.vv    v14, v22, v23, 1        \n\t"
4399
4400                "vsetvli      t0, x0, e64, m1         \n\t"
4401                "vslidedown.vi  v16, v1, 2            \n\t"
4402
4403                "vsetvli      t0, x0, e32, m1         \n\t"
4404                "vxor.vv      v24, v16, v16           \n\t"
4405                "vxor.vv      v26, v16, v16           \n\t"
4406                "vxor.vv      v28, v16, v16           \n\t"
4407                "vxor.vv      v30, v16, v16           \n\t"
4408
4409                // ---- i8 main dot products ----
4410                // vmadot: A × unsigned Bh × 16 → fp16 accumulate
4411                "vmadot       v24, v1, v8, i8         \n\t"  // N0..7
4412                "vmadot       v26, v1, v10, i8        \n\t"  // N8..15
4413                "vmadot       v28, v1, v12, i8        \n\t"  // N16..23
4414                "vmadot       v30, v1, v14, i8        \n\t"  // N24..31
4415                // vmadot: A × unsigned Bh × 1 → fp16 accumulate
4416                "vmadot       v24, v16, v9, i8        \n\t"  // N0..7
4417                "vmadot       v26, v16, v11, i8       \n\t"  // N8..15
4418                "vmadot       v28, v16, v13, i8       \n\t"  // N16..23
4419                "vmadot       v30, v16, v15, i8       \n\t"  // N24..31
4420
4421                "vsetvli      t0, x0, e8, m1          \n\t"
4422                "vle8.v       v1, (s4)                \n\t"  // Bzp
4423                "addi         s4, s4, 32+128*4+128    \n\t"
4424
4425                "vsetvli      t0, x0, e8, m1          \n\t"
4426                "vpack.vv     v16, v24, v26, 2        \n\t"  // v16 = N0..15
4427                "vpack.vv     v18, v28, v30, 2        \n\t"  // v18 = N16..31
4428                "vpack.vv     v24, v16, v18, 3        \n\t"  // v24 = N0..31
4429
4430                "vwaddu.vx    v28, v1, x0             \n\t"  // uint8 -> uint16
4431
4432                "vsetvli      t0, x0, e16, m1         \n\t"
4433                "vwmul.vx     v30, v28, t2            \n\t"
4434
4435                // b_scale fp16 -> fp32
4436                "vsetvli      t0, x0, e16, mf2        \n\t"
4437                "vfwcvt.f.f.v v28, v2                 \n\t"
4438                "vsetvli      t0, x0, e32, m1         \n\t"
4439                "vadd.vv      v24, v24, v30           \n\t"
4440
4441                // a_scale * b_scale;
4442                "vsetvli      t0, x0, e32, m1         \n\t"
4443                "vfmul.vf     v30, v28, f0            \n\t"
4444                "vfcvt.f.x.v  v26, v24                \n\t"
4445                "vsetvli      t0, x0, e32, m1         \n\t"
4446                // static_cast<float>(qsum) * a_scale * b_scale;
4447                "vfmacc.vv    v3, v30, v26            \n\t"
4448
4449                "addi         t3, t3, -1              \n\t"
4450                "bgtz         t3, BLK_LPST%=          \n\t"
4451                "BLK_LPND%=:                          \n\t"
4452                "vsetvli      t0, %[NBLKS], e32, m1   \n\t"
4453                "vse32.v      v3, (%[pC])             \n\t"
4454                "FUNC_END%=:                          \n\t"
4455                :
4456                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
4457                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6", "s7", "v0", "v1",
4458                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v16", "v17", "v18", "v19",
4459                  "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31");
4460        }
4461    }
4462}
4463
4464void gemm_kernel_i8i5_m4(size_t          blk_len,
4465                         const uint8_t * quant_a_ptr,
4466                         const uint8_t * quant_b_data,
4467                         const uint8_t * quant_b_zp,
4468                         float *         c_ptr,
4469                         size_t          count_m,
4470                         size_t          count_n,
4471                         size_t          k_blks,
4472                         size_t          ldc) {
4473    constexpr size_t NB_COLS = 32;
4474
4475    GGML_UNUSED(count_m);
4476    GGML_UNUSED(blk_len);
4477
4478    // This kernel computes a 4x32 output tile. For each K32 block we decode the
4479    // packed Q5 weights once and reuse the decoded vectors across the 4 A rows.
4480    constexpr size_t B_Q50_BLK_STRIDE = sizeof(nrow_block_q5_0<NB_COLS>);
4481    constexpr size_t B_Q51_BLK_STRIDE = sizeof(nrow_block_q5_1<NB_COLS>);
4482
4483    if (quant_b_zp) {
4484        // Q5_1 block layout per K32/N32 tile:
4485        //   [scale_fp16 x 32][zp_u8 x 32][qh high-bit mask x 128B][qs low nibbles x 512B]
4486        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
4487            size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
4488            uint8_t * b_data  = (uint8_t *) quant_b_data + (ni / NB_COLS) * k_blks * B_Q51_BLK_STRIDE;
4489            uint8_t * a_data  = (uint8_t *) quant_a_ptr;
4490            float *   dst_c   = c_ptr + ni;
4491            size_t    cnt     = k_blks;
4492
4493            asm volatile(
4494                // v4-v7 are the fp32 accumulators for rows 0..3 of the current N32 tile.
4495                "vsetvli        t0, x0, e32, m1         \n\t"
4496                "vxor.vv        v4, v4, v4              \n\t"
4497                "vxor.vv        v5, v5, v5              \n\t"
4498                "vxor.vv        v6, v6, v6              \n\t"
4499                "vxor.vv        v7, v7, v7              \n\t"
4500
4501                ".align 4                               \n\t"
4502                "BLK_LOOP%=:                            \n\t"
4503                // Load the 4 A-row scales/sums for this K32 block and build row data pointers.
4504                "flw            fa0, 0(%[A])            \n\t"
4505                "flw            fa1, 4(%[A])            \n\t"
4506                "flw            fa2, 8(%[A])            \n\t"
4507                "flw            fa3, 12(%[A])           \n\t"
4508                "lh             s1, 16(%[A])            \n\t"
4509                "lh             s2, 18(%[A])            \n\t"
4510                "lh             s3, 20(%[A])            \n\t"
4511                "lh             s4, 22(%[A])            \n\t"
4512                "addi           t3, %[A], 24            \n\t"
4513                "addi           t4, t3, 32              \n\t"
4514                "addi           t5, t3, 64              \n\t"
4515                "addi           t6, t3, 96              \n\t"
4516                "addi           %[A], %[A], 152         \n\t"
4517
4518                // B-side pointers:
4519                //   t1 -> zp stream, t2 -> qh bitmask stream, s5 -> qs low-nibble stream.
4520                "addi           t1, %[B], 64            \n\t"
4521                "addi           t2, %[B], 96            \n\t"
4522                "addi           s5, %[B], 224           \n\t"
4523                "vsetvli        t0, x0, e8, mf2         \n\t"
4524                "vle8.v         v2, (%[B])              \n\t"
4525                "vsetvli        t0, x0, e8, m1          \n\t"
4526                "vle8.v         v0, (t2)                \n\t"
4527                "vl4r.v         v8, (s5)                \n\t"
4528                "addi           %[B], %[B], 736         \n\t"
4529
4530                // Decode Q5 payload once for the whole tile:
4531                //   1) split `qs` low/high nibbles,
4532                //   2) repack into bytes,
4533                //   3) use the `qh` mask to inject bit4 (+16) where needed,
4534                //   4) expand into the vmadot-friendly layout reused by all 4 rows.
4535                "vand.vi        v12, v8, 0xF            \n\t"
4536                "vand.vi        v13, v9, 0xF            \n\t"
4537                "vand.vi        v14, v10, 0xF           \n\t"
4538                "vand.vi        v15, v11, 0xF           \n\t"
4539                "vsrl.vi        v8, v8, 4               \n\t"
4540                "vsrl.vi        v9, v9, 4               \n\t"
4541                "vsrl.vi        v10, v10, 4             \n\t"
4542                "vsrl.vi        v11, v11, 4             \n\t"
4543
4544                "vpack.vv       v16, v12, v8, 0         \n\t"
4545                "vpack.vv       v18, v13, v9, 0         \n\t"
4546                "li             t2, 16                  \n\t"
4547                "vpack.vv       v20, v14, v10, 0        \n\t"
4548                "vpack.vv       v22, v15, v11, 0        \n\t"
4549
4550                "vsetvli        t0, x0, e8, m8          \n\t"
4551                "vadd.vx        v16, v16, t2, v0.t      \n\t"
4552
4553                "vsetvli        t0, x0, e8, m1          \n\t"
4554                "vupack.vv      v8, v16, v17, 1         \n\t"
4555                "vupack.vv      v10, v18, v19, 1        \n\t"
4556                "vupack.vv      v12, v20, v21, 1        \n\t"
4557                "vupack.vv      v14, v22, v23, 1        \n\t"
4558
4559                // Convert per-column fp16 scales once; the same scale vector is shared by all 4 rows.
4560                "vsetvli        t0, x0, e16, mf2        \n\t"
4561                "vfwcvt.f.f.v   v18, v2                 \n\t"
4562                "vsetvli        t0, x0, e8, m1          \n\t"
4563                "vle8.v         v3, (t1)                \n\t"
4564                "vsetvli        t0, x0, e8, m1          \n\t"
4565
4566                // Row 0: dot(A0, decoded_q5) + a_sum0 * zp, then scale by A/B scales.
4567                // The widen/mul correction sequence intentionally matches the proven m1 Q5_1 path.
4568                "vle8.v         v1, (t3)                \n\t"
4569                "vsetvli        t0, x0, e64, m1         \n\t"
4570                "vupack.vv      v16, v1, v2, 1          \n\t"
4571                "vsetvli        t0, x0, e32, m1         \n\t"
4572                "vxor.vv        v24, v24, v24           \n\t"
4573                "vxor.vv        v26, v26, v26           \n\t"
4574                "vxor.vv        v28, v28, v28           \n\t"
4575                "vxor.vv        v30, v30, v30           \n\t"
4576                "vmadot         v24, v16, v8, i8        \n\t"
4577                "vmadot         v26, v16, v10, i8       \n\t"
4578                "vmadot         v28, v16, v12, i8       \n\t"
4579                "vmadot         v30, v16, v14, i8       \n\t"
4580                "vmadot         v24, v17, v9, i8        \n\t"
4581                "vmadot         v26, v17, v11, i8       \n\t"
4582                "vmadot         v28, v17, v13, i8       \n\t"
4583                "vmadot         v30, v17, v15, i8       \n\t"
4584                "vpack.vv       v16, v24, v26, 2        \n\t"
4585                "vpack.vv       v20, v28, v30, 2        \n\t"
4586                "vpack.vv       v24, v16, v20, 3        \n\t"
4587                "vpack.vv       v26, v17, v21, 3        \n\t"
4588                "vsetvli        t0, x0, e8, m1          \n\t"
4589                "vwaddu.vx      v28, v3, x0             \n\t"
4590                "vsetvli        t0, x0, e16, m1         \n\t"
4591                "vwmul.vx       v12, v28, s1            \n\t"
4592                "vwmul.vx       v14, v28, s2            \n\t"
4593                "vwmul.vx       v20, v28, s3            \n\t"
4594                "vwmul.vx       v22, v28, s4            \n\t"
4595                "vsetvli        t0, x0, e32, m1         \n\t"
4596                "vadd.vv        v24, v24, v12           \n\t"
4597                "vadd.vv        v25, v25, v14           \n\t"
4598                "vadd.vv        v26, v26, v20           \n\t"
4599                "vadd.vv        v27, v27, v22           \n\t"
4600                "vfcvt.f.x.v    v12, v24                \n\t"
4601                "vfcvt.f.x.v    v14, v25                \n\t"
4602                "vfcvt.f.x.v    v20, v26                \n\t"
4603                "vfcvt.f.x.v    v22, v27                \n\t"
4604                "vfmul.vv       v12, v12, v18           \n\t"
4605                "vfmul.vv       v14, v14, v18           \n\t"
4606                "vfmul.vv       v20, v20, v18           \n\t"
4607                "vfmul.vv       v22, v22, v18           \n\t"
4608                "vfmacc.vf      v4, fa0, v12            \n\t"
4609                "vfmacc.vf      v5, fa1, v14            \n\t"
4610                "vfmacc.vf      v6, fa2, v20            \n\t"
4611                "vfmacc.vf      v7, fa3, v22            \n\t"
4612
4613                "addi           %[BK], %[BK], -1        \n\t"
4614                "bgtz           %[BK], BLK_LOOP%=       \n\t"
4615
4616                // Tail-aware store for the final N tile (`nb_real` may be < 32).
4617                "vsetvli        t0, %[NBLKS], e32, m1   \n\t"
4618                "add            t1, %[LDC], %[DST]      \n\t"
4619                "vse32.v        v4, (%[DST])            \n\t"
4620                "vse32.v        v5, (t1)                \n\t"
4621                "add            t2, t1, %[LDC]          \n\t"
4622                "vse32.v        v6, (t2)                \n\t"
4623                "add            t3, t2, %[LDC]          \n\t"
4624                "vse32.v        v7, (t3)                \n\t"
4625                : [A] "+r"(a_data), [B] "+r"(b_data), [BK] "+r"(cnt)
4626                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [NBLKS] "r"(nb_real)
4627                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "t5", "t6", "s1", "s2", "s3", "s4", "s5", "v0", "v1",
4628                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16",
4629                  "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30",
4630                  "v31", "fa0", "fa1", "fa2", "fa3");
4631        }
4632    } else {
4633        // Q5_0 block layout per K32/N32 tile:
4634        //   [scale_fp16 x 32][qh high-bit mask x 128B][qs low nibbles x 512B]
4635        // There is no explicit zp stream; the implicit midpoint correction is +16.
4636        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
4637            size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
4638            uint8_t * b_data  = (uint8_t *) quant_b_data + (ni / NB_COLS) * k_blks * B_Q50_BLK_STRIDE;
4639            uint8_t * a_data  = (uint8_t *) quant_a_ptr;
4640            float *   dst_c   = c_ptr + ni;
4641            size_t    cnt     = k_blks;
4642
4643            asm volatile(
4644                // v4-v7 are the fp32 accumulators for rows 0..3 of the current N32 tile.
4645                "vsetvli        t0, x0, e32, m1         \n\t"
4646                "vxor.vv        v4, v4, v4              \n\t"
4647                "vxor.vv        v5, v5, v5              \n\t"
4648                "vxor.vv        v6, v6, v6              \n\t"
4649                "vxor.vv        v7, v7, v7              \n\t"
4650
4651                ".align 4                               \n\t"
4652                "BLK_LOOP%=:                            \n\t"
4653                // Load the 4 A-row scales/sums for this K32 block and build row data pointers.
4654                "flw            fa0, 0(%[A])            \n\t"
4655                "flw            fa1, 4(%[A])            \n\t"
4656                "flw            fa2, 8(%[A])            \n\t"
4657                "flw            fa3, 12(%[A])           \n\t"
4658                "lh             s1, 16(%[A])            \n\t"
4659                "lh             s2, 18(%[A])            \n\t"
4660                "lh             s3, 20(%[A])            \n\t"
4661                "lh             s4, 22(%[A])            \n\t"
4662                "addi           t3, %[A], 24            \n\t"
4663                "addi           t4, t3, 32              \n\t"
4664                "addi           t5, t3, 64              \n\t"
4665                "addi           t6, t3, 96              \n\t"
4666                "addi           %[A], %[A], 152         \n\t"
4667
4668                // B-side pointers:
4669                //   t1 -> qh bitmask stream, t2 -> qs low-nibble stream.
4670                "addi           t1, %[B], 64            \n\t"
4671                "addi           t2, %[B], 192           \n\t"
4672                "vsetvli        t0, x0, e8, mf2         \n\t"
4673                "vle8.v         v2, (%[B])              \n\t"
4674                "vsetvli        t0, x0, e8, m1          \n\t"
4675                "vle8.v         v0, (t1)                \n\t"
4676                "vl4r.v         v8, (t2)                \n\t"
4677                "addi           %[B], %[B], 704         \n\t"
4678
4679                // Decode Q5 payload once for the whole tile and expand it into the vmadot layout.
4680                "vand.vi        v12, v8, 0xF            \n\t"
4681                "vand.vi        v13, v9, 0xF            \n\t"
4682                "vand.vi        v14, v10, 0xF           \n\t"
4683                "vand.vi        v15, v11, 0xF           \n\t"
4684                "vsrl.vi        v8, v8, 4               \n\t"
4685                "vsrl.vi        v9, v9, 4               \n\t"
4686                "vsrl.vi        v10, v10, 4             \n\t"
4687                "vsrl.vi        v11, v11, 4             \n\t"
4688
4689                "vpack.vv       v16, v12, v8, 0         \n\t"
4690                "vpack.vv       v18, v13, v9, 0         \n\t"
4691                "li             t2, 16                  \n\t"
4692                "vpack.vv       v20, v14, v10, 0        \n\t"
4693                "vpack.vv       v22, v15, v11, 0        \n\t"
4694
4695                "vsetvli        t0, x0, e8, m8          \n\t"
4696                "vadd.vx        v16, v16, t2, v0.t      \n\t"
4697
4698                "vsetvli        t0, x0, e8, m1          \n\t"
4699                "vupack.vv      v8, v16, v17, 1         \n\t"
4700                "vupack.vv      v10, v18, v19, 1        \n\t"
4701                "vupack.vv      v12, v20, v21, 1        \n\t"
4702                "vupack.vv      v14, v22, v23, 1        \n\t"
4703
4704                // Convert per-column fp16 scales once; the same scale vector is shared by all 4 rows.
4705                "vsetvli        t0, x0, e16, mf2        \n\t"
4706                "vfwcvt.f.f.v   v18, v2                 \n\t"
4707                "vsetvli        t0, x0, e8, m1          \n\t"
4708
4709                // Row 0: dot(A0, decoded_q5) + a_sum0 * 16 (implicit Q5_0 midpoint correction).
4710                "vle8.v         v1, (t3)                \n\t"
4711                "vsetvli        t0, x0, e64, m1         \n\t"
4712                "vupack.vv      v16, v1, v2, 1          \n\t"
4713                "vsetvli        t0, x0, e32, m1         \n\t"
4714                "vxor.vv        v24, v24, v24           \n\t"
4715                "vxor.vv        v26, v26, v26           \n\t"
4716                "vxor.vv        v28, v28, v28           \n\t"
4717                "vxor.vv        v30, v30, v30           \n\t"
4718                "vmadot         v24, v16, v8, i8        \n\t"
4719                "vmadot         v26, v16, v10, i8       \n\t"
4720                "vmadot         v28, v16, v12, i8       \n\t"
4721                "vmadot         v30, v16, v14, i8       \n\t"
4722                "vmadot         v24, v17, v9, i8        \n\t"
4723                "vmadot         v26, v17, v11, i8       \n\t"
4724                "vmadot         v28, v17, v13, i8       \n\t"
4725                "vmadot         v30, v17, v15, i8       \n\t"
4726                "vpack.vv       v16, v24, v26, 2        \n\t"
4727                "slli           s1, s1, 4               \n\t"
4728                "vpack.vv       v20, v28, v30, 2        \n\t"
4729                "slli           s2, s2, 4               \n\t"
4730                "vpack.vv       v24, v16, v20, 3        \n\t"
4731                "slli           s3, s3, 4               \n\t"
4732                "vpack.vv       v26, v17, v21, 3        \n\t"
4733                "slli           s4, s4, 4               \n\t"
4734                "vadd.vx        v24, v24, s1            \n\t"
4735                "vadd.vx        v25, v25, s2            \n\t"
4736                "vadd.vx        v26, v26, s3            \n\t"
4737                "vadd.vx        v27, v27, s4            \n\t"
4738                "vfcvt.f.x.v    v24, v24                \n\t"
4739                "vfcvt.f.x.v    v25, v25                \n\t"
4740                "vfcvt.f.x.v    v26, v26                \n\t"
4741                "vfcvt.f.x.v    v27, v27                \n\t"
4742                "vfmul.vv       v24, v24, v18           \n\t"
4743                "vfmul.vv       v25, v25, v18           \n\t"
4744                "vfmul.vv       v26, v26, v18           \n\t"
4745                "vfmul.vv       v27, v27, v18           \n\t"
4746                "vfmacc.vf      v4, fa0, v24            \n\t"
4747                "vfmacc.vf      v5, fa1, v25            \n\t"
4748                "vfmacc.vf      v6, fa2, v26            \n\t"
4749                "vfmacc.vf      v7, fa3, v27            \n\t"
4750
4751                "addi           %[BK], %[BK], -1        \n\t"
4752                "bgtz           %[BK], BLK_LOOP%=       \n\t"
4753
4754                // Tail-aware store for the final N tile (`nb_real` may be < 32).
4755                "vsetvli        t0, %[NBLKS], e32, m1   \n\t"
4756                "add            t1, %[LDC], %[DST]      \n\t"
4757                "vse32.v        v4, (%[DST])            \n\t"
4758                "vse32.v        v5, (t1)                \n\t"
4759                "add            t2, t1, %[LDC]          \n\t"
4760                "vse32.v        v6, (t2)                \n\t"
4761                "add            t3, t2, %[LDC]          \n\t"
4762                "vse32.v        v7, (t3)                \n\t"
4763                : [A] "+r"(a_data), [B] "+r"(b_data), [BK] "+r"(cnt)
4764                : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [NBLKS] "r"(nb_real)
4765                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "t5", "t6", "s1", "s2", "s3", "s4", "v0", "v1", "v2",
4766                  "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
4767                  "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31",
4768                  "fa0", "fa1", "fa2", "fa3");
4769        }
4770    }
4771}
4772
4773void gemm_kernel_i8i8_m1(size_t          blk_len,
4774                         const uint8_t * quant_a_ptr,
4775                         const uint8_t * quant_b_data,
4776                         const uint8_t * quant_b_zp,
4777                         float *         c_ptr,
4778                         size_t          count_m,
4779                         size_t          count_n,
4780                         size_t          k_blks,
4781                         size_t          ldc) {
4782    for (size_t n = 0; n < count_n; n += 32) {
4783        size_t    nblks         = (count_n - n) > 32 ? 32 : count_n - n;
4784        uint8_t * QuantBDataPtr = (uint8_t *) quant_b_data +      //
4785                                  n * k_blks * blk_len +          // b data
4786                                  n * k_blks * sizeof(_Float16);  // scale
4787        float * CPtr = c_ptr + n;
4788        size_t  cnt  = k_blks;
4789
4790        // A format Version_1 (FP32 SCALE FOR Normal VMADOTins of IME2)
4791        // A M1K32 int8    256bit
4792        // Ascale fp32 * 1  32bit
4793        // || scl*1(fp32) | Asum(int16) | blk0 || scl*1(fp32) | Asum(int16) | blk0 || ...
4794        // || Element                          || Element                          || ...
4795        // B format
4796        // B N8K32 int4    2048bit
4797        //   4VRF, N32K32, 8192bit
4798        // Bscale fp16 * N32 512bit;
4799        // || scl*32..(fp16) | blk0 blk1 ... blk31 || scl*32..(fp16) | blk0 blk1 ... blk31 || ...
4800        // || Element                              || Element                              || ...
4801
4802        //bias always be nullptr
4803        __asm__ volatile(
4804
4805            // t3 = k/32
4806            "mv           t3, %[BCK]              \n\t"
4807            "mv           t4, %[NBLKS]            \n\t"
4808            "mv           s2, %[pA]               \n\t"  // s2 = pASCL
4809            "addi         s3, %[pA], 4+2          \n\t"  // s3 = pAData, (pA+AScl+ASum)
4810            "mv           s4, %[pB]               \n\t"  // s4 = pBSCL
4811            "addi         s5, %[pB], 32*2         \n\t"  // s5 = pBdata;
4812            "mv           s6, %[pC]               \n\t"
4813
4814            "vsetvli      t0, x0, e32, m1         \n\t"
4815            "vxor.vv      v2, v0, v0              \n\t"  // clear acc
4816
4817            // ordinary vmadot: vle*6 flw*1 vecIns*64 vmadot*8
4818            ".align 4                             \n\t"
4819            "_K_LPST%=:                           \n\t"
4820
4821            "vsetvli      t0, x0, e8, m1          \n\t"
4822            "vl4r.v       v4, (s5)                \n\t"  // B Data 4VRF * 8Row * 32
4823            "addi         s5, s5, 128*4           \n\t"
4824            "vl4r.v       v8, (s5)                \n\t"  // B Data 4VRF * 8Row * 32
4825            "addi         s5, s5, 128*4+64        \n\t"
4826
4827            "vsetvli      t0, x0, e8, mf2         \n\t"
4828            "vle8.v       v0, (s4)                \n\t"  // B Scale 4VRF*8Row*FP16 = 512bit
4829            "addi         s4, s4, 64+128*8        \n\t"
4830
4831            "vsetvli      t0, x0, e8, mf4         \n\t"
4832            "vle8.v       v3, (s3)                \n\t"  // A Data M1*K32*int8 = 256bit
4833            "addi         s3, s3, 32+6            \n\t"
4834
4835            "flw          f0, (s2)                \n\t"  // A Scale fp32
4836            "addi         s2, s2, 6+32            \n\t"  // AScale + Asum(FP32+i16)
4837
4838            "vsetvli      t0, zero, e32, m1       \n\t"
4839            "vupack.vv    v24, v4, v5, 1          \n\t"
4840            "vupack.vv    v26, v6, v7, 1          \n\t"
4841            "vupack.vv    v28, v8, v9, 1          \n\t"
4842            "vupack.vv    v30, v10, v11, 1        \n\t"
4843
4844            "vslidedown.vi  v4, v3, 4             \n\t"
4845
4846            "vxor.vv      v16, v16, v16           \n\t"
4847            "vxor.vv      v18, v16, v16           \n\t"
4848            "vxor.vv      v20, v16, v16           \n\t"
4849            "vxor.vv      v22, v16, v16           \n\t"
4850
4851            "vmadot       v16, v3, v24, i8         \n\t"  // M0 N0 - N7 INT32(256bit)
4852            "vmadot       v18, v3, v26, i8         \n\t"  // M0 N8 - N15
4853            "vmadot       v20, v3, v28, i8         \n\t"  // M0 N16 - N23
4854            "vmadot       v22, v3, v30, i8         \n\t"  // M0 N24 - N31
4855
4856            "vmadot       v16, v4, v25, i8         \n\t"
4857            "vmadot       v18, v4, v27, i8         \n\t"
4858            "vmadot       v20, v4, v29, i8         \n\t"
4859            "vmadot       v22, v4, v31, i8         \n\t"
4860
4861            "vpack.vv     v24, v16, v18, 2        \n\t"
4862            "vpack.vv     v26, v20, v22, 2        \n\t"
4863            "vpack.vv     v16, v24, v26, 3        \n\t"
4864
4865            // b_scale fp16 -> fp32
4866            "vsetvli      t0, x0, e16, mf2        \n\t"
4867            "vfwcvt.f.f.v v24, v0                 \n\t"
4868            // mac result i32 -> fp32
4869            "vsetvli      t0, x0, e32, m1         \n\t"
4870            "vfcvt.f.x.v  v26, v16                \n\t"
4871            // a_scale * b_scale;
4872            "vfmul.vf     v1, v24, f0             \n\t"
4873            // static_cast<float>(qsum) * a_scale * b_scale;
4874            "vfmacc.vv    v2, v1, v26             \n\t"
4875
4876            "addi         t3, t3, -1              \n\t"
4877            "bgtz         t3, _K_LPST%=           \n\t"
4878            "_K_LPND%=:                           \n\t"
4879
4880            //-----------------------------------------
4881            // STORE Equal 32N-------------------------
4882            "_ST32%=:                             \n\t"
4883            "vsetvli      t0, t4, e32, m1         \n\t"
4884            "vse32.v      v2, (s6)                \n\t"  // M0 [N0 : N32]; FP32(1024bit)
4885
4886            "_FUNC_END%=:                         \n\t"
4887
4888            :
4889            : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
4890            : "cc", "t0", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6");
4891    }
4892}
4893
4894void gemm_kernel_i8i8_m4(size_t          blk_len,
4895                         const uint8_t * quant_a_ptr,
4896                         const uint8_t * quant_b_data,
4897                         const uint8_t * quant_b_zp,
4898                         float *         c_ptr,
4899                         size_t          count_m,
4900                         size_t          count_n,
4901                         size_t          k_blks,
4902                         size_t          ldc) {
4903    int64_t b_data_stride = k_blks * sizeof(ggml_fp16_t) + k_blks * blk_len;
4904    for (size_t ni = 0; ni < count_n; ni += 32) {
4905        uint8_t * b_data = (uint8_t *) quant_b_data + ni * b_data_stride;
4906        int8_t *  a_data = (int8_t *) quant_a_ptr;
4907        float *   dst_c  = c_ptr + ni;
4908
4909        asm volatile(
4910            "vsetvli        t0, x0, e32, m1       \n\t"
4911            "vxor.vv        v28, v28, v28         \n\t"
4912            "vxor.vv        v29, v29, v29         \n\t"
4913            "vxor.vv        v30, v30, v30         \n\t"
4914            "vxor.vv        v31, v31, v31         \n\t"
4915
4916            ".align 4                             \n\t"
4917            "BLK_LOOP%=:                          \n\t"
4918            // load scale A
4919            "flw            fa0, (%[A])           \n\t"
4920            "flw            fa1, 4(%[A])          \n\t"
4921            "flw            fa2, 8(%[A])          \n\t"
4922            "flw            fa3, 12(%[A])         \n\t"
4923            "addi           %[A], %[A], 16+8      \n\t"  // Ascl+Asum; FP32*4+i16*4
4924
4925            // load scale B
4926            "vsetvli        t0, x0, e16, mf2      \n\t"
4927            "vle16.v        v12, (%[B])           \n\t"
4928            "addi           %[B], %[B], 64        \n\t"
4929            "vfwcvt.f.f.v   v14, v12              \n\t"
4930
4931            "vsetvli        t0, x0, e8, m1        \n\t"
4932            "vl1r.v         v0, (%[A])            \n\t"
4933            "addi           %[A], %[A], 128       \n\t"  // 4*32@i8
4934            "vl4r.v         v4, (%[B])            \n\t"  // 32*32@i8
4935            "addi           %[B], %[B], 512       \n\t"
4936            "vl4r.v         v8, (%[B])            \n\t"  // 32*32@i8
4937            "addi           %[B], %[B], 512       \n\t"
4938
4939            "vsetvli        t0, zero, e32, m1     \n\t"
4940            "vupack.vv      v2, v0, v0, 1         \n\t"
4941
4942            "vupack.vv      v24, v4, v5, 1        \n\t"
4943            "vupack.vv      v26, v6, v7, 1        \n\t"
4944            "vupack.vv      v4, v8, v9, 1         \n\t"
4945            "vupack.vv      v6, v10, v11, 1       \n\t"
4946
4947            // init the accumu to asum * zp
4948            "vsetvli        t0, x0, e32, m1       \n\t"
4949            "vxor.vv        v16, v16, v16         \n\t"
4950            "vxor.vv        v18, v16, v16         \n\t"
4951            "vxor.vv        v20, v16, v16         \n\t"
4952            "vxor.vv        v22, v16, v16         \n\t"
4953
4954            // i4 * i4 vmadot
4955            "vsetvli        t0, x0, e32, m1       \n\t"
4956            "vmadot         v16, v2, v24, i8      \n\t"
4957            "vmadot         v18, v2, v26, i8      \n\t"
4958            "vmadot         v20, v2, v4, i8       \n\t"
4959            "vmadot         v22, v2, v6, i8       \n\t"
4960            "vmadot         v16, v3, v25, i8      \n\t"
4961            "vmadot         v18, v3, v27, i8      \n\t"
4962            "vmadot         v20, v3, v5, i8       \n\t"
4963            "vmadot         v22, v3, v7, i8       \n\t"
4964
4965            "vpack.vv       v0, v16, v18, 2       \n\t"
4966            "vpack.vv       v2, v20, v22, 2       \n\t"
4967            "vpack.vv       v16, v0, v2, 3        \n\t"
4968            "vpack.vv       v18, v1, v3, 3        \n\t"
4969
4970            "vfcvt.f.x.v    v16, v16              \n\t"
4971            "vfcvt.f.x.v    v17, v17              \n\t"
4972            "vfcvt.f.x.v    v18, v18              \n\t"
4973            "vfcvt.f.x.v    v19, v19              \n\t"
4974
4975            // mul scale
4976            "vfmul.vv       v16, v16, v14         \n\t"
4977            "vfmul.vv       v17, v17, v14         \n\t"
4978            "vfmul.vv       v18, v18, v14         \n\t"
4979            "vfmul.vv       v19, v19, v14         \n\t"
4980
4981            "addi           %[BK], %[BK], -1      \n\t"
4982            "vfmacc.vf      v28, fa0, v16         \n\t"
4983            "vfmacc.vf      v29, fa1, v17         \n\t"
4984            "vfmacc.vf      v30, fa2, v18         \n\t"
4985            "vfmacc.vf      v31, fa3, v19         \n\t"
4986
4987            "bgtz           %[BK], BLK_LOOP%=     \n\t"
4988
4989            // save
4990            "vsetvli        t0, x0, e32, m1       \n\t"
4991            "add            t2, %[LDC], %[DST]    \n\t"
4992            "vse32.v        v28, (%[DST])         \n\t"
4993            "add            t3, %[LDC], t2        \n\t"
4994            "vse32.v        v29, (t2)             \n\t"
4995            "add            t2, %[LDC], t3        \n\t"
4996            "vse32.v        v30, (t3)             \n\t"
4997            "vse32.v        v31, (t2)             \n\t"
4998            : [A] "+r"(a_data), [B] "+r"(b_data)
4999            : [DST] "r"(dst_c), [LDC] "r"(ldc * 4), [BK] "r"(k_blks)
5000            : "t0", "t1", "t2", "t3", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12",
5001              "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27",
5002              "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "fa3");
5003    }
5004}
5005
5006void moe_m2_gemm_kernel_i8i4_impl(size_t           blk_len,
5007                                  const uint8_t ** quant_a_ptr,
5008                                  const uint8_t *  quant_b_data,
5009                                  const uint8_t *  quant_b_zp,
5010                                  float **         c_ptr,
5011                                  size_t           count_m,
5012                                  size_t           count_n,
5013                                  size_t           k_blks,
5014                                  size_t           ldc) {
5015#if 0
5016    moe_gemm_kernel_i8i4_mrow_ref<2, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks,
5017                                     ldc);
5018#else
5019    int64_t b_data_stride =
5020        k_blks * (sizeof(ggml_fp16_t) + 16 * sizeof(int8_t) + (quant_b_zp != NULL ? sizeof(int8_t) : 0));
5021    if (quant_b_zp == NULL) {
5022        for (size_t ni = 0; ni < count_n; ni += 32) {
5023            uint8_t * b_data  = (uint8_t *) quant_b_data + ni * b_data_stride;
5024            int8_t *  a_data0 = (int8_t *) quant_a_ptr[0];
5025            int8_t *  a_data1 = (int8_t *) quant_a_ptr[1];
5026            float *   dst_c0  = (float *) c_ptr[0] + ni;
5027            float *   dst_c1  = (float *) c_ptr[1] + ni;
5028
5029            asm volatile(
5030                "vsetvli        t0, x0, e16, m1         \n\t"
5031                "vxor.vv        v28, v28, v28           \n\t"
5032                "vxor.vv        v29, v29, v29           \n\t"
5033                "vmv.v.i        v0, 1                   \n\t"  // init the scale
5034                "vsll.vi        v1, v0, 4               \n\t"
5035                "vfcvt.f.x.v    v0, v0                  \n\t"
5036                "vfcvt.f.x.v    v1, v1                  \n\t"
5037                "mv             t3, %[BK]               \n\t"
5038
5039                ".align 4                               \n\t"
5040                "BLK_LOOP%=:                            \n\t"
5041                // load scale A0
5042                "flw            fa0, (%[A0])            \n\t"  // A0 scale
5043                "lh             t1, 4(%[A0])            \n\t"  // A0 asum
5044                "addi           %[A0], %[A0], 6         \n\t"
5045
5046                // load scale B
5047                "vsetvli        t0, x0, e16, mf2        \n\t"
5048                "vle16.v        v12, (%[B])             \n\t"
5049                "addi           %[B], %[B], 64          \n\t"
5050                "vsetvli        t0, x0, e16, m1         \n\t"
5051                "vpack.vv       v14, v12, v12, 3        \n\t"
5052
5053                // load scale A1
5054                "flw            fa1, (%[A1])            \n\t"  // A1 scale
5055                "lh             t2, 4(%[A1])            \n\t"  // A1 asum
5056                "addi           %[A1], %[A1], 6         \n\t"
5057                "vsetvli        t0, x0, e16, m1         \n\t"
5058                "vmv.v.x        v10, t1                 \n\t"
5059                "vmv.v.x        v11, t2                 \n\t"
5060
5061                "vpack.vv       v18, v10, v11, 1        \n\t"
5062                "vsll.vi        v18, v18, 3             \n\t"  // mul 8
5063                "vfcvt.f.x.v    v18, v18                \n\t"
5064
5065                "vsetvli        t0, x0, e8, mf4         \n\t"  // A0 data
5066                "vle8.v         v16, (%[A0])            \n\t"
5067                "addi           %[A0], %[A0], 32        \n\t"  // 1*32@i8
5068                "vle8.v         v20, (%[A1])            \n\t"
5069                "addi           %[A1], %[A1], 32        \n\t"  // 1*32@i8
5070
5071                "vl4r.v         v4, (%[B])              \n\t"  // 32*32@i4
5072                "addi           %[B], %[B], 512         \n\t"
5073
5074                "vsrl.vi        v17, v16, 4             \n\t"
5075                "vsrl.vi        v21, v20, 4             \n\t"
5076                "vsetvli        t0, x0, e8, m1          \n\t"
5077                "vnpack4.vv     v2, v16, v20, 2         \n\t"  // low  u4
5078                "vnpack4.vv     v3, v17, v21, 2         \n\t"  // high s4
5079
5080                // init the accumu to asum * zp
5081                "vsetvli        t0, x0, e16, m1         \n\t"
5082                "vor.vv         v19, v18, v18           \n\t"
5083                "vor.vv         v20, v18, v18           \n\t"
5084                "vor.vv         v21, v18, v18           \n\t"
5085
5086                // i4 * i4 vmadot
5087                "vsetvli        t0, x0, e16, m1         \n\t"
5088                "vmadotsu.hp    v18, v3, v4, v1, 0, i4  \n\t"  // high 4
5089                "vmadotsu.hp    v19, v3, v5, v1, 0, i4  \n\t"
5090                "vmadotsu.hp    v20, v3, v6, v1, 0, i4  \n\t"
5091                "vmadotsu.hp    v21, v3, v7, v1, 0, i4  \n\t"
5092                "vmadotu.hp     v18, v2, v4, v0, 0, i4  \n\t"  // low 4
5093                "vmadotu.hp     v19, v2, v5, v0, 0, i4  \n\t"
5094                "vmadotu.hp     v20, v2, v6, v0, 0, i4  \n\t"
5095                "vmadotu.hp     v21, v2, v7, v0, 0, i4  \n\t"
5096
5097                "vpack.vv       v8, v18, v19, 1         \n\t"
5098                "vpack.vv       v12, v20, v21, 1        \n\t"
5099                "vpack.vv       v20, v8, v12, 2         \n\t"
5100
5101                "vfwmul.vv      v16, v20, v14           \n\t"
5102
5103                "vsetvli        t0, x0, e32, m1         \n\t"
5104
5105                "addi           t3, t3, -1              \n\t"
5106                "vfmacc.vf      v28, fa0, v16           \n\t"
5107                "vfmacc.vf      v29, fa1, v17           \n\t"
5108
5109                "bgtz           t3, BLK_LOOP%=          \n\t"
5110
5111                // save
5112                "vsetvli        t0, x0, e32, m1         \n\t"
5113                "vse32.v        v28, (%[DST0])          \n\t"
5114                "vse32.v        v29, (%[DST1])          \n\t"
5115                : [A0] "+r"(a_data0), [A1] "+r"(a_data1), [B] "+r"(b_data)
5116                : [DST0] "r"(dst_c0), [DST1] "r"(dst_c1), [BK] "r"(k_blks)
5117                : "t0", "t1", "t2", "t3", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11",
5118                  "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25",
5119                  "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "fa3");
5120        }
5121    } else {
5122#    if 0
5123        moe_gemm_kernel_i8i4_mrow_ref<2, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5124                                         k_blks, ldc);
5125#    else
5126        for (size_t ni = 0; ni < count_n; ni += 32) {
5127            uint8_t * b_data  = (uint8_t *) quant_b_data + ni * b_data_stride;
5128            int8_t *  a_data0 = (int8_t *) quant_a_ptr[0];
5129            int8_t *  a_data1 = (int8_t *) quant_a_ptr[1];
5130            float *   dst_c0  = (float *) c_ptr[0] + ni;
5131            float *   dst_c1  = (float *) c_ptr[1] + ni;
5132
5133            asm volatile(
5134                "vsetvli        t0, x0, e16, m1         \n\t"
5135                "vxor.vv        v28, v28, v28           \n\t"
5136                "vxor.vv        v29, v29, v29           \n\t"
5137                "vmv.v.i        v0, 1                   \n\t"  // init the scale
5138                "vsll.vi        v1, v0, 4               \n\t"
5139                "vfcvt.f.x.v    v0, v0                  \n\t"
5140                "vfcvt.f.x.v    v1, v1                  \n\t"
5141                "mv             t3, %[BK]               \n\t"
5142
5143                ".align 4                               \n\t"
5144                "BLK_LOOP%=:                            \n\t"
5145                // load scale A0
5146                "flw            fa0, (%[A0])            \n\t"  // A0 scale
5147                "lh             t1, 4(%[A0])            \n\t"  // A0 asum
5148                "addi           %[A0], %[A0], 6         \n\t"
5149
5150                // load scale B
5151                "vsetvli        t0, x0, e16, mf2        \n\t"
5152                "vle16.v        v12, (%[B])             \n\t"
5153                "addi           %[B], %[B], 64          \n\t"
5154                "vsetvli        t0, x0, e16, m1         \n\t"
5155                "vpack.vv       v14, v12, v12, 3        \n\t"
5156
5157                // load scale A1
5158                "flw            fa1, (%[A1])            \n\t"  // A1 scale
5159                "lh             t2, 4(%[A1])            \n\t"  // A1 asum
5160                "addi           %[A1], %[A1], 6         \n\t"
5161
5162                // load zp
5163                "vsetvli        t0, x0, e8, mf4         \n\t"
5164                "vle8.v         v8, (%[B])              \n\t"
5165                "addi           %[B], %[B], 32          \n\t"
5166                "vwaddu.vx      v10, v8, x0             \n\t"
5167
5168                "vsetvli        t0, x0, e8, mf4         \n\t"  // A0 data
5169                "vle8.v         v16, (%[A0])            \n\t"
5170                "addi           %[A0], %[A0], 32        \n\t"  // 1*32@i8
5171                "vle8.v         v20, (%[A1])            \n\t"
5172                "addi           %[A1], %[A1], 32        \n\t"  // 1*32@i8
5173
5174                "vl4r.v         v4, (%[B])              \n\t"  // 32*32@i4
5175                "addi           %[B], %[B], 512         \n\t"
5176
5177                "vsrl.vi        v17, v16, 4             \n\t"
5178                "vsrl.vi        v21, v20, 4             \n\t"
5179                "vsetvli        t0, x0, e8, m1          \n\t"
5180                "vnpack4.vv     v2, v16, v20, 2         \n\t"  // low  u4
5181                "vnpack4.vv     v3, v17, v21, 2         \n\t"  // high s4
5182
5183                // init the accumu to asum * zp
5184                "vsetvli        t0, x0, e16, m1         \n\t"
5185                "vxor.vv        v18, v18, v18           \n\t"
5186                "vxor.vv        v19, v19, v19           \n\t"
5187                "vxor.vv        v20, v20, v20           \n\t"
5188                "vxor.vv        v21, v21, v21           \n\t"
5189
5190                // i4 * i4 vmadot
5191                "vsetvli        t0, x0, e16, m1         \n\t"
5192                "vmadotsu.hp    v18, v3, v4, v1, 0, i4  \n\t"  // high 4
5193                "vmadotsu.hp    v19, v3, v5, v1, 0, i4  \n\t"
5194                "vmadotsu.hp    v20, v3, v6, v1, 0, i4  \n\t"
5195                "vmadotsu.hp    v21, v3, v7, v1, 0, i4  \n\t"
5196                "vmadotu.hp     v18, v2, v4, v0, 0, i4  \n\t"  // low 4
5197                "vmadotu.hp     v19, v2, v5, v0, 0, i4  \n\t"
5198                "vmadotu.hp     v20, v2, v6, v0, 0, i4  \n\t"
5199                "vmadotu.hp     v21, v2, v7, v0, 0, i4  \n\t"
5200
5201                "vpack.vv       v8, v18, v19, 1         \n\t"
5202                "vpack.vv       v12, v20, v21, 1        \n\t"
5203                "vpack.vv       v20, v8, v12, 2         \n\t"
5204                // asum*zp
5205                "vsetvli        t0, x0, e16, mf2        \n\t"
5206                "vwmul.vx       v2, v10, t1             \n\t"
5207                "vwmul.vx       v4, v10, t2             \n\t"
5208
5209                "vsetvli        t0, x0, e32, m1         \n\t"
5210
5211                "vfcvt.f.x.v    v2, v2                  \n\t"
5212                "vfcvt.f.x.v    v4, v4                  \n\t"
5213
5214                "vsetvli        t0, x0, e16, m1         \n\t"
5215                "vfwcvt.f.f.v   v16, v20                \n\t"
5216
5217                "vfwcvt.f.f.v   v18, v14                \n\t"
5218
5219                // +asum*zp
5220                "vsetvli        t0, x0, e32, m1         \n\t"
5221                "vfadd.vv       v16, v16, v2            \n\t"
5222                "vfadd.vv       v17, v17, v4            \n\t"
5223                "vfmul.vv       v16, v16, v18           \n\t"
5224                "vfmul.vv       v17, v17, v18           \n\t"
5225
5226                "addi           t3, t3, -1              \n\t"
5227                "vfmacc.vf      v28, fa0, v16           \n\t"
5228                "vfmacc.vf      v29, fa1, v17           \n\t"
5229
5230                "bgtz           t3, BLK_LOOP%=          \n\t"
5231
5232                // save
5233                "vsetvli        t0, x0, e32, m1         \n\t"
5234                "vse32.v        v28, (%[DST0])          \n\t"
5235                "vse32.v        v29, (%[DST1])          \n\t"
5236                : [A0] "+r"(a_data0), [A1] "+r"(a_data1), [B] "+r"(b_data)
5237                : [DST0] "r"(dst_c0), [DST1] "r"(dst_c1), [BK] "r"(k_blks)
5238                : "t0", "t1", "t2", "t3", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11",
5239                  "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25",
5240                  "v26", "v27", "v28", "v29", "v30", "v31", "fa0", "fa1", "fa2", "fa3");
5241        }
5242#    endif
5243    }
5244#endif
5245}
5246
5247void moe_m2_gemm_kernel_i8i5_impl(size_t           blk_len,
5248                                  const uint8_t ** quant_a_ptr,
5249                                  const uint8_t *  quant_b_data,
5250                                  const uint8_t *  quant_b_zp,
5251                                  float **         c_ptr,
5252                                  size_t           count_m,
5253                                  size_t           count_n,
5254                                  size_t           k_blks,
5255                                  size_t           ldc) {
5256    constexpr size_t NB_COLS          = 32;
5257    constexpr size_t B_Q50_BLK_STRIDE = sizeof(nrow_block_q5_0<NB_COLS>);
5258    constexpr size_t B_Q51_BLK_STRIDE = sizeof(nrow_block_q5_1<NB_COLS>);
5259
5260    GGML_UNUSED(blk_len);
5261    GGML_UNUSED(count_m);
5262    GGML_UNUSED(ldc);
5263
5264    if (quant_b_zp == NULL) {
5265        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
5266            size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
5267            uint8_t * b_data  = (uint8_t *) quant_b_data + (ni / NB_COLS) * k_blks * B_Q50_BLK_STRIDE;
5268            int8_t *  a_data0 = (int8_t *) quant_a_ptr[0];
5269            int8_t *  a_data1 = (int8_t *) quant_a_ptr[1];
5270            float *   dst_c0  = (float *) c_ptr[0] + ni;
5271            float *   dst_c1  = (float *) c_ptr[1] + ni;
5272
5273            asm volatile(
5274                "mv             t4, %[BK]               \n\t"
5275                "vsetvli        t0, x0, e32, m1         \n\t"
5276                "vxor.vv        v2, v0, v0              \n\t"
5277                "vxor.vv        v3, v0, v0              \n\t"
5278
5279                ".align 4                               \n\t"
5280                "BLK_LOOP%=:                            \n\t"
5281                // ---- load B scale/Bh/Bs and advance to the next q5_0 k-block ----
5282                "vsetvli        t0, x0, e8, mf2         \n\t"
5283                "vle8.v         v1, (%[B])              \n\t"  // v1 = scale_fp16 × 32
5284                "addi           %[B], %[B], 64          \n\t"
5285                "vsetvli        t0, x0, e8, m1          \n\t"
5286                "vle8.v         v0, (%[B])              \n\t"  // v0 = Bh N32K32 1-bit packed
5287                "addi           %[B], %[B], 128         \n\t"
5288                "vl4r.v         v8, (%[B])              \n\t"  // v8..v11 = Bs N32K32 i4
5289                "addi           %[B], %[B], 512         \n\t"
5290
5291                // ---- load A0/A1 header then payload, each block stride = 38B ----
5292                "flw            f0, (%[A0])             \n\t"  // f0 = A0 scale (fp32)
5293                "lh             t2, 4(%[A0])            \n\t"  // t2 = A0 asum (int16)
5294                "addi           %[A0], %[A0], 6         \n\t"
5295                "flw            f1, (%[A1])             \n\t"  // f1 = A1 scale (fp32)
5296                "lh             t3, 4(%[A1])            \n\t"  // t3 = A1 asum (int16)
5297                "addi           %[A1], %[A1], 6         \n\t"
5298                "vsetvli        t0, x0, e8, mf4         \n\t"
5299                "vle8.v         v4, (%[A0])             \n\t"  // v4 = A0 M1K32 int8
5300                "addi           %[A0], %[A0], 32        \n\t"
5301                "vle8.v         v5, (%[A1])             \n\t"  // v5 = A1 M1K32 int8
5302                "addi           %[A1], %[A1], 32        \n\t"
5303
5304                //// ---- A nibble unpacking ----
5305                "vsetvli        t0, x0, e8, m1          \n\t"
5306                "vand.vi        v12, v8, 0xF            \n\t"  //8bit(lo4) //[8*32]
5307                "vand.vi        v13, v9, 0xF            \n\t"
5308                "vand.vi        v14, v10, 0xF           \n\t"
5309                "vand.vi        v15, v11, 0xF           \n\t"
5310                "vsrl.vi        v8, v8, 4               \n\t"  //8bit(hi4)
5311                "vsrl.vi        v9, v9, 4               \n\t"
5312                "vsrl.vi        v10, v10, 4             \n\t"
5313                "vsrl.vi        v11, v11, 4             \n\t"
5314
5315                "slli           t2, t2, 4               \n\t"  // a_sum * 16;
5316                "slli           t3, t3, 4               \n\t"
5317                // [4*32]*2
5318                "vsetvli        t0, x0, e8, m1          \n\t"
5319                "vpack.vv       v16, v12, v8, 0         \n\t"
5320                "vpack.vv       v18, v13, v9, 0         \n\t"
5321                "vpack.vv       v20, v14, v10, 0        \n\t"
5322                "vpack.vv       v22, v15, v11, 0        \n\t"
5323
5324                "li             t1, 16                  \n\t"
5325                "vsetvli        t0, x0, e8, m8          \n\t"
5326                "vadd.vx        v16, v16, t1, v0.t      \n\t"
5327
5328                // [4*32]*2 -> [8*16]
5329                "vsetvli        t0, x0, e8, m1          \n\t"
5330                "vupack.vv      v8, v16, v17, 1         \n\t"
5331                "vupack.vv      v10, v18, v19, 1        \n\t"
5332                "vupack.vv      v12, v20, v21, 1        \n\t"
5333                "vupack.vv      v14, v22, v23, 1        \n\t"
5334
5335                "vpack.vv       v6, v4, v5, 2           \n\t"
5336
5337                // init the accumu to asum * zp
5338                "vsetvli        t0, x0, e32, m1         \n\t"
5339                "vxor.vv        v24, v16, v16           \n\t"
5340                "vxor.vv        v26, v16, v16           \n\t"
5341                "vupack.vv      v4, v6, v7, 1           \n\t"
5342                "vxor.vv        v28, v16, v16           \n\t"
5343                "vxor.vv        v30, v16, v16           \n\t"
5344
5345                // ---- i8 main dot products ----
5346                // vmadot: A × unsigned Bh × 16 → fp16 accumulate
5347                "vmadot         v24, v4, v8, i8         \n\t"  // N0..7
5348                "vmadot         v26, v4, v10, i8        \n\t"  // N8..15
5349                "vmadot         v28, v4, v12, i8        \n\t"  // N16..23
5350                "vmadot         v30, v4, v14, i8        \n\t"  // N24..31
5351                // vmadot: A × unsigned Bh × 1 → fp16 accumulate
5352                "vmadot         v24, v5, v9, i8         \n\t"  // N0..7
5353                "vmadot         v26, v5, v11, i8        \n\t"  // N8..15
5354                "vmadot         v28, v5, v13, i8        \n\t"  // N16..23
5355                "vmadot         v30, v5, v15, i8        \n\t"  // N24..31
5356
5357                "vpack.vv       v16, v24, v26, 2        \n\t"  // v16 = N0..15
5358                "vpack.vv       v18, v28, v30, 2        \n\t"  // v18 = N16..31
5359                "vpack.vv       v24, v16, v18, 3        \n\t"  // v24 = N0..31
5360
5361                "vadd.vx        v24, v24, t2            \n\t"
5362                "vadd.vx        v25, v25, t3            \n\t"
5363                // b_scale fp16 -> fp32
5364                "vsetvli        t0, x0, e16, mf2        \n\t"
5365                "vfwcvt.f.f.v   v28, v1                 \n\t"
5366
5367                // a_scale * b_scale;
5368                "vsetvli        t0, x0, e32, m1         \n\t"
5369                "vfcvt.f.x.v    v26, v24                \n\t"
5370                "vfcvt.f.x.v    v27, v25                \n\t"
5371                "vfmul.vf       v30, v28, f0            \n\t"
5372                "vfmul.vf       v31, v28, f1            \n\t"
5373                // static_cast<float>(qsum) * a_scale * b_scale;
5374                "vfmacc.vv      v2, v30, v26            \n\t"
5375                "vfmacc.vv      v3, v31, v27            \n\t"
5376
5377                "addi           t4, t4, -1              \n\t"
5378                "bgtz           t4, BLK_LOOP%=          \n\t"
5379
5380                "vsetvli        t0, %[NR], e32, m1      \n\t"
5381                "vse32.v        v2, (%[DST0])           \n\t"
5382                "vse32.v        v3, (%[DST1])           \n\t"
5383                : [A0] "+r"(a_data0), [A1] "+r"(a_data1), [B] "+r"(b_data)
5384                : [DST0] "r"(dst_c0), [DST1] "r"(dst_c1), [BK] "r"(k_blks), [NR] "r"(nb_real)
5385                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8",
5386                  "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22",
5387                  "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "f0", "f1");
5388        }
5389    } else {
5390        for (size_t ni = 0; ni < count_n; ni += NB_COLS) {
5391            size_t    nb_real = std::min<size_t>(NB_COLS, count_n - ni);
5392            uint8_t * b_data  = (uint8_t *) quant_b_data + (ni / NB_COLS) * k_blks * B_Q51_BLK_STRIDE;
5393            int8_t *  a_data0 = (int8_t *) quant_a_ptr[0];
5394            int8_t *  a_data1 = (int8_t *) quant_a_ptr[1];
5395            float *   dst_c0  = (float *) c_ptr[0] + ni;
5396            float *   dst_c1  = (float *) c_ptr[1] + ni;
5397
5398            asm volatile(
5399                "mv             t4, %[BK]               \n\t"
5400                "vsetvli        t0, x0, e32, m1         \n\t"
5401                "vxor.vv        v2, v0, v0              \n\t"
5402                "vxor.vv        v3, v0, v0              \n\t"
5403                "addi           t5, %[B], 64            \n\t"  // t5 = zp   (32B)
5404                "addi           t6, %[B], 96            \n\t"  // t6 = qh   (128B)
5405                "addi           s1, %[B], 224           \n\t"  // s1 = qs   (512B)
5406
5407                ".align 4                               \n\t"
5408                "BLK_LOOP%=:                            \n\t"
5409                // ---- load B scale/zp/Bh/Bs and advance to the next q5_1 k-block ----
5410                "vsetvli        t0, x0, e8, mf2         \n\t"
5411                "vle8.v         v1, (%[B])              \n\t"  // v1 = scale_fp16 × 32
5412                "addi           %[B], %[B], 736         \n\t"
5413                "vsetvli        t0, x0, e8, m1          \n\t"
5414                "vle8.v         v0, (t6)                \n\t"  // v0 = Bh N32K32 1-bit packed
5415                "addi           t6, t6, 736             \n\t"
5416                "vl4r.v         v8, (s1)                \n\t"  // v8..v11 = Bs N32K32 i4
5417                "addi           s1, s1, 736             \n\t"
5418
5419                // ---- load A0/A1 header then payload, each block stride = 38B ----
5420                "flw            f0, (%[A0])             \n\t"  // f0 = A0 scale (fp32)
5421                "lh             t2, 4(%[A0])            \n\t"  // t2 = A0 asum (int16)
5422                "addi           %[A0], %[A0], 6         \n\t"
5423                "flw            f1, (%[A1])             \n\t"  // f1 = A1 scale (fp32)
5424                "lh             t3, 4(%[A1])            \n\t"  // t3 = A1 asum (int16)
5425                "addi           %[A1], %[A1], 6         \n\t"
5426                "vsetvli        t0, x0, e8, mf4         \n\t"
5427                "vle8.v         v4, (%[A0])             \n\t"  // v4 = A0 M1K32 int8
5428                "addi           %[A0], %[A0], 32        \n\t"
5429                "vle8.v         v5, (%[A1])             \n\t"  // v5 = A1 M1K32 int8
5430                "addi           %[A1], %[A1], 32        \n\t"
5431
5432                //// ---- A nibble unpacking ----
5433                "vsetvli        t0, x0, e8, m1          \n\t"
5434                "vand.vi        v12, v8, 0xF            \n\t"  //8bit(lo4) //[8*32]
5435                "vand.vi        v13, v9, 0xF            \n\t"
5436                "vand.vi        v14, v10, 0xF           \n\t"
5437                "vand.vi        v15, v11, 0xF           \n\t"
5438                "vsrl.vi        v8, v8, 4               \n\t"  //8bit(hi4)
5439                "vsrl.vi        v9, v9, 4               \n\t"
5440                "vsrl.vi        v10, v10, 4             \n\t"
5441                "vsrl.vi        v11, v11, 4             \n\t"
5442
5443                // q5_1 uses explicit zp, so keep a_sum unshifted here.
5444                // [4*32]*2
5445                "vpack.vv       v16, v12, v8, 0         \n\t"
5446                "vpack.vv       v18, v13, v9, 0         \n\t"
5447                "vpack.vv       v20, v14, v10, 0        \n\t"
5448                "vpack.vv       v22, v15, v11, 0        \n\t"
5449
5450                "li             t1, 16                  \n\t"
5451                "vsetvli        t0, x0, e8, m8          \n\t"
5452                "vadd.vx        v16, v16, t1, v0.t      \n\t"
5453
5454                // [4*32]*2 -> [8*16]
5455                "vsetvli        t0, x0, e8, m1          \n\t"
5456                "vupack.vv      v8, v16, v17, 1         \n\t"
5457                "vupack.vv      v10, v18, v19, 1        \n\t"
5458                "vupack.vv      v12, v20, v21, 1        \n\t"
5459                "vupack.vv      v14, v22, v23, 1        \n\t"
5460
5461                "vpack.vv       v6, v4, v5, 2           \n\t"
5462
5463                // init the accumu to asum * zp
5464                "vsetvli        t0, x0, e32, m1         \n\t"
5465                "vxor.vv        v24, v16, v16           \n\t"
5466                "vxor.vv        v26, v16, v16           \n\t"
5467                "vupack.vv      v4, v6, v7, 1           \n\t"
5468                "vxor.vv        v28, v16, v16           \n\t"
5469                "vxor.vv        v30, v16, v16           \n\t"
5470
5471                // ---- i8 main dot products ----
5472                // vmadot: A × unsigned Bh × 16 → fp16 accumulate
5473                "vmadot         v24, v4, v8, i8         \n\t"  // N0..7
5474                "vmadot         v26, v4, v10, i8        \n\t"  // N8..15
5475                "vmadot         v28, v4, v12, i8        \n\t"  // N16..23
5476                "vmadot         v30, v4, v14, i8        \n\t"  // N24..31
5477                // vmadot: A × unsigned Bh × 1 → fp16 accumulate
5478                "vmadot         v24, v5, v9, i8         \n\t"  // N0..7
5479                "vmadot         v26, v5, v11, i8        \n\t"  // N8..15
5480                "vmadot         v28, v5, v13, i8        \n\t"  // N16..23
5481                "vmadot         v30, v5, v15, i8        \n\t"  // N24..31
5482
5483                "vsetvli        t0, x0, e8, mf4         \n\t"
5484                "vle8.v         v4, (t5)                \n\t"  // v4 = Bzp N32 uint8
5485                "addi           t5, t5, 736             \n\t"
5486
5487                "vsetvli        t0, x0, e8, m1          \n\t"
5488                "vpack.vv       v16, v24, v26, 2        \n\t"  // v16 = N0..15
5489                "vpack.vv       v18, v28, v30, 2        \n\t"  // v18 = N16..31
5490                "vpack.vv       v24, v16, v18, 3        \n\t"  // v24 = N0..31
5491
5492                "vsetvli        t0, x0, e8, mf4         \n\t"
5493                "vwaddu.vx      v28, v4, x0             \n\t"
5494
5495                "vsetvli        t0, x0, e16, mf2        \n\t"
5496                "vwmul.vx       v30, v28, t2            \n\t"
5497                "vwmul.vx       v31, v28, t3            \n\t"
5498
5499                // b_scale fp16 -> fp32
5500                "vfwcvt.f.f.v   v28, v1                 \n\t"
5501
5502                "vsetvli        t0, x0, e32, m1         \n\t"
5503                "vadd.vv        v24, v24, v30           \n\t"
5504                "vadd.vv        v25, v25, v31           \n\t"
5505
5506                // a_scale * b_scale;
5507                "vfcvt.f.x.v    v26, v24                \n\t"
5508                "vfcvt.f.x.v    v27, v25                \n\t"
5509                "vfmul.vf       v30, v28, f0            \n\t"
5510                "vfmul.vf       v31, v28, f1            \n\t"
5511                // static_cast<float>(qsum) * a_scale * b_scale;
5512                "vfmacc.vv      v2, v30, v26            \n\t"
5513                "vfmacc.vv      v3, v31, v27            \n\t"
5514
5515                "addi           t4, t4, -1              \n\t"
5516                "bgtz           t4, BLK_LOOP%=          \n\t"
5517
5518                "vsetvli        t0, %[NR], e32, m1      \n\t"
5519                "vse32.v        v2, (%[DST0])           \n\t"
5520                "vse32.v        v3, (%[DST1])           \n\t"
5521                : [A0] "+r"(a_data0), [A1] "+r"(a_data1), [B] "+r"(b_data)
5522                : [DST0] "r"(dst_c0), [DST1] "r"(dst_c1), [BK] "r"(k_blks), [NR] "r"(nb_real)
5523                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "t5", "t6", "s1", "v0", "v1", "v2", "v3", "v4", "v5",
5524                  "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", "v20",
5525                  "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "f0", "f1");
5526        }
5527    }
5528}
5529
5530size_t gemm_kernel_i8i2k(size_t          blk_len,
5531                         const uint8_t * quant_a_ptr,
5532                         const uint8_t * quant_b_data,
5533                         const uint8_t * quant_b_zp,
5534                         float *         c_ptr,
5535                         size_t          count_m,
5536                         size_t          count_n,
5537                         size_t          k_blks,
5538                         size_t          ldc) {
5539    if (count_m >= 4) {
5540#if 0
5541        gemm_kernel_i8i2k_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5542#else
5543        gemm_kernel_i8i2k_m4(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5544#endif
5545        return 4;
5546    } else {
5547#if 0
5548        gemm_kernel_i8i2k_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks,
5549                                      ldc);
5550#else
5551        gemm_kernel_i8i2k_m1(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5552#endif
5553        return 1;
5554    }
5555}
5556
5557size_t gemm_kernel_i8i3k(size_t          blk_len,
5558                         const uint8_t * quant_a_ptr,
5559                         const uint8_t * quant_b_data,
5560                         const uint8_t * quant_b_zp,
5561                         float *         c_ptr,
5562                         size_t          count_m,
5563                         size_t          count_n,
5564                         size_t          k_blks,
5565                         size_t          ldc) {
5566    if (count_m >= 4) {
5567#if 0
5568        gemm_kernel_i8i3k_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5569#else
5570        gemm_kernel_i8i3k_m4(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5571#endif
5572        return 4;
5573    } else {
5574#if 0
5575        gemm_kernel_i8i3k_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5576#else
5577        gemm_kernel_i8i3k_m1(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_m, count_n, k_blks, ldc);
5578#endif
5579        return 1;
5580    }
5581}
5582
5583size_t gemm_kernel_i8i4(size_t          blk_len,
5584                        const uint8_t * quant_a_ptr,
5585                        const uint8_t * quant_b_data,
5586                        const uint8_t * quant_b_zp,
5587                        float *         c_ptr,
5588                        size_t          count_m,
5589                        size_t          count_n,
5590                        size_t          k_blks,
5591                        size_t          ldc) {
5592    if (count_m >= 4) {
5593#if 0
5594        gemm_kernel_i8i4_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5595                                         k_blks, ldc);
5596#else
5597        gemm_kernel_i8i4_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5598#endif
5599        return 4;
5600    } else {
5601#if 0
5602        gemm_kernel_i8i4_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5603                                         k_blks, ldc);
5604#else
5605        gemm_kernel_i8i4_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5606#endif
5607        return 1;
5608    }
5609}
5610
5611size_t gemm_kernel_i8i4_hp(size_t          blk_len,
5612                           const uint8_t * quant_a_ptr,
5613                           const uint8_t * quant_b_data,
5614                           const uint8_t * quant_b_zp,
5615                           float *         c_ptr,
5616                           size_t          count_m,
5617                           size_t          count_n,
5618                           size_t          k_blks,
5619                           size_t          ldc) {
5620    if (count_m >= 4) {
5621#if 0
5622        gemm_kernel_i8i4_hp_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5623                                            k_blks, ldc);
5624#else
5625        gemm_kernel_i8i4_hp_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5626#endif
5627        return 4;
5628    } else {
5629#if 0
5630        gemm_kernel_i8i4_hp_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5631                                            k_blks, ldc);
5632#else
5633        gemm_kernel_i8i4_hp_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5634#endif
5635        return 1;
5636    }
5637}
5638
5639size_t moe_m2_gemm_kernel_i8i4(size_t           blk_len,
5640                               const uint8_t ** quant_a_ptr,
5641                               const uint8_t *  quant_b_data,
5642                               const uint8_t *  quant_b_zp,
5643                               float **         c_ptr,
5644                               size_t           count_m,
5645                               size_t           count_n,
5646                               size_t           k_blks,
5647                               size_t           ldc) {
5648    moe_m2_gemm_kernel_i8i4_impl(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5649    return 2;
5650}
5651
5652size_t gemm_kernel_i8i8(size_t          blk_len,
5653                        const uint8_t * quant_a_ptr,
5654                        const uint8_t * quant_b_data,
5655                        const uint8_t * quant_b_zp,
5656                        float *         c_ptr,
5657                        size_t          count_m,
5658                        size_t          count_n,
5659                        size_t          k_blks,
5660                        size_t          ldc) {
5661    if (count_m >= 4) {
5662#if 0
5663        gemm_kernel_i8i8_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5664                                         k_blks, ldc);
5665#else
5666        gemm_kernel_i8i8_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5667#endif
5668        return 4;
5669    } else {
5670#if 0
5671        gemm_kernel_i8i8_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5672                                         k_blks, ldc);
5673#else
5674        gemm_kernel_i8i8_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5675#endif
5676        return 1;
5677    }
5678}
5679
5680size_t gemm_kernel_i8mxfp4(size_t          blk_len,
5681                           const uint8_t * quant_a_ptr,
5682                           const uint8_t * quant_b_data,
5683                           const uint8_t * quant_b_zp,
5684                           float *         c_ptr,
5685                           size_t          count_m,
5686                           size_t          count_n,
5687                           size_t          k_blks,
5688                           size_t          ldc) {
5689    if (count_m >= 4) {
5690#if 1
5691        gemm_kernel_i8mxfp4_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5692                                            k_blks, ldc);
5693#else
5694        gemm_kernel_i8mxfp4_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5695#endif
5696        return 4;
5697    } else {
5698#if 1
5699        gemm_kernel_i8mxfp4_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5700                                            k_blks, ldc);
5701#else
5702        gemm_kernel_i8mxfp4_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5703#endif
5704        return 1;
5705    }
5706}
5707
5708size_t moe_m2_gemm_kernel_i8mxfp4(size_t           blk_len,
5709                                  const uint8_t ** quant_a_ptr,
5710                                  const uint8_t *  quant_b_data,
5711                                  const uint8_t *  quant_b_zp,
5712                                  float **         c_ptr,
5713                                  size_t           count_m,
5714                                  size_t           count_n,
5715                                  size_t           k_blks,
5716                                  size_t           ldc) {
5717    //moe_m2_gemm_kernel_i8mxfp4_impl(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5718    return 2;
5719}
5720
5721size_t gemm_kernel_i8i5(size_t          blk_len,
5722                        const uint8_t * quant_a_ptr,
5723                        const uint8_t * quant_b_data,
5724                        const uint8_t * quant_b_zp,
5725                        float *         c_ptr,
5726                        size_t          count_m,
5727                        size_t          count_n,
5728                        size_t          k_blks,
5729                        size_t          ldc) {
5730    if (count_m >= 4) {
5731#if 0
5732        gemm_kernel_i8i5_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5733                                         k_blks, ldc);
5734#else
5735        gemm_kernel_i8i5_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5736#endif
5737        return 4;
5738    } else {
5739#if 0
5740        gemm_kernel_i8i5_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5741                                         k_blks, ldc);
5742#else
5743        gemm_kernel_i8i5_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5744#endif
5745        return 1;
5746    }
5747}
5748
5749size_t moe_m2_gemm_kernel_i8i5(size_t           blk_len,
5750                               const uint8_t ** quant_a_ptr,
5751                               const uint8_t *  quant_b_data,
5752                               const uint8_t *  quant_b_zp,
5753                               float **         c_ptr,
5754                               size_t           count_m,
5755                               size_t           count_n,
5756                               size_t           k_blks,
5757                               size_t           ldc) {
5758#if 0
5759    moe_gemm_kernel_i8i5_mrow_ref<2, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
5760                                         k_blks, ldc);
5761#else
5762    moe_m2_gemm_kernel_i8i5_impl(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
5763#endif
5764    return 2;
5765}
5766
5767}  // namespace ime2
5768}  // namespace spacemit_kernels