hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

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

Hao-Chen2337ggml-cpu : fix missing Q5_0 dispatch in SpaceMiT backend (#26792)08659901c2026-08-09

1#define GGML_COMMON_IMPL_CPP
2#define GGML_COMMON_DECL_CPP
3
4#include "ime.h"
5
6#include "binary-ops.h"
7#include "common.h"
8#include "ggml-backend-impl.h"
9#include "ggml-common.h"
10#include "ggml-cpu.h"
11#include "ime_env.h"
12#include "ime_kernels.h"
13#include "ops.h"
14#include "repack.h"
15#include "rvv_kernels.h"
16#include "spine_mem_pool.h"
17#include "traits.h"
18#include "vec.h"
19
20#include <fcntl.h>
21#include <sys/mman.h>
22#include <unistd.h>
23
24#include <algorithm>
25#include <atomic>
26#include <cassert>
27#include <cerrno>
28#include <cmath>
29#include <cstdio>  // for GGML_ASSERT
30#include <stdexcept>
31#include <thread>
32// clang-format off
33#if defined(__riscv)
34
35#if !defined(__riscv_v) || !defined(__riscv_v_intrinsic)
36#error "riscv v extension or v_intrinsic not enabled"
37#else
38#include <riscv_vector.h>
39#endif
40
41#if !defined(__riscv_zfh) || !defined(__riscv_zvfh)
42#error "riscv zfh extension not enabled, GGML_RV_ZFH and GGML_RV_ZVFH must be defined to 1"
43#endif
44
45#if !defined(__riscv_zba)
46#error "riscv zba extension not enabled, GGML_RV_ZBA must be defined to 1"
47#endif
48
49#if defined(RISCV64_SPACEMIT_IME1) || defined(RISCV64_SPACEMIT_IME2)
50#else
51#error "RISCV64_SPACEMIT_IME1 or RISCV64_SPACEMIT_IME2 not defined"
52#endif
53
54#else
55
56#error "riscv not enabled in this build"
57
58#endif
59
60#if defined(__GNUC__)
61#pragma GCC diagnostic ignored "-Woverlength-strings"
62#pragma GCC diagnostic ignored "-Wcast-qual"
63#pragma GCC diagnostic ignored "-Wunused-parameter"
64#endif
65
66// clang-format on
67
68extern "C" {
69extern void ggml_threadpool_chunk_set(struct ggml_threadpool * tp, int value);
70extern int  ggml_threadpool_chunk_add(struct ggml_threadpool * tp, int value);
71}
72
73namespace ggml::cpu::riscv64_spacemit {
74
75struct TLSContext {
76    int       cpu_id{ -1 };
77    cpu_set_t cpuset;
78    void *    tcm_buffer{ nullptr };
79    size_t    tcm_buffer_size{ 0 };
80};
81
82thread_local TLSContext tls_context;
83
84template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> constexpr size_t get_repacked_block_type_size() {
85    if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0>) {
86        return sizeof(block_q8_0);
87    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0>) {
88        return sizeof(block_q4_0) * INTER_SIZE / QK4_0;
89    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_1> || std::is_same_v<BLOC_TYPE, block_q4_K>) {
90        return (sizeof(block_q4_0) + sizeof(uint8_t)) * INTER_SIZE / QK4_1;
91    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K>) {
92        return sizeof(spacemit_kernels::nrow_block_q2_k<1>);
93    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q3_K>) {
94        return sizeof(spacemit_kernels::nrow_block_q3_k<1>);
95    } else if constexpr (std::is_same_v<BLOC_TYPE, block_mxfp4>) {
96        return sizeof(spacemit_kernels::nrow_block_mxfp4<1>);
97    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_1> || std::is_same_v<BLOC_TYPE, block_q5_K>) {
98        return sizeof(spacemit_kernels::nrow_block_q5_1<1>);
99    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_0>) {
100        return sizeof(spacemit_kernels::nrow_block_q5_0<1>);
101    } else {
102        assert(false);
103        return 0;
104    }
105}
106
107template <typename BLOC_TYPE> constexpr bool block_type_has_zp() {
108    if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0> ||
109                  std::is_same_v<BLOC_TYPE, block_q3_K> || std::is_same_v<BLOC_TYPE, block_q4_0> ||
110                  std::is_same_v<BLOC_TYPE, block_mxfp4> || std::is_same_v<BLOC_TYPE, block_q5_0>) {
111        return false;
112    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_1> || std::is_same_v<BLOC_TYPE, block_q4_K> ||
113                         std::is_same_v<BLOC_TYPE, block_q2_K> || std::is_same_v<BLOC_TYPE, block_q5_1> ||
114                         std::is_same_v<BLOC_TYPE, block_q5_K>) {
115        return true;
116    } else {
117        assert(false);
118        return false;
119    }
120}
121
122class tensor_traits_base : public ggml::cpu::tensor_traits {
123  public:
124    virtual int repack(ggml_tensor * t, const void * data, size_t data_size) = 0;
125};
126
127template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_traits : public tensor_traits_base {
128    bool work_size(int /* n_threads */, const ggml_tensor * op, size_t & size) override {
129        switch (op->op) {
130            case GGML_OP_MUL_MAT:
131                {
132                    int64_t src1_nelements = ggml_nelements(op->src[1]);
133
134                    if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K> || std::is_same_v<BLOC_TYPE, block_q3_K>) {
135                        size =
136                            spacemit_kernels::div_round_up(src1_nelements, QK_K) * spacemit_kernels::q8k_blk_size(QK_K);
137                    } else if constexpr (INTER_SIZE == QK4_0) {
138                        size = spacemit_kernels::div_round_up(src1_nelements, QK4_0) *
139                               spacemit_kernels::q8_blk_size(QK4_0, true);
140                    } else if constexpr (INTER_SIZE == 256) {
141                        size = spacemit_kernels::div_round_up(src1_nelements, 256) *
142                               spacemit_kernels::q8_hp_blk_size(256, true, true);
143                    } else {
144                        GGML_ABORT("unsupported block type");
145                    }
146
147                    size = GGML_PAD(size, sizeof(int64_t));
148
149                    return true;
150                }
151            case GGML_OP_MUL_MAT_ID:
152                {
153                    int64_t src1_nelements = ggml_nelements(op->src[1]);
154
155                    if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K> || std::is_same_v<BLOC_TYPE, block_q3_K>) {
156                        size =
157                            spacemit_kernels::div_round_up(src1_nelements, QK_K) * spacemit_kernels::q8k_blk_size(QK_K);
158                    } else if constexpr (INTER_SIZE == QK4_0) {
159                        size = spacemit_kernels::div_round_up(src1_nelements, QK4_0) *
160                               spacemit_kernels::q8_blk_size(QK4_0, true);
161                    } else if constexpr (INTER_SIZE == 256) {
162                        size = spacemit_kernels::div_round_up(src1_nelements, 256) *
163                               spacemit_kernels::q8_hp_blk_size(256, true, true);
164                    } else {
165                        GGML_ABORT("unsupported block type");
166                    }
167
168                    size = GGML_PAD(size, sizeof(int64_t));
169
170                    const int64_t ne02 = op->src[0]->ne[2];  // n_as, n_expert
171                    const int64_t ne12 = op->src[1]->ne[2];  // n_tokens
172
173                    const size_t sizeof_mmid_row_mapping = sizeof(int64_t);
174                    size += sizeof_mmid_row_mapping * ne02 * (ne12 + 1) + (ne02 + 1) * sizeof(int64_t);
175
176                    size = GGML_PAD(size, sizeof(int64_t));
177
178                    return true;
179                }
180            default:
181                // GGML_ABORT("fatal error");
182                break;
183        }
184        return false;
185    }
186
187    bool compute_forward(ggml_compute_params * params, ggml_tensor * op) override {
188        switch (op->op) {
189            case GGML_OP_MUL_MAT:
190                switch (op->src[0]->type) {
191                    case GGML_TYPE_Q2_K:
192                    case GGML_TYPE_Q3_K:
193                    case GGML_TYPE_Q4_0:
194                    case GGML_TYPE_Q4_1:
195                    case GGML_TYPE_Q4_K:
196                    case GGML_TYPE_Q6_K:
197                    case GGML_TYPE_Q8_0:
198                    case GGML_TYPE_Q5_0:
199                    case GGML_TYPE_Q5_1:
200                    case GGML_TYPE_Q5_K:
201                        //case GGML_TYPE_MXFP4:
202                        forward_mul_mat(params, op);
203                        return true;
204                    default:
205                        // GGML_ABORT("fatal error: unsupported type for src0 in MUL_MAT");
206                        return false;
207                }
208                break;
209            case GGML_OP_MUL_MAT_ID:
210                switch (op->src[0]->type) {
211                    case GGML_TYPE_Q2_K:
212                    case GGML_TYPE_Q3_K:
213                    case GGML_TYPE_Q4_0:
214                    case GGML_TYPE_Q4_1:
215                    case GGML_TYPE_Q4_K:
216                    case GGML_TYPE_Q6_K:
217                    case GGML_TYPE_Q8_0:
218                    case GGML_TYPE_Q5_0:
219                    case GGML_TYPE_Q5_1:
220                    case GGML_TYPE_Q5_K:
221                        //case GGML_TYPE_MXFP4:
222                        forward_mul_mat_id(params, op);
223                        return true;
224                    default:
225                        // GGML_ABORT("fatal error: unsupported type for src0 in MUL_MAT_ID");
226                        return false;
227                }
228                break;
229            default:
230                // GGML_ABORT("fatal error");
231                break;
232        }
233        return false;
234    }
235
236    void forward_mul_mat(ggml_compute_params * params, ggml_tensor * op) {
237        constexpr size_t a_blk_len = INTER_SIZE;
238        constexpr size_t b_blk_len = INTER_SIZE;
239
240        const ggml_tensor * src0 = op->src[0];
241        const ggml_tensor * src1 = op->src[1];
242        ggml_tensor *       dst  = op;
243
244        GGML_TENSOR_BINARY_OP_LOCALS
245
246        int ith = params->ith;
247        int nth = params->nth;
248
249        [[maybe_unused]] const enum ggml_type type = src0->type;
250
251        void *        w_data  = (void *) src0->data;
252        const float * feature = (const float *) src1->data;
253        float *       output  = (float *) dst->data;
254
255        const int64_t gemm_m = ne11 * ne12 * ne13;
256        const int64_t gemm_k = ne10;
257        const int64_t gemm_n = ne01;
258
259        spacemit_kernels::quantize_a_row_def       quantize_a_row_i8;
260        spacemit_kernels::quantize_a_row_def       quantize_a_4row_i8;
261        spacemit_kernels::gemm_kernel_quantize_def gemm_kernel;
262        bool                                       set_kernel_impl = false;
263
264        int64_t block_stride_a = spacemit_kernels::q8_blk_size(a_blk_len);
265
266#if defined(RISCV64_SPACEMIT_IME2)
267        if (!set_kernel_impl && (global_spine_env_info.use_ime2)) {
268            quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8;
269            quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8;
270            block_stride_a     = spacemit_kernels::q8_blk_size(a_blk_len, true);
271
272            if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0>) {
273                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i8;
274                set_kernel_impl = true;
275            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
276                                 std::is_same_v<BLOC_TYPE, block_q4_K>) {
277                if constexpr (INTER_SIZE == 256) {
278                    gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i4_hp;
279                    quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8_hp;
280                    quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8_hp;
281                    block_stride_a     = spacemit_kernels::q8_hp_blk_size(a_blk_len, true, true);
282                    set_kernel_impl    = true;
283                } else {
284                    gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i4;
285                    quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8;
286                    quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8;
287                    block_stride_a     = spacemit_kernels::q8_blk_size(a_blk_len, true);
288                    set_kernel_impl    = true;
289                }
290            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K>) {
291                quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8k;
292                quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8k;
293                block_stride_a     = spacemit_kernels::q8k_blk_size(a_blk_len);
294
295                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i2k;
296                set_kernel_impl = true;
297            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q3_K>) {
298                quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8k;
299                quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8k;
300                block_stride_a     = spacemit_kernels::q8k_blk_size(a_blk_len);
301
302                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i3k;
303                set_kernel_impl = true;
304            } else if constexpr (std::is_same_v<BLOC_TYPE, block_mxfp4>) {
305                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8mxfp4;
306                set_kernel_impl = true;
307            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_1> || std::is_same_v<BLOC_TYPE, block_q5_K> ||
308                                 std::is_same_v<BLOC_TYPE, block_q5_0>) {
309                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i5;
310                set_kernel_impl = true;
311            }
312        }
313#endif
314
315#if defined(RISCV64_SPACEMIT_IME1)
316        if (!set_kernel_impl && (global_spine_env_info.use_ime1)) {
317            quantize_a_row_i8  = spacemit_kernels::ime1::quantize_a_row_i8;
318            quantize_a_4row_i8 = spacemit_kernels::ime1::quantize_a_4row_i8;
319
320            if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
321                          std::is_same_v<BLOC_TYPE, block_q4_K>) {
322                gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i4;
323                set_kernel_impl = true;
324            }
325        }
326#endif
327        if (!set_kernel_impl) {
328            GGML_ABORT("no kernel implementation found for the block type");
329        }
330
331        const int64_t a_k_blks = spacemit_kernels::div_round_up(gemm_k, a_blk_len);
332        const int64_t b_k_blks = spacemit_kernels::div_round_up(gemm_k, b_blk_len);
333
334        const int64_t row_stride_a        = a_k_blks * block_stride_a;
335        const int64_t gemm_workspace_size = GGML_PAD(gemm_m * row_stride_a, alignof(int64_t));
336
337        if (ith == 0 && params->wsize < gemm_workspace_size) {
338            GGML_ABORT("wsize less than gemm_workspace_size");
339        }
340
341        uintptr_t ws_ptr = reinterpret_cast<uintptr_t>(params->wdata);
342
343        void *        tcm_buffer      = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer;
344        const int64_t tcm_buffer_size = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size;
345
346        auto * quant_a_buffer = reinterpret_cast<uint8_t *>(ws_ptr);
347
348        constexpr int64_t row_align = 4;
349        const int64_t     row_blks  = spacemit_kernels::div_round_up(gemm_m, row_align);
350
351        const int64_t row_stride_b      = b_k_blks * get_repacked_block_type_size<BLOC_TYPE, INTER_SIZE, NB_COLS>();
352        const int64_t per_mb_rows_wsize = row_align * row_stride_a;
353        const int64_t per_nb_cols_wsize = NB_COLS * row_stride_b;
354
355        const int64_t barrier_idx = static_cast<int64_t>(ith / 2);
356
357        GGML_ASSERT(global_spine_env_info.init_barrier != nullptr);
358        GGML_ASSERT(barrier_idx < spine_init_barrier_count);
359        spine_barrier_t * cur_barrier = &global_spine_env_info.init_barrier[barrier_idx];
360
361        if (gemm_m == 1) {
362            int task_per_thread = spacemit_kernels::div_round_up(a_k_blks, nth);
363            int a_blk_start     = ith * task_per_thread;
364            int a_blk_end       = std::min(a_blk_start + task_per_thread, (int) a_k_blks);
365            if (a_blk_start < a_blk_end) {
366                quantize_a_row_i8(a_blk_len, feature + a_blk_start * a_blk_len, (a_blk_end - a_blk_start) * a_blk_len,
367                                  quant_a_buffer + a_blk_start * block_stride_a);
368            }
369        } else {
370            int task_per_thread = spacemit_kernels::div_round_up(row_blks, nth);
371            int m_row_blk_start = ith * task_per_thread;
372            int m_row_blk_end   = std::min(m_row_blk_start + task_per_thread, (int) row_blks);
373            for (int m_row_blk = m_row_blk_start; m_row_blk < m_row_blk_end; m_row_blk++) {
374                int m_idx             = m_row_blk * row_align;
375                int rows_tobe_handled = (gemm_m - m_idx) > row_align ? row_align : (gemm_m - m_idx);
376
377                if (rows_tobe_handled == row_align && quantize_a_4row_i8 != nullptr) {
378                    const float * a_row_ptr       = feature + m_idx * gemm_k;
379                    auto *        quant_a_row_ptr = quant_a_buffer + m_idx * row_stride_a;
380                    quantize_a_4row_i8(a_blk_len, a_row_ptr, gemm_k, quant_a_row_ptr);
381                } else {
382                    while (rows_tobe_handled) {
383                        const float * a_row_ptr       = feature + m_idx * gemm_k;
384                        auto *        quant_a_row_ptr = quant_a_buffer + m_idx * row_stride_a;
385                        quantize_a_row_i8(a_blk_len, a_row_ptr, gemm_k, quant_a_row_ptr);
386                        rows_tobe_handled -= 1;
387                        m_idx += 1;
388                    }
389                }
390            }
391        }
392
393        ggml_barrier(params->threadpool);
394
395        const int64_t gemm_m_stride     = gemm_n / gemm_m > 64 ? gemm_m : 16;
396        const int64_t gemm_m_blocked    = spacemit_kernels::div_round_up(gemm_m, gemm_m_stride);
397        const int64_t max_gemm_n_stride = spacemit_kernels::div_round_up(gemm_n * gemm_m_blocked, nth);
398
399        int64_t gemm_n_stride = gemm_n;
400        if (max_gemm_n_stride < gemm_n) {
401            gemm_n_stride =
402                std::min(gemm_n_stride, spacemit_kernels::div_round_up(max_gemm_n_stride, NB_COLS) * NB_COLS);
403        }
404
405        if (gemm_n_stride == gemm_n && tcm_buffer != nullptr && per_mb_rows_wsize <= tcm_buffer_size) {
406            for (int64_t m_start = ith * row_align; m_start < gemm_m; m_start += row_align * nth) {
407                uint8_t * b_col    = reinterpret_cast<uint8_t *>(w_data);
408                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
409
410                int64_t m_row_real = std::min(gemm_m - m_start, row_align);
411
412                spacemit_kernels::rvv::memcpy1d(tcm_buffer, quant_a_buffer + m_start * row_stride_a,
413                                                m_row_real * row_stride_a);
414
415                int64_t n_blk_real = 0;
416                for (int64_t ni = 0; ni < gemm_n; ni += n_blk_real, b_col += n_blk_real * row_stride_b) {
417                    n_blk_real = std::min(gemm_n - ni, (int64_t) NB_COLS);
418
419                    uint8_t * a_row_ptr = (uint8_t *) tcm_buffer;
420                    float *   c_blk     = output + m_start * gemm_n + ni;
421
422                    int32_t rows_remaining = m_row_real;
423
424                    while (rows_remaining > 0) {
425                        auto rows_handled = gemm_kernel(b_blk_len, a_row_ptr, b_col, b_col_zp, c_blk, rows_remaining,
426                                                        n_blk_real, b_k_blks, gemm_n);
427
428                        c_blk += rows_handled * gemm_n;
429                        a_row_ptr += rows_handled * row_stride_a;
430
431                        rows_remaining -= rows_handled;
432                    }
433                }
434            }
435        } else if (tcm_buffer != nullptr && per_nb_cols_wsize <= tcm_buffer_size) {
436            uint8_t * a_row = quant_a_buffer;
437            uint8_t * b_col = reinterpret_cast<uint8_t *>(tcm_buffer);
438            if ((gemm_workspace_size + per_nb_cols_wsize) <= tcm_buffer_size) {
439                a_row = (uint8_t *) tcm_buffer;
440                b_col = reinterpret_cast<uint8_t *>(tcm_buffer) + gemm_workspace_size;
441            }
442            uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
443
444            int64_t ni      = ith * NB_COLS;
445            int64_t nb_real = std::min(gemm_n - ni, NB_COLS);
446
447            if (ith % 2 == 0 && nb_real > 0) {
448                spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(w_data) + ni * row_stride_b,
449                                                nb_real * row_stride_b);
450                if (a_row != quant_a_buffer) {
451                    spacemit_kernels::rvv::memcpy1d(a_row, quant_a_buffer, gemm_workspace_size);
452                }
453            }
454
455            spine_barrier_wait(cur_barrier);
456
457            if (ith % 2 != 0 && nb_real > 0) {
458                if (a_row != quant_a_buffer) {
459                    spacemit_kernels::rvv::memcpy1d(a_row, quant_a_buffer, gemm_workspace_size);
460                }
461                spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(w_data) + ni * row_stride_b,
462                                                nb_real * row_stride_b);
463            }
464
465            for (; ni < gemm_n; ni += NB_COLS * nth) {
466                int64_t rows_remaining = gemm_m;
467                float * c_blk          = output + ni;
468                auto *  a_row_cur      = a_row;
469
470                if (ith % 2 != 0) {
471                    spine_barrier_wait(cur_barrier);
472                }
473
474                while (rows_remaining > 0) {
475                    auto rows_handled = gemm_kernel(b_blk_len, a_row_cur, b_col, b_col_zp, c_blk, rows_remaining,
476                                                    nb_real, b_k_blks, gemm_n);
477
478                    c_blk += rows_handled * gemm_n;
479                    a_row_cur += rows_handled * row_stride_a;
480
481                    rows_remaining -= rows_handled;
482                }
483
484                if (ith % 2 == 0) {
485                    spine_barrier_wait(cur_barrier);
486                }
487
488                const int64_t next_ni = ni + NB_COLS * nth;
489                if (next_ni < gemm_n) {
490                    nb_real = std::min(gemm_n - next_ni, NB_COLS);
491                    spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(w_data) + next_ni * row_stride_b,
492                                                    nb_real * row_stride_b);
493                }
494            }
495        } else {
496            const int64_t task_count_m = spacemit_kernels::div_round_up(gemm_m, gemm_m_stride);
497            const int64_t task_count_n = spacemit_kernels::div_round_up(gemm_n, gemm_n_stride);
498
499            int64_t task_count      = task_count_m * task_count_n;
500            int64_t task_per_thread = (task_count + nth - 1) / nth;
501            int64_t start           = ith * task_per_thread;
502            int64_t end             = std::min((ith + 1) * task_per_thread, task_count);
503            for (int64_t compute_idx = start; compute_idx < end; compute_idx++) {
504                const auto tid_n = compute_idx / task_count_m;
505                const auto tid_m = compute_idx % task_count_m;
506
507                const int64_t m_start = tid_m * gemm_m_stride;
508                const int64_t m_count = std::min(gemm_m - m_start, (int64_t) gemm_m_stride);
509
510                const int64_t n_start = tid_n * gemm_n_stride;
511                const int64_t n_count = std::min(gemm_n - n_start, (int64_t) gemm_n_stride);
512
513                const int64_t n_blk = m_count == 1 ? n_count : NB_COLS;
514
515                uint8_t * b_col    = reinterpret_cast<uint8_t *>(w_data) + n_start * row_stride_b;
516                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
517
518                int64_t n_blk_real = 0;
519                for (int64_t ni = 0; ni < n_count; ni += n_blk_real, b_col += n_blk_real * row_stride_b) {
520                    n_blk_real = std::min(n_count - ni, n_blk);
521
522                    uint8_t * a_row = quant_a_buffer + m_start * row_stride_a;
523
524                    float * c_blk = output + m_start * gemm_n + n_start + ni;
525
526                    int64_t rows_remaining = m_count;
527
528                    uint8_t * b_col_cur    = b_col;
529                    uint8_t * b_col_zp_cur = b_col_zp;
530
531                    while (rows_remaining > 0) {
532                        auto rows_handled = gemm_kernel(b_blk_len, a_row, b_col_cur, b_col_zp_cur, c_blk,
533                                                        rows_remaining, n_blk_real, b_k_blks, gemm_n);
534
535                        c_blk += rows_handled * gemm_n;
536                        a_row += rows_handled * row_stride_a;
537
538                        rows_remaining -= rows_handled;
539                    }
540                }
541            }
542        }
543    }
544
545    void forward_mul_mat_id(ggml_compute_params * params, ggml_tensor * op) {
546        constexpr size_t a_blk_len = INTER_SIZE;
547        constexpr size_t b_blk_len = INTER_SIZE;
548
549        const ggml_tensor * src0 = op->src[0];
550        const ggml_tensor * src1 = op->src[1];
551        const ggml_tensor * ids  = op->src[2];
552        ggml_tensor *       dst  = op;
553
554        GGML_TENSOR_BINARY_OP_LOCALS
555
556        int ith = params->ith;
557        int nth = params->nth;
558
559        // row groups
560        const int n_ids = ids->ne[0];  // n_expert_used
561        const int n_as  = ne02;        // n_expert
562
563        struct mmid_row_mapping {
564            int32_t i1;
565            int32_t i2;
566        };
567
568        spacemit_kernels::quantize_a_row_def           quantize_a_row_i8;
569        spacemit_kernels::gemm_kernel_quantize_def     gemm_kernel;
570        spacemit_kernels::moe_gemm_kernel_quantize_def moe_gemm_kernel_m2;
571        bool                                           set_kernel_impl = false;
572        size_t                                         block_stride_a  = spacemit_kernels::q8_blk_size(QK4_0);
573
574#if defined(RISCV64_SPACEMIT_IME2)
575        if (!set_kernel_impl && (global_spine_env_info.use_ime2)) {
576            quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8;
577            block_stride_a    = spacemit_kernels::q8_blk_size(QK4_0, true);
578
579            if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0>) {
580                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i8;
581                set_kernel_impl = true;
582            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
583                                 std::is_same_v<BLOC_TYPE, block_q4_K>) {
584                if constexpr (INTER_SIZE == 256) {
585                    gemm_kernel       = spacemit_kernels::ime2::gemm_kernel_i8i4_hp;
586                    quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8_hp;
587                    block_stride_a    = spacemit_kernels::q8_hp_blk_size(a_blk_len, true, true);
588                    set_kernel_impl   = true;
589                } else {
590                    gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i4;
591                    moe_gemm_kernel_m2 = spacemit_kernels::ime2::moe_m2_gemm_kernel_i8i4;
592                    quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8;
593                    block_stride_a     = spacemit_kernels::q8_blk_size(a_blk_len, true);
594                    set_kernel_impl    = true;
595                }
596            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K>) {
597                quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8k;
598                block_stride_a    = spacemit_kernels::q8k_blk_size(a_blk_len);
599                gemm_kernel       = spacemit_kernels::ime2::gemm_kernel_i8i2k;
600                set_kernel_impl   = true;
601            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q3_K>) {
602                quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8k;
603                block_stride_a    = spacemit_kernels::q8k_blk_size(a_blk_len);
604                gemm_kernel       = spacemit_kernels::ime2::gemm_kernel_i8i3k;
605                set_kernel_impl   = true;
606            } else if constexpr (std::is_same_v<BLOC_TYPE, block_mxfp4>) {
607                gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8mxfp4;
608                moe_gemm_kernel_m2 = spacemit_kernels::ime2::moe_m2_gemm_kernel_i8mxfp4;
609                set_kernel_impl    = true;
610            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_1> || std::is_same_v<BLOC_TYPE, block_q5_K> ||
611                                 std::is_same_v<BLOC_TYPE, block_q5_0>) {
612                gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i5;
613                moe_gemm_kernel_m2 = spacemit_kernels::ime2::moe_m2_gemm_kernel_i8i5;
614                set_kernel_impl    = true;
615            }
616        }
617#endif
618
619#if defined(RISCV64_SPACEMIT_IME1)
620        if (!set_kernel_impl && (global_spine_env_info.use_ime1)) {
621            quantize_a_row_i8 = spacemit_kernels::ime1::quantize_a_row_i8;
622
623            if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
624                          std::is_same_v<BLOC_TYPE, block_q4_K>) {
625                gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i4;
626                set_kernel_impl = true;
627            }
628        }
629#endif
630        if (!set_kernel_impl) {
631            GGML_ABORT("no kernel implementation found for the block type");
632        }
633
634        const size_t a_k_blks = spacemit_kernels::div_round_up(ne10, a_blk_len);
635        const size_t b_k_blks = spacemit_kernels::div_round_up(ne10, b_blk_len);
636
637        const size_t nbw1                = a_k_blks * block_stride_a;
638        const size_t nbw2                = ne11 * nbw1;
639        const size_t nbw3                = nbw2 * ne12;
640        const size_t gemm_workspace_size = GGML_PAD(nbw3, alignof(int64_t));
641
642        const uintptr_t ws_ptr         = reinterpret_cast<uintptr_t>(params->wdata);
643        auto *          quant_a_buffer = reinterpret_cast<uint8_t *>(ws_ptr);
644
645        if (ne11 == 1) {
646            for (int64_t ii = ith; ii < ne12 * a_k_blks; ii += nth) {
647                int64_t i12       = ii / a_k_blks;
648                int64_t ak_blk_id = ii % a_k_blks;
649                quantize_a_row_i8(a_blk_len, (float *) ((char *) src1->data + i12 * nb12) + ak_blk_id * a_blk_len,
650                                  a_blk_len, quant_a_buffer + i12 * nbw2 + ak_blk_id * block_stride_a);
651            }
652        } else {
653            for (int64_t ii = ith; ii < ne12 * ne11; ii += nth) {
654                int64_t i12 = ii / ne11;
655                int64_t i11 = ii % ne11;
656                quantize_a_row_i8(a_blk_len, (float *) ((char *) src1->data + i12 * nb12 + i11 * nb11), ne10,
657                                  quant_a_buffer + i12 * nbw2 + i11 * nbw1);
658            }
659        }
660
661#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) *ne12 + (i1)]
662
663        int64_t *          matrix_row_counts       = (int64_t *) (ws_ptr + gemm_workspace_size);
664        int32_t *          valid_ep_count          = (int32_t *) (matrix_row_counts + n_as);
665        int32_t *          valid_act_count         = (int32_t *) (valid_ep_count + 1);
666        int64_t *          valid_matrix_row_counts = (int64_t *) (valid_act_count + 1);
667        mmid_row_mapping * matrix_rows             = (mmid_row_mapping *) (valid_matrix_row_counts + n_as);
668
669        if (ith == 0) {
670            // initialize matrix_row_counts
671            memset(matrix_row_counts, 0, n_as * sizeof(int64_t));
672
673            // group rows by src0 matrix
674            for (int32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
675                for (int32_t id = 0; id < n_ids; ++id) {
676                    const int32_t i02 =
677                        *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
678
679                    GGML_ASSERT(i02 >= 0 && i02 < n_as);
680
681                    MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
682                    matrix_row_counts[i02] += 1;
683                }
684            }
685
686            int32_t valid_ep_count_t  = 0;
687            int32_t valid_act_count_t = 0;
688            for (int cur_a = 0; cur_a < n_as; ++cur_a) {
689                const int64_t cne1 = matrix_row_counts[cur_a];
690                if (cne1 == 0) {
691                    continue;
692                }
693                valid_matrix_row_counts[valid_ep_count_t] = cur_a;
694                valid_act_count_t += cne1;
695                valid_ep_count_t += 1;
696            }
697            valid_ep_count[0]  = valid_ep_count_t;
698            valid_act_count[0] = valid_act_count_t;
699        }
700
701        const int64_t barrier_idx = static_cast<int64_t>(ith / 2);
702
703        GGML_ASSERT(global_spine_env_info.init_barrier != nullptr);
704        GGML_ASSERT(barrier_idx < spine_init_barrier_count);
705        spine_barrier_t * cur_barrier = &global_spine_env_info.init_barrier[barrier_idx];
706
707        ggml_barrier(params->threadpool);
708
709        const size_t row_stride_b      = b_k_blks * get_repacked_block_type_size<BLOC_TYPE, INTER_SIZE, NB_COLS>();
710        const size_t expert_b_stride   = ne01 * row_stride_b;
711        const size_t per_nb_cols_wsize = NB_COLS * row_stride_b;
712
713        std::array<const uint8_t *, 2> src_workspaces;
714        std::array<float *, 2>         dst_workspaces;
715
716        auto *     tcm_buffer      = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer;
717        const auto tcm_buffer_size = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size;
718
719        const auto valid_ep_count_t  = valid_ep_count[0];
720        const auto valid_act_count_t = valid_act_count[0];
721
722        int nth_es = 1;
723        int nth_n  = nth;
724
725        int ith_es = ith % nth_es;
726        int ith_n  = (ith / nth_es) % nth_n;
727
728        if (valid_ep_count_t % nth == 0 && tcm_buffer != nullptr && valid_ep_count_t == n_as &&
729            valid_act_count_t == n_as && per_nb_cols_wsize <= tcm_buffer_size) {
730            for (int64_t valid_id = ith; valid_id < valid_ep_count_t; valid_id += nth) {
731                const int64_t cur_a = valid_matrix_row_counts[valid_id];
732
733                auto * src0_cur = (uint8_t *) src0->data + cur_a * expert_b_stride;
734
735                mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, 0);
736                const int        id          = row_mapping.i1;
737                const int64_t    i11         = id % ne11;
738                const int64_t    i12         = row_mapping.i2;
739                const int64_t    i1          = id;
740                const int64_t    i2          = i12;
741
742                auto *  src1_col = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
743                float * c_blk    = (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2));
744
745                uint8_t * a_row = src1_col;
746                uint8_t * b_col = reinterpret_cast<uint8_t *>(tcm_buffer);
747                if ((nbw1 + per_nb_cols_wsize) <= tcm_buffer_size) {
748                    a_row = (uint8_t *) tcm_buffer;
749                    b_col = reinterpret_cast<uint8_t *>(tcm_buffer) + nbw1;
750                }
751                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
752
753                if (ith % 2 == 0) {
754                    spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(src0_cur), per_nb_cols_wsize);
755
756                    if (a_row != src1_col) {
757                        spacemit_kernels::rvv::memcpy1d(a_row, src1_col, nbw1);
758                    }
759                }
760
761                spine_barrier_wait(cur_barrier);
762
763                if (ith % 2 != 0) {
764                    if (a_row != src1_col) {
765                        spacemit_kernels::rvv::memcpy1d(a_row, src1_col, nbw1);
766                    }
767
768                    spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(src0_cur), per_nb_cols_wsize);
769                }
770
771                int64_t nb_real = std::min(ne01, NB_COLS);
772                for (int64_t ni = 0; ni < ne01; ni += NB_COLS) {
773                    if (ith % 2 != 0) {
774                        spine_barrier_wait(cur_barrier);
775                    }
776
777                    gemm_kernel(b_blk_len, a_row, b_col, b_col_zp, c_blk + ni, 1, nb_real, b_k_blks, ne01);
778
779                    if (ith % 2 == 0) {
780                        spine_barrier_wait(cur_barrier);
781                    }
782
783                    const int64_t next_ni = ni + NB_COLS;
784                    if (next_ni < ne01) {
785                        nb_real = std::min(ne01 - next_ni, NB_COLS);
786                        spacemit_kernels::rvv::memcpy1d(
787                            b_col, reinterpret_cast<uint8_t *>(src0_cur) + next_ni * row_stride_b, per_nb_cols_wsize);
788                    }
789                }
790            }
791        } else {
792            for (int64_t valid_id = ith_es; valid_id < valid_ep_count_t; valid_id += nth_es) {
793                const int64_t cur_a = valid_matrix_row_counts[valid_id];
794                const int64_t cne1  = matrix_row_counts[cur_a];
795
796                int64_t src1_cur_start = 0;
797                int64_t src1_cur_end   = cne1;
798
799                int64_t src0_cur_start = (ith_n * ne01) / nth_n;
800                int64_t src0_cur_end   = MIN(((ith_n + 1) * ne01) / nth_n, ne01);
801
802                if (src1_cur_start >= src1_cur_end || src0_cur_start >= src0_cur_end) {
803                    continue;
804                }
805
806                src0_cur_start =
807                    (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start;
808                src0_cur_end =
809                    (src0_cur_end % NB_COLS) ? src0_cur_end + NB_COLS - (src0_cur_end % NB_COLS) : src0_cur_end;
810
811                auto *    src0_cur = (uint8_t *) src0->data + cur_a * expert_b_stride + src0_cur_start * row_stride_b;
812                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? src0_cur : nullptr;
813
814                size_t extra_tcm_buffer_size = tcm_buffer_size;
815                void * extra_tcm_buffer      = tcm_buffer;
816                if (tcm_buffer != nullptr && (src1_cur_end - src1_cur_start) >= 4 &&
817                    (src0_cur_end - src0_cur_start) * row_stride_b <= tcm_buffer_size) {
818                    spacemit_kernels::rvv::memcpy1d(tcm_buffer, src0_cur,
819                                                    (src0_cur_end - src0_cur_start) * row_stride_b);
820                    src0_cur = reinterpret_cast<uint8_t *>(tcm_buffer);
821                    b_col_zp = block_type_has_zp<BLOC_TYPE>() ? src0_cur : nullptr;
822                    extra_tcm_buffer_size -= (src0_cur_end - src0_cur_start) * row_stride_b;
823                    extra_tcm_buffer = reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(tcm_buffer) +
824                                                                (src0_cur_end - src0_cur_start) * row_stride_b);
825                }
826
827                int ir1 = src1_cur_start;
828
829                if (extra_tcm_buffer_size >= nbw1 && extra_tcm_buffer != nullptr) {
830                    int64_t quant_a_tile_size = extra_tcm_buffer_size / nbw1;
831                    do {
832                        quant_a_tile_size = MIN(quant_a_tile_size, src1_cur_end - ir1);
833
834                        uint8_t * quant_a_tile_buffer = reinterpret_cast<uint8_t *>(extra_tcm_buffer);
835
836                        int iir1 = ir1;
837                        for (; iir1 < (ir1 + quant_a_tile_size); ++iir1) {
838                            mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, iir1);
839
840                            const int id = row_mapping.i1;  // selected expert index
841
842                            const int64_t i11 = id % ne11;
843                            const int64_t i12 = row_mapping.i2;  // row index in src1
844
845                            auto * src1_col = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
846                            spacemit_kernels::rvv::memcpy1d(quant_a_tile_buffer, src1_col, nbw1);
847                            quant_a_tile_buffer = quant_a_tile_buffer + nbw1;
848                        }
849
850                        quant_a_tile_buffer = reinterpret_cast<uint8_t *>(extra_tcm_buffer);
851                        iir1                = ir1;
852
853                        if (moe_gemm_kernel_m2 != nullptr) {
854                            for (; iir1 < (ir1 + quant_a_tile_size - 1); iir1 += 2, quant_a_tile_buffer += 2 * nbw1) {
855                                mmid_row_mapping row_mapping_0 = MMID_MATRIX_ROW(cur_a, iir1);
856                                mmid_row_mapping row_mapping_1 = MMID_MATRIX_ROW(cur_a, iir1 + 1);
857
858                                src_workspaces[0] = quant_a_tile_buffer;
859                                src_workspaces[1] = quant_a_tile_buffer + nbw1;
860
861                                dst_workspaces[0] =
862                                    (float *) ((char *) dst->data + (row_mapping_0.i1 * nb1 + row_mapping_0.i2 * nb2)) +
863                                    src0_cur_start;
864                                dst_workspaces[1] = (float *) ((char *) dst->data +
865                                                               ((row_mapping_1.i1) * nb1 + (row_mapping_1.i2) * nb2)) +
866                                                    src0_cur_start;
867                                moe_gemm_kernel_m2(b_blk_len, src_workspaces.data(), src0_cur, b_col_zp,
868                                                   dst_workspaces.data(), 1, src0_cur_end - src0_cur_start, b_k_blks,
869                                                   ne01);
870                            }
871                        }
872
873                        for (; iir1 < (ir1 + quant_a_tile_size); iir1++, quant_a_tile_buffer += nbw1) {
874                            mmid_row_mapping row_mapping_0 = MMID_MATRIX_ROW(cur_a, iir1);
875
876                            gemm_kernel(
877                                b_blk_len, quant_a_tile_buffer, src0_cur, b_col_zp,
878                                (float *) ((char *) dst->data + (row_mapping_0.i1 * nb1 + row_mapping_0.i2 * nb2)) +
879                                    src0_cur_start,
880                                1, src0_cur_end - src0_cur_start, b_k_blks, ne01);
881                        }
882
883                        ir1 += quant_a_tile_size;
884                    } while (ir1 < src1_cur_end);
885                } else {
886                    if (moe_gemm_kernel_m2 != nullptr) {
887                        for (; ir1 < src1_cur_end - 1; ir1 += 2) {
888                            for (int iir1 = 0; iir1 < 2; ++iir1) {
889                                mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1 + iir1);
890
891                                const int id = row_mapping.i1;  // selected expert index
892
893                                const int64_t i11 = id % ne11;
894                                const int64_t i12 = row_mapping.i2;  // row index in src1
895
896                                const int64_t i1 = id;               // selected expert index
897                                const int64_t i2 = i12;              // row
898
899                                src_workspaces[iir1] = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
900
901                                dst_workspaces[iir1] =
902                                    (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start;
903                            }
904
905                            moe_gemm_kernel_m2(b_blk_len, src_workspaces.data(), src0_cur, b_col_zp,
906                                               dst_workspaces.data(), 1, src0_cur_end - src0_cur_start, b_k_blks, ne01);
907                        }
908                    }
909
910                    for (; ir1 < src1_cur_end; ir1++) {
911                        mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1);
912
913                        const int id = row_mapping.i1;  // selected expert index
914
915                        const int64_t i11 = id % ne11;
916                        const int64_t i12 = row_mapping.i2;  // row index in src1
917
918                        const int64_t i1 = id;               // selected expert index
919                        const int64_t i2 = i12;              // row
920
921                        auto * src1_col = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
922
923                        gemm_kernel(b_blk_len, src1_col, src0_cur, b_col_zp,
924                                    (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start, 1,
925                                    src0_cur_end - src0_cur_start, b_k_blks, ne01);
926                    }
927                }
928            }
929        }
930#undef MMID_MATRIX_ROW
931    }
932
933    int repack(ggml_tensor * t, const void * data, size_t data_size) override {
934        GGML_LOG_DEBUG("%s: repack tensor %s with %s_%dx%d\n", __func__, t->name, ggml_type_name(t->type),
935                       (int) NB_COLS, (int) INTER_SIZE);
936        return ggml::cpu::riscv64_spacemit::repack<BLOC_TYPE, INTER_SIZE, NB_COLS>(t, data, data_size);
937    }
938};
939
940class tensor_traits_common : public tensor_traits_base {
941    bool work_size(int n_threads, const ggml_tensor * op, size_t & size) override {
942        switch (op->op) {
943            case GGML_OP_FLASH_ATTN_EXT:
944                {
945                    const int     n_tasks = n_threads;
946                    const int64_t neq2    = op->src[0]->ne[2];  // number of query heads
947                    const int64_t DK      = op->src[1]->ne[0];
948                    const int64_t DV      = op->src[2]->ne[0];  // DV
949
950                    // Tiled flash attention scratch (tile sizes defined in common.h)
951                    // Per-thread: Q_q + KQ + mask + VKQ32 + V32 + K_f32 + padding
952                    size_t prefill = sizeof(float) *
953                                     (GGML_FA_TILE_Q * DK + 2 * GGML_FA_TILE_Q * GGML_FA_TILE_KV + GGML_FA_TILE_Q * DV +
954                                      GGML_FA_TILE_KV * DV + GGML_FA_TILE_KV * DK) *
955                                     n_tasks;
956
957                    // Decode path: n_kv_chunks = n_tasks (one chunk per thread)
958                    // Per-thread: VKQ accmulator (DV), partial M, partial S + intra-thread scratch for V, Q and VKQ
959                    size_t n_chunks = n_tasks;
960                    size_t decode   = sizeof(float) * (neq2 * n_chunks * (2 + DV) + n_tasks * (DK + 2 * DV));
961
962                    size = MAX(prefill, decode);
963                }
964                return true;
965            default:
966                break;
967        }
968        return false;
969    }
970
971    bool compute_forward(ggml_compute_params * params, ggml_tensor * op) override {
972        switch (op->op) {
973            case GGML_OP_NORM:
974                switch (op->src[0]->type) {
975                    case GGML_TYPE_F32:
976                        spacemit_kernels::rvv::forward_norm_f32(params, op);
977                        return true;
978                    default:
979                        GGML_ABORT("fatal error");
980                }
981            case GGML_OP_RMS_NORM:
982                switch (op->src[0]->type) {
983                    case GGML_TYPE_F32:
984                        spacemit_kernels::rvv::forward_rms_norm_f32(params, op);
985                        return true;
986                    default:
987                        GGML_ABORT("fatal error");
988                }
989            case GGML_OP_ADD:
990                switch (op->src[0]->type) {
991                    case GGML_TYPE_F32:
992                        spacemit_kernels::rvv::forward_binary<GGML_OP_ADD, float>(params, op);
993                        return true;
994                    case GGML_TYPE_F16:
995                        spacemit_kernels::rvv::forward_binary<GGML_OP_ADD, _Float16>(params, op);
996                        return true;
997                    default:
998                        ggml_compute_forward_add(params, op);
999                        return true;
1000                }
1001            case GGML_OP_SUB:
1002                switch (op->src[0]->type) {
1003                    case GGML_TYPE_F32:
1004                        spacemit_kernels::rvv::forward_binary<GGML_OP_SUB, float>(params, op);
1005                        return true;
1006                    case GGML_TYPE_F16:
1007                        spacemit_kernels::rvv::forward_binary<GGML_OP_SUB, _Float16>(params, op);
1008                        return true;
1009                    default:
1010                        ggml_compute_forward_sub(params, op);
1011                        return true;
1012                }
1013            case GGML_OP_MUL:
1014                switch (op->src[0]->type) {
1015                    case GGML_TYPE_F32:
1016                        spacemit_kernels::rvv::forward_binary<GGML_OP_MUL, float>(params, op);
1017                        return true;
1018                    case GGML_TYPE_F16:
1019                        spacemit_kernels::rvv::forward_binary<GGML_OP_MUL, _Float16>(params, op);
1020                        return true;
1021                    default:
1022                        ggml_compute_forward_mul(params, op);
1023                        return true;
1024                }
1025            case GGML_OP_DIV:
1026                switch (op->src[0]->type) {
1027                    case GGML_TYPE_F32:
1028                        spacemit_kernels::rvv::forward_binary<GGML_OP_DIV, float>(params, op);
1029                        return true;
1030                    case GGML_TYPE_F16:
1031                        spacemit_kernels::rvv::forward_binary<GGML_OP_DIV, _Float16>(params, op);
1032                        return true;
1033                    default:
1034                        ggml_compute_forward_div(params, op);
1035                        return true;
1036                }
1037            case GGML_OP_FLASH_ATTN_EXT:
1038                forward_flash_attn_ext_f16(params, op);
1039                return true;
1040            case GGML_OP_CONT:
1041                {
1042                    const ggml_tensor * src0 = op->src[0];
1043                    if (op->type == src0->type && op->nb[0] != src0->nb[0] && op->nb[0] == src0->nb[1] &&
1044                        op->ne[3] * op->ne[2] * op->nb[2] == src0->ne[3] * src0->ne[2] * src0->nb[2]) {
1045                        spacemit_kernels::rvv::forward_cont_with_permute(params, op);
1046                    } else {
1047                        ggml_compute_forward_cont(params, op);
1048                    }
1049                    return true;
1050                }
1051            case GGML_OP_CPY:
1052                {
1053                    const ggml_tensor * src0 = op->src[0];
1054                    if (op->type == src0->type && op->nb[0] == src0->nb[1] && src0->nb[0] != src0->nb[1] &&
1055                        ggml_nelements(src0) == ggml_nelements(op)) {
1056                        spacemit_kernels::rvv::forward_cpy_with_permute(params, op);
1057                    } else {
1058                        ggml_compute_forward_cpy(params, op);
1059                    }
1060                    return true;
1061                }
1062            case GGML_OP_REPEAT:
1063                {
1064                    const bool rows_equal         = ggml_nrows(op->src[0]) == ggml_nrows(op);
1065                    const bool broadcast_or_equal = op->src[0]->ne[0] == 1 || op->src[0]->ne[0] == op->ne[0];
1066
1067                    if (rows_equal && broadcast_or_equal) {
1068                        switch (op->src[0]->type) {
1069                            case GGML_TYPE_F32:
1070                                spacemit_kernels::rvv::forward_repeat_nrows<int32_t>(params, op);
1071                                return true;
1072                            case GGML_TYPE_F16:
1073                                spacemit_kernels::rvv::forward_repeat_nrows<int16_t>(params, op);
1074                                return true;
1075                            default:
1076                                break;
1077                        }
1078                    }
1079
1080                    if (op->src[0]->ne[1] == 1 && op->src[0]->ne[0] == op->ne[0]) {
1081                        switch (op->src[0]->type) {
1082                            case GGML_TYPE_F32:
1083                                spacemit_kernels::rvv::forward_repeat_dim1<int32_t>(params, op);
1084                                return true;
1085                            case GGML_TYPE_F16:
1086                                spacemit_kernels::rvv::forward_repeat_dim1<int16_t>(params, op);
1087                                return true;
1088                            default:
1089                                break;
1090                        }
1091                    }
1092
1093                    ggml_compute_forward_repeat(params, op);
1094                }
1095                return true;
1096            case GGML_OP_SUM_ROWS:
1097                {
1098                    if (op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) {
1099                        spacemit_kernels::rvv::forward_sum_rows<float>(params, op);
1100                    } else {
1101                        ggml_compute_forward_sum_rows(params, op);
1102                    }
1103                }
1104                return true;
1105            case GGML_OP_GET_ROWS:
1106                {
1107                    if (op->src[0]->type == op->type) {
1108                        switch (op->src[0]->type) {
1109                            case GGML_TYPE_F32:
1110                                spacemit_kernels::rvv::forward_get_rows<int32_t>(params, op);
1111                                return true;
1112                            case GGML_TYPE_F16:
1113                                spacemit_kernels::rvv::forward_get_rows<int16_t>(params, op);
1114                                return true;
1115                            default:
1116                                break;
1117                        }
1118                    }
1119
1120                    ggml_compute_forward_get_rows(params, op);
1121                }
1122                return true;
1123            case GGML_OP_CONCAT:
1124                {
1125                    const int32_t dim = ggml_get_op_params_i32(op, 0);
1126                    if (dim == 0 && op->type == op->src[0]->type) {
1127                        switch (op->src[0]->type) {
1128                            case GGML_TYPE_F32:
1129                                spacemit_kernels::rvv::forward_concat<int32_t>(params, op);
1130                                return true;
1131                            case GGML_TYPE_F16:
1132                                spacemit_kernels::rvv::forward_concat<int16_t>(params, op);
1133                                return true;
1134                            default:
1135                                break;
1136                        }
1137                    }
1138
1139                    ggml_compute_forward_concat(params, op);
1140                }
1141                return true;
1142            // TODO For GGML_OP_GATED_DELTA_NET
1143            // case GGML_OP_GATED_DELTA_NET:
1144            //     return true;
1145            default:
1146                break;
1147        }
1148        return false;
1149    }
1150
1151    void forward_flash_attn_ext_f16(const ggml_compute_params * params, ggml_tensor * dst) {
1152        const ggml_tensor * q = dst->src[0];
1153        const ggml_tensor * k = dst->src[1];
1154        const ggml_tensor * v = dst->src[2];
1155
1156        GGML_TENSOR_LOCALS(int64_t, neq, q, ne)
1157        GGML_TENSOR_LOCALS(size_t, nbq, q, nb)
1158        GGML_TENSOR_LOCALS(int64_t, nek, k, ne)
1159        GGML_TENSOR_LOCALS(size_t, nbk, k, nb)
1160        GGML_TENSOR_LOCALS(int64_t, nev, v, ne)
1161        GGML_TENSOR_LOCALS(size_t, nbv, v, nb)
1162        GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
1163        GGML_TENSOR_LOCALS(size_t, nb, dst, nb)
1164
1165        const int64_t DK = nek0;
1166        const int64_t DV = nev0;
1167
1168        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_DEFAULT);
1169        const bool supported_types = (q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
1170        const bool supported_shape = (DK > 0 && DK <= 128 && DV > 0 && DV <= 128);
1171        const bool supported_vlen  = (__riscv_vlenb() == 128);
1172
1173        if (!(supported_prec && supported_types && supported_shape && supported_vlen)) {
1174            ggml_compute_forward_flash_attn_ext(params, dst);
1175            return;
1176        }
1177
1178        // total rows in q
1179        const int64_t nr = neq1 * neq2 * neq3;
1180
1181        // rows per thread
1182        const int ith = params->ith;
1183        const int nth = params->nth;
1184
1185        static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q;
1186        const bool               use_tiled = !params->use_ref && (neq1 >= Q_TILE_SZ);
1187
1188        // 4x chunks per thread
1189        // int     nth_scaled = nth * 4;
1190        // int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled;
1191        // int64_t nchunk     = (nr + chunk_size - 1) / chunk_size;
1192
1193        // if (nth == 1 || nchunk < nth) {
1194        //     nchunk = nth;
1195        // }
1196
1197        int64_t nchunk = nth;
1198
1199        if (ith == 0) {
1200            // Every thread starts at ith, so the first unprocessed chunk is nth.  This save a bit of coordination right at the start.
1201            ggml_threadpool_chunk_set(params->threadpool, nth);
1202        }
1203
1204        ggml_barrier(params->threadpool);
1205
1206        // The number of elements in each chunk
1207        const int64_t dr = (nr + nchunk - 1) / nchunk;
1208
1209        // The first chunk comes from our thread_id, the rest will get auto-assigned.
1210        int current_chunk = ith;
1211
1212        while (current_chunk < nchunk) {
1213            const int64_t ir0 = dr * current_chunk;
1214            const int64_t ir1 = MIN(ir0 + dr, nr);
1215
1216            if (use_tiled) {
1217                spacemit_kernels::rvv::forward_flash_attn_ext_f16_tiled_vlen1024_vf16(
1218                    params, dst, ir0, ir1, ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer,
1219                    ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size);
1220            } else {
1221                spacemit_kernels::rvv::forward_flash_attn_ext_f16_one_chunk_vlen1024_vf16(
1222                    params, dst, ir0, ir1, ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer,
1223                    ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size);
1224            }
1225
1226            current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);
1227        }
1228    }
1229
1230    int repack(ggml_tensor * t, const void * data, size_t data_size) override {
1231        memcpy(t->data, data, data_size);
1232        return 0;
1233    }
1234};
1235
1236// Impl By IME1
1237static const tensor_traits<block_q4_0, 32, 16>  q4_0_16x32_q8_0;
1238static const tensor_traits<block_q4_1, 32, 16>  q4_1_16x32_q8_0;
1239static const tensor_traits<block_q4_K, 32, 16>  q4_k_16x32_q8_0;
1240// Impl By IME2
1241static const tensor_traits<block_q2_K, 256, 32> q2_k_32x256_q8_0;
1242static const tensor_traits<block_q3_K, 256, 32> q3_k_32x256_q8_0;
1243static const tensor_traits<block_q4_0, 32, 32>  q4_0_32x32_q8_0;
1244static const tensor_traits<block_q4_1, 32, 32>  q4_1_32x32_q8_0;
1245static const tensor_traits<block_q4_0, 256, 32> q4_0_32x256_q8_0;
1246static const tensor_traits<block_q4_1, 256, 32> q4_1_32x256_q8_0;
1247static const tensor_traits<block_q4_K, 32, 32>  q4_k_32x32_q8_0;
1248static const tensor_traits<block_q6_K, 32, 32>  q6_k_32x32_q8_0;
1249static const tensor_traits<block_q8_0, 32, 32>  q8_0_32x32_q8_0;
1250static const tensor_traits<block_mxfp4, 32, 32> mxfp4_32x32_q8_0;
1251static const tensor_traits<block_q5_K, 32, 32>  q5_k_32x32_q8_0;
1252static const tensor_traits<block_q5_1, 32, 32>  q5_1_32x32_q8_0;
1253static const tensor_traits<block_q5_0, 32, 32>  q5_0_32x32_q8_0;
1254// Impl By RVV
1255static const tensor_traits_common               rvv_impl;
1256
1257}  // namespace ggml::cpu::riscv64_spacemit
1258
1259static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack_type(const ggml_tensor * cur) {
1260    switch (cur->type) {
1261        case GGML_TYPE_Q2_K:
1262            {
1263#if defined(RISCV64_SPACEMIT_IME2)
1264                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1265                    return &ggml::cpu::riscv64_spacemit::q2_k_32x256_q8_0;
1266                }
1267#endif
1268            }
1269            break;
1270        case GGML_TYPE_Q3_K:
1271            {
1272#if defined(RISCV64_SPACEMIT_IME2)
1273                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1274                    return &ggml::cpu::riscv64_spacemit::q3_k_32x256_q8_0;
1275                }
1276#endif
1277            }
1278            break;
1279        case GGML_TYPE_Q4_0:
1280            {
1281#if defined(RISCV64_SPACEMIT_IME2)
1282                if (cur->ne[1] % 32 == 0 && cur->ne[0] % 256 == 0 &&
1283                    (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1284                    return &ggml::cpu::riscv64_spacemit::q4_0_32x256_q8_0;
1285                }
1286
1287                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1288                    return &ggml::cpu::riscv64_spacemit::q4_0_32x32_q8_0;
1289                }
1290#endif
1291
1292#if defined(RISCV64_SPACEMIT_IME1)
1293                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
1294                    return &ggml::cpu::riscv64_spacemit::q4_0_16x32_q8_0;
1295                }
1296#endif
1297            }
1298            break;
1299        case GGML_TYPE_Q4_1:
1300            {
1301#if defined(RISCV64_SPACEMIT_IME2)
1302                // TODO
1303                // if (cur->ne[1] % 32 == 0 && cur->ne[0] % 256 == 0 &&
1304                //     (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1305                //     return &ggml::cpu::riscv64_spacemit::q4_1_32x256_q8_0;
1306                // }
1307
1308                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1309                    return &ggml::cpu::riscv64_spacemit::q4_1_32x32_q8_0;
1310                }
1311#endif
1312
1313#if defined(RISCV64_SPACEMIT_IME1)
1314                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
1315                    return &ggml::cpu::riscv64_spacemit::q4_1_16x32_q8_0;
1316                }
1317#endif
1318            }
1319            break;
1320        case GGML_TYPE_Q4_K:
1321            {
1322#if defined(RISCV64_SPACEMIT_IME2)
1323                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1324                    return &ggml::cpu::riscv64_spacemit::q4_k_32x32_q8_0;
1325                }
1326#endif
1327
1328#if defined(RISCV64_SPACEMIT_IME1)
1329                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
1330                    return &ggml::cpu::riscv64_spacemit::q4_k_16x32_q8_0;
1331                }
1332#endif
1333            }
1334            break;
1335        case GGML_TYPE_Q6_K:
1336            {
1337#if defined(RISCV64_SPACEMIT_IME2)
1338                if ((ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1339                    return &ggml::cpu::riscv64_spacemit::q6_k_32x32_q8_0;
1340                }
1341#endif
1342            }
1343            break;
1344        case GGML_TYPE_Q8_0:
1345            {
1346#if defined(RISCV64_SPACEMIT_IME2)
1347                if ((ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1348                    return &ggml::cpu::riscv64_spacemit::q8_0_32x32_q8_0;
1349                }
1350#endif
1351            }
1352            break;
1353        case GGML_TYPE_MXFP4:
1354            {
1355#if defined(RISCV64_SPACEMIT_IME2)
1356                // TODO
1357                // if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1358                //     return &ggml::cpu::riscv64_spacemit::mxfp4_32x32_q8_0;
1359                // }
1360#endif
1361            }
1362            break;
1363        case GGML_TYPE_Q5_K:
1364            {
1365#if defined(RISCV64_SPACEMIT_IME2)
1366                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1367                    return &ggml::cpu::riscv64_spacemit::q5_k_32x32_q8_0;
1368                }
1369#endif
1370            }
1371            break;
1372        case GGML_TYPE_Q5_1:
1373            {
1374#if defined(RISCV64_SPACEMIT_IME2)
1375                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1376                    return &ggml::cpu::riscv64_spacemit::q5_1_32x32_q8_0;
1377                }
1378#endif
1379            }
1380            break;
1381        case GGML_TYPE_Q5_0:
1382            {
1383#if defined(RISCV64_SPACEMIT_IME2)
1384                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
1385                    return &ggml::cpu::riscv64_spacemit::q5_0_32x32_q8_0;
1386                }
1387#endif
1388            }
1389            break;
1390        default:
1391            break;
1392    }
1393
1394    return nullptr;
1395}
1396
1397static enum ggml_status ggml_backend_riscv64_spacemit_buffer_init_tensor(ggml_backend_buffer_t buffer,
1398                                                                         ggml_tensor *         tensor) {
1399    tensor->extra =
1400        (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_riscv64_spacemit_get_optimal_repack_type(tensor));
1401
1402    GGML_UNUSED(buffer);
1403
1404    return GGML_STATUS_SUCCESS;
1405}
1406
1407static void ggml_backend_riscv64_spacemit_buffer_free_buffer(ggml_backend_buffer_t buffer) {
1408    GGML_ASSERT(buffer);
1409
1410    void * base = buffer->context;
1411    if (base == nullptr) {
1412        return;
1413    }
1414
1415    ggml::cpu::riscv64_spacemit::spine_mem_pool_free(base);
1416}
1417
1418static void * ggml_backend_riscv64_spacemit_buffer_get_base(ggml_backend_buffer_t buffer) {
1419    GGML_ASSERT(buffer);
1420
1421    void * base = buffer->context;
1422    GGML_ASSERT(base != nullptr);
1423    return base;
1424}
1425
1426static void ggml_backend_riscv64_spacemit_buffer_memset_tensor(ggml_backend_buffer_t buffer,
1427                                                               ggml_tensor *         tensor,
1428                                                               uint8_t               value,
1429                                                               size_t                offset,
1430                                                               size_t                size) {
1431    GGML_ASSERT(tensor);
1432    memset((char *) tensor->data + offset, value, size);
1433
1434    GGML_UNUSED(buffer);
1435}
1436
1437static void ggml_backend_riscv64_spacemit_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
1438    GGML_ASSERT(buffer);
1439
1440    void * base = buffer->context;
1441    GGML_ASSERT(base != nullptr);
1442    memset(base, value, buffer->size);
1443}
1444
1445static void ggml_backend_riscv64_spacemit_buffer_set_tensor(ggml_backend_buffer_t buffer,
1446                                                            ggml_tensor *         tensor,
1447                                                            const void *          data,
1448                                                            size_t                offset,
1449                                                            size_t                size) {
1450    GGML_ASSERT(offset == 0);
1451    GGML_ASSERT(size == ggml_nbytes(tensor));
1452
1453    auto tensor_traits = (ggml::cpu::riscv64_spacemit::tensor_traits_base *) tensor->extra;
1454    if (tensor_traits) {
1455        auto OK = tensor_traits->repack(tensor, data, size);
1456        GGML_ASSERT(OK == 0);
1457    }
1458
1459    GGML_UNUSED(buffer);
1460}
1461
1462static const ggml_backend_buffer_i ggml_backend_riscv64_spacemit_buffer_i = {
1463    /* .free_buffer     = */ ggml_backend_riscv64_spacemit_buffer_free_buffer,
1464    /* .get_base        = */ ggml_backend_riscv64_spacemit_buffer_get_base,
1465    /* .init_tensor     = */ ggml_backend_riscv64_spacemit_buffer_init_tensor,
1466    /* .memset_tensor   = */ ggml_backend_riscv64_spacemit_buffer_memset_tensor,
1467    /* .set_tensor      = */ ggml_backend_riscv64_spacemit_buffer_set_tensor,
1468    /* .get_tensor      = */ nullptr,
1469    /* .set_tensor_2d   = */ nullptr,
1470    /* .get_tensor_2d   = */ nullptr,
1471    /* .cpy_tensor      = */ nullptr,
1472    /* .clear           = */ ggml_backend_riscv64_spacemit_buffer_clear,
1473    /* .reset           = */ nullptr,
1474};
1475
1476static const char * ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
1477    return "CPU_RISCV64_SPACEMIT";
1478
1479    GGML_UNUSED(buft);
1480}
1481
1482static ggml_backend_buffer_t ggml_backend_cpu_riscv64_spacemit_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft,
1483                                                                                        size_t size) {
1484    void * base = ggml::cpu::riscv64_spacemit::spine_mem_pool_alloc(size, 64);
1485    if (base == nullptr) {
1486        return nullptr;
1487    }
1488
1489    return ggml_backend_buffer_init(buft, ggml_backend_riscv64_spacemit_buffer_i, base, size);
1490}
1491
1492static size_t ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
1493    return 64;
1494
1495    GGML_UNUSED(buft);
1496}
1497
1498static size_t ggml_backend_cpu_riscv64_spacemit_nbytes(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
1499    for (int i = 0; i < GGML_MAX_DIMS; ++i) {
1500        if (tensor->ne[i] <= 0) {
1501            return 0;
1502        }
1503    }
1504
1505    GGML_UNUSED(buft);
1506
1507    const auto plain_nbytes = [&]() {
1508        size_t total = ggml_type_size(tensor->type);
1509        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
1510            total += (tensor->ne[i] - 1) * tensor->nb[i];
1511        }
1512        return total;
1513    };
1514
1515    const size_t blck_size = ggml_blck_size(tensor->type);
1516    if (blck_size == 1) {
1517        return plain_nbytes();
1518    }
1519
1520    const size_t row_nbytes = tensor->ne[0] * tensor->nb[0] / blck_size;
1521
1522    const auto add_strided_nbytes = [&](size_t total, size_t src_block_size, size_t dst_block_size) {
1523        for (int i = 1; i < GGML_MAX_DIMS; ++i) {
1524            total += (tensor->ne[i] - 1) * (tensor->nb[i] / src_block_size) * dst_block_size;
1525        }
1526        return total;
1527    };
1528
1529    const auto remap_block_nbytes = [&](size_t src_block_size, size_t dst_block_size, int64_t padded_rows = 0) {
1530        GGML_ASSERT(row_nbytes % src_block_size == 0);
1531
1532        size_t total =
1533            add_strided_nbytes((row_nbytes / src_block_size) * dst_block_size, src_block_size, dst_block_size);
1534
1535        if (padded_rows > 0 && tensor->ne[1] % padded_rows != 0) {
1536            total += (padded_rows - tensor->ne[1] % padded_rows) * (tensor->nb[1] / src_block_size) * dst_block_size;
1537        }
1538
1539        return total;
1540    };
1541
1542    size_t nbytes = row_nbytes;
1543    switch (tensor->type) {
1544        case GGML_TYPE_Q4_K:
1545            nbytes = remap_block_nbytes(sizeof(block_q4_K), sizeof(block_q4_1) * 8);
1546            break;
1547        case GGML_TYPE_Q6_K:
1548            nbytes = remap_block_nbytes(sizeof(block_q6_K), sizeof(block_q8_0) * 8, 32);
1549            break;
1550        case GGML_TYPE_Q8_0:
1551            nbytes = remap_block_nbytes(sizeof(block_q8_0), sizeof(block_q8_0), 32);
1552            break;
1553        case GGML_TYPE_Q2_K:
1554            nbytes = remap_block_nbytes(sizeof(block_q2_K), sizeof(spacemit_kernels::nrow_block_q2_k<1>));
1555            break;
1556        case GGML_TYPE_Q3_K:
1557            nbytes = remap_block_nbytes(sizeof(block_q3_K), sizeof(spacemit_kernels::nrow_block_q3_k<1>));
1558            break;
1559        case GGML_TYPE_MXFP4:
1560            nbytes = remap_block_nbytes(sizeof(block_mxfp4), sizeof(spacemit_kernels::nrow_block_mxfp4<1>));
1561            break;
1562        case GGML_TYPE_Q5_K:
1563            nbytes = remap_block_nbytes(sizeof(block_q5_K), sizeof(spacemit_kernels::nrow_block_q5_1<1>) * 8);
1564            break;
1565        case GGML_TYPE_Q5_1:
1566            nbytes = remap_block_nbytes(sizeof(block_q5_1), sizeof(spacemit_kernels::nrow_block_q5_1<1>));
1567            break;
1568        case GGML_TYPE_Q5_0:
1569            nbytes = remap_block_nbytes(sizeof(block_q5_0), sizeof(spacemit_kernels::nrow_block_q5_0<1>));
1570            break;
1571        default:
1572            nbytes = add_strided_nbytes(row_nbytes, 1, 1);
1573            break;
1574    }
1575
1576    return nbytes;
1577}
1578
1579namespace ggml::cpu::riscv64_spacemit {
1580
1581class extra_buffer_type : ggml::cpu::extra_buffer_type {
1582    bool supports_op(ggml_backend_dev_t, const ggml_tensor * op) override {
1583        switch (op->op) {
1584            case GGML_OP_MUL_MAT:
1585                if (op->src[0]->buffer && (ggml_n_dims(op->src[0]) == 2) &&
1586                    op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type() &&
1587                    ggml_riscv64_spacemit_get_optimal_repack_type(op->src[0])) {
1588                    if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
1589                        return false;
1590                    }
1591                    if (op->src[1]->type == GGML_TYPE_F32) {
1592                        return true;
1593                    }
1594                }
1595                break;
1596            case GGML_OP_MUL_MAT_ID:
1597                if (op->src[0]->buffer && (ggml_n_dims(op->src[0]) == 3) &&
1598                    op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type() &&
1599                    ggml_riscv64_spacemit_get_optimal_repack_type(op->src[0])) {
1600                    if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
1601                        return false;
1602                    }
1603                    if (op->src[1]->type == GGML_TYPE_F32) {
1604                        return true;
1605                    }
1606                }
1607                break;
1608            default:
1609                // GGML_ABORT("fatal error");
1610                break;
1611        }
1612        return false;
1613    }
1614
1615    ggml::cpu::tensor_traits * get_tensor_traits(const ggml_tensor * op) override {
1616        switch (op->op) {
1617            case GGML_OP_MUL_MAT:
1618            case GGML_OP_MUL_MAT_ID:
1619                if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type()) {
1620                    return (ggml::cpu::tensor_traits *) op->src[0]->extra;
1621                }
1622                break;
1623            case GGML_OP_NORM:
1624            case GGML_OP_RMS_NORM:
1625            case GGML_OP_ADD:
1626            case GGML_OP_SUB:
1627            case GGML_OP_MUL:
1628            case GGML_OP_DIV:
1629            case GGML_OP_FLASH_ATTN_EXT:
1630            case GGML_OP_CONT:
1631            case GGML_OP_CPY:
1632            case GGML_OP_REPEAT:
1633            case GGML_OP_SUM_ROWS:
1634            case GGML_OP_GET_ROWS:
1635            case GGML_OP_CONCAT:
1636                // case GGML_OP_GATED_DELTA_NET:
1637                return (ggml::cpu::tensor_traits *) (&ggml::cpu::riscv64_spacemit::rvv_impl);
1638            default:
1639                // GGML_ABORT("fatal error");
1640                break;
1641        }
1642
1643        return nullptr;
1644    }
1645};
1646
1647}  // namespace ggml::cpu::riscv64_spacemit
1648
1649ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void) {
1650    static ggml_backend_buffer_type ggml_backend_cpu_buffer_type_riscv64_spacemit = {
1651  /* .iface    = */
1652        {
1653         /* .get_name         = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name,
1654         /* .alloc_buffer     = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_alloc_buffer,
1655         /* .get_alignment    = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment,
1656         /* .get_max_size     = */ nullptr,
1657         /* .get_alloc_size   = */ ggml_backend_cpu_riscv64_spacemit_nbytes,
1658         /* .is_host          = */ nullptr,
1659         },
1660 /* .device  = */
1661        ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0),
1662 /* .context = */
1663        new ggml::cpu::riscv64_spacemit::extra_buffer_type(),
1664    };
1665
1666    return &ggml_backend_cpu_buffer_type_riscv64_spacemit;
1667}
1668
1669extern "C" {
1670static int bind_ai_thread() {
1671    int  fd, bytes;
1672    char str[32];
1673
1674    fd = open("/proc/set_ai_thread", O_WRONLY);
1675    if (fd < 0) {
1676        GGML_LOG_ERROR("try open /proc/set_ai_thread failed\n");
1677        return -1;
1678    }
1679
1680    snprintf(str, 16, "%d", 0);
1681    bytes = write(fd, str, strlen(str));
1682    if (bytes < 0) {
1683        GGML_LOG_ERROR("try write /proc/set_ai_thread failed\n");
1684        close(fd);
1685        return -1;
1686    }
1687
1688    close(fd);
1689    return 0;
1690}
1691
1692void ggml_backend_cpu_riscv64_spacemit_set_numa_thread_affinity(int thread_n) {
1693    int cpu_id = sched_getcpu();
1694    if (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2 &&
1695        !((1 << cpu_id) & ggml::cpu::riscv64_spacemit::global_spine_env_info.cpu_mask)) {
1696        GGML_PRINT_DEBUG("bind_ai_thread for thread %d, pid %d\n", thread_n, getpid());
1697        bind_ai_thread();
1698    }
1699
1700    if (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_tcm &&
1701        ggml::cpu::riscv64_spacemit::tls_context.cpu_id == -1) {
1702        CPU_ZERO(&(ggml::cpu::riscv64_spacemit::tls_context.cpuset));
1703        pthread_t    main_thread     = pthread_self();
1704        const auto & perfer_core_ids = ggml::cpu::riscv64_spacemit::global_spine_env_info.perfer_core_ids;
1705        if (thread_n < 0 || static_cast<size_t>(thread_n) >= perfer_core_ids.size()) {
1706            GGML_ABORT("thread_n %d exceeds perfer_core_ids size %zu\n", thread_n, perfer_core_ids.size());
1707        }
1708        auto perfer_cpu_id = perfer_core_ids[static_cast<size_t>(thread_n)];
1709        CPU_SET(perfer_cpu_id, &(ggml::cpu::riscv64_spacemit::tls_context.cpuset));
1710        int s =
1711            pthread_setaffinity_np(main_thread, sizeof(cpu_set_t), &(ggml::cpu::riscv64_spacemit::tls_context.cpuset));
1712        if (s != 0) {
1713            GGML_ABORT("set thread affinity error for thread_n %d, cpu_id %d\n", thread_n, perfer_cpu_id);
1714        }
1715
1716        int ai_cpu_id = perfer_cpu_id - ggml::cpu::riscv64_spacemit::global_spine_env_info.aicpu_id_offset;
1717        ggml::cpu::riscv64_spacemit::tls_context.cpu_id = ai_cpu_id;
1718        ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer =
1719            ggml::cpu::riscv64_spacemit::spine_mem_pool_tcm_mem_get(ai_cpu_id);
1720        ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size =
1721            ggml::cpu::riscv64_spacemit::global_spine_env_info.tcm_blk_size;
1722    }
1723
1724    if (ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer != nullptr) {
1725        void * rt =
1726            ggml::cpu::riscv64_spacemit::spine_mem_pool_tcm_mem_wait(ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
1727        if (rt == nullptr) {
1728            GGML_ABORT("wait tcm buffer failed for cpu_id: %d", ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
1729        }
1730    }
1731}
1732
1733void ggml_backend_cpu_riscv64_spacemit_clear_numa_thread_affinity_threaded(int thread_n) {
1734    if (ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer != nullptr) {
1735        auto rt = ggml::cpu::riscv64_spacemit::spine_mem_pool_tcm_mem_release(
1736            ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
1737        if (rt != 0) {
1738            GGML_ABORT("release tcm buffer failed for cpu_id: %d", ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
1739        }
1740    }
1741}
1742}