@novastera-oss/llamarn 0.4.1 → 0.4.3-beta4

package/cpp/llama.cpp/ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -0,0 +1,1025 @@
+#define GGML_COMMON_IMPL_CPP
+#define GGML_COMMON_DECL_CPP
+
+#include "ime.h"
+
+#include "ggml-backend-impl.h"
+#include "ggml-common.h"
+#include "ggml-cpu.h"
+#include "ime_kernels.h"
+#include "traits.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdio>  // for GGML_ASSERT
+#include <stdexcept>
+#include <thread>
+
+// clang-format off
+#if defined(__riscv)
+
+#if !defined(__riscv_v) || !defined(__riscv_v_intrinsic)
+#error "riscv v extension or v_intrinsic not enabled"
+#else
+#include <riscv_vector.h>
+#endif
+
+#if !defined(__riscv_zfh)
+#error "riscv zfh extension not enabled"
+#endif
+
+#if defined(RISCV64_SPACEMIT_IME1)
+#else
+#error "RISCV64_SPACEMIT_IME1 not defined"
+#endif
+
+#else
+
+#error "riscv not enabled in this build"
+
+#endif
+
+#if defined(__GNUC__)
+#pragma GCC diagnostic ignored "-Woverlength-strings"
+#pragma GCC diagnostic ignored "-Wcast-qual"
+#pragma GCC diagnostic ignored "-Wunused-parameter"
+#endif
+
+#if defined(RISCV64_SPACEMIT_IME1)
+#define QGEMM_STRIDEN_THREAD_ALIGN 16
+#else
+#define QGEMM_STRIDEN_THREAD_ALIGN 32
+#endif
+
+// clang-format on
+
+struct qnbitgemm_spacemit_ime_args {
+    const float *     a_ptr               = nullptr;
+    size_t            lda                 = 0;
+    const std::byte * packed_quant_b_data = nullptr;
+    const float *     quant_b_scale       = nullptr;
+    const void *      quant_b_zp          = nullptr;
+    const float *     quant_b_blksum      = nullptr;
+    const float *     bias                = nullptr;
+    float *           c_ptr               = nullptr;
+    size_t            ldc                 = 0;
+};
+
+constexpr size_t div_round_up(size_t up, size_t down) {
+    return (up + down - 1) / down;
+}
+
+constexpr size_t q8_blk_size(size_t blk_len) {
+    const size_t blk_size = sizeof(float) + blk_len * sizeof(int8_t);
+    // Currently, the strictest alignment requirement of a block is for a float.
+    // Ensure contiguous blocks are suitably aligned.
+    assert(blk_size % alignof(float) == 0);
+    return blk_size;
+}
+
+namespace ggml::cpu::riscv64_spacemit {
+
+const int num_ai_cores = std::thread::hardware_concurrency() / 2;
+
+}  // namespace ggml::cpu::riscv64_spacemit
+
+static void sqnbitgemm_spacemit_ime_i8i4(const size_t                        blk_len,
+                                         const size_t                        gemm_k,
+                                         const qnbitgemm_spacemit_ime_args * gemm_args,
+                                         void * const                        per_gemm_ws,
+                                         const size_t                        m_start,
+                                         const size_t                        m_count,
+                                         const size_t                        n_start,
+                                         const size_t                        n_count) {
+    constexpr size_t scale_stride = sizeof(uint16_t);
+    constexpr size_t blk_bitwidth = 4;
+
+    const size_t k_blks = div_round_up(gemm_k, blk_len);
+
+    const size_t      lda         = k_blks * q8_blk_size(blk_len);
+    const size_t      ldc         = gemm_args->ldc;
+    const size_t      ldb         = k_blks * (blk_len * blk_bitwidth / 8);
+    const std::byte * quant_a_ptr = static_cast<const std::byte *>(per_gemm_ws) + m_start * lda;
+
+    const size_t      zero_point_stride   = gemm_args->quant_b_zp != nullptr ? sizeof(uint8_t) : 0;
+    const size_t      packed_b_stride     = ldb + k_blks * (scale_stride + zero_point_stride);
+    const std::byte * packed_quant_b_data = gemm_args->packed_quant_b_data + n_start * packed_b_stride;
+
+    float * c_ptr = gemm_args->c_ptr + m_start * ldc + n_start;
+
+    size_t       count_n               = 0;
+    const size_t compute_block_count_n = m_count == 1 ? n_count : 16;
+    for (size_t n = 0; n < n_count; n += count_n) {
+        count_n = std::min(n_count - n, compute_block_count_n);
+
+        const std::byte * a_row    = quant_a_ptr;
+        const std::byte * b_col    = packed_quant_b_data + n * packed_b_stride;
+        const std::byte * b_col_zp = (zero_point_stride != 0) ? b_col : nullptr;
+        float *           c_blk    = c_ptr + n;
+
+        int32_t rows_remaining = m_count;
+
+        while (rows_remaining > 0) {
+            const auto rows_handled = sqnbitgemm_spacemit_ime::ime1::gemm_kernel_i8i4(
+                blk_len, a_row, b_col, nullptr, b_col_zp, c_blk, rows_remaining, count_n, gemm_k, k_blks, ldc, nullptr,
+                scale_stride);
+
+            c_blk += rows_handled * ldc;
+            a_row += rows_handled * lda;
+
+            rows_remaining -= rows_handled;
+        }
+    }
+}
+
+template <int K> constexpr int QK_0() {
+    if constexpr (K == 4) {
+        return QK4_0;
+    }
+    if constexpr (K == 8) {
+        return QK8_0;
+    }
+    return -1;
+}
+
+template <int K, int N> struct block {
+    ggml_half d[N];                         // deltas for N qK_0 blocks
+    uint8_t   qs[(QK_0<K>() * N * K) / 8];  // quants for N qK_0 blocks
+};
+
+template <int K, int N> struct block_with_zp {
+    ggml_half d[N];                         // deltas for N qK_1 blocks
+    uint8_t   zp[N];                        // zero points for N qK_1 blocks
+    uint8_t   qs[(QK_0<K>() * N * K) / 8];  // quants for N qK_1 blocks
+};
+
+// control size
+static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 8, "wrong block<4,16> size/padding");
+static_assert(sizeof(block_with_zp<4, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 8 + 16 * sizeof(uint8_t),
+              "wrong block_with_zp<4,16> size/padding");
+static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK4_0 * 16, "wrong block<8,16> size/padding");
+
+using block_q4_0x16 = block<4, 16>;
+using block_q4_1x16 = block_with_zp<4, 16>;
+using block_q8_0x16 = block<8, 16>;
+
+static block_q4_0x16 make_block_q4_0x16(block_q4_0 * in, unsigned int blck_size_interleave) {
+    block_q4_0x16 out;
+    GGML_ASSERT(QK4_0 / blck_size_interleave == 2);
+
+    for (int i = 0; i < 16; i++) {
+        out.d[i] = in[i].d;
+    }
+
+    for (int i = 0; i < 16; i++) {
+        // [0, 15], in.d & 0x0F
+        for (int j = 0; j < QK4_0 / 4; j++) {
+            //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+            //dst [b0 b8] ......... [b7 b15]
+            out.qs[i * QK4_0 / 4 + j] = (in[i].qs[j] & 0x0F) | ((in[i].qs[j + QK4_0 / 4] & 0x0F) << 4);
+        }
+    }
+
+    for (int i = 0; i < 16; i++) {
+        // [16, 31], in.d & 0xF0
+        for (int j = 0; j < QK4_0 / 4; j++) {
+            //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+            //dst [b16 b24] ......... [b23 b31]
+            out.qs[4 * QK4_0 + i * QK4_0 / 4 + j] = ((in[i].qs[j] & 0xF0) >> 4) | (in[i].qs[j + QK4_0 / 4] & 0xF0);
+        }
+    }
+
+    return out;
+}
+
+static block_q4_1x16 make_block_q4_1x16(block_q4_1 * in, unsigned int blck_size_interleave) {
+    block_q4_1x16 out;
+    GGML_ASSERT(QK4_1 / blck_size_interleave == 2);
+
+    for (int i = 0; i < 16; i++) {
+        float d   = GGML_FP16_TO_FP32(in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d);
+        float m   = GGML_FP16_TO_FP32(in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.m);
+        float mid = -std::nearbyintf(m / d);
+        mid       = std::min(15.0f, std::max(0.0f, mid));
+        out.d[i]  = GGML_FP32_TO_FP16(d);
+        out.zp[i] = static_cast<uint8_t>(mid);
+    }
+
+    for (int i = 0; i < 16; i++) {
+        // [0, 15], in.d & 0x0F
+        for (int j = 0; j < QK4_1 / 4; j++) {
+            //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+            //dst [b0 b8] ......... [b7 b15]
+            out.qs[i * QK4_1 / 4 + j] = (in[i].qs[j] & 0x0F) | ((in[i].qs[j + QK4_1 / 4] & 0x0F) << 4);
+        }
+    }
+
+    for (int i = 0; i < 16; i++) {
+        // [16, 31], in.d & 0xF0
+        for (int j = 0; j < QK4_1 / 4; j++) {
+            //src [b0 b16] ......... [b8 b24] ......... [b15 b31]
+            //dst [b16 b24] ......... [b23 b31]
+            out.qs[4 * QK4_1 + i * QK4_1 / 4 + j] = ((in[i].qs[j] & 0xF0) >> 4) | (in[i].qs[j + QK4_1 / 4] & 0xF0);
+        }
+    }
+
+    return out;
+}
+
+static int repack_q4_0_to_q4_0_16_bl(struct ggml_tensor *       t,
+                                     int                        interleave_block,
+                                     const void * GGML_RESTRICT data,
+                                     size_t                     data_size) {
+    GGML_ASSERT(t->type == GGML_TYPE_Q4_0);
+    GGML_ASSERT(interleave_block == 16);
+
+    constexpr int nrows_interleaved = 16;
+
+    block_q4_0x16 *    dst = (block_q4_0x16 *) t->data;
+    const block_q4_0 * src = (const block_q4_0 *) data;
+    block_q4_0         dst_tmp[16];
+    int                nrow    = ggml_nrows(t);
+    int                nblocks = t->ne[0] / QK4_0;
+
+    GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0));
+
+    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_0 != 0) {
+        return -1;
+    }
+
+    for (int b = 0; b < nrow; b += nrows_interleaved) {
+        for (int64_t x = 0; x < nblocks; x++) {
+            for (int i = 0; i < nrows_interleaved; i++) {
+                dst_tmp[i] = src[x + i * nblocks];
+            }
+            *dst++ = make_block_q4_0x16(dst_tmp, interleave_block);
+        }
+        src += nrows_interleaved * nblocks;
+    }
+    return 0;
+
+    GGML_UNUSED(data_size);
+}
+
+static int repack_q4_1_to_q4_1_16_bl(struct ggml_tensor *       t,
+                                     int                        interleave_block,
+                                     const void * GGML_RESTRICT data,
+                                     size_t                     data_size) {
+    GGML_ASSERT(t->type == GGML_TYPE_Q4_1);
+    GGML_ASSERT(interleave_block == 16);
+
+    constexpr int nrows_interleaved = 16;
+
+    block_q4_1x16 *    dst = (block_q4_1x16 *) t->data;
+    const block_q4_1 * src = (const block_q4_1 *) data;
+    block_q4_1         dst_tmp[16];
+    int                nrow    = ggml_nrows(t);
+    int                nblocks = t->ne[0] / QK4_1;
+
+    GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_1));
+
+    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_1 != 0) {
+        return -1;
+    }
+
+    for (int b = 0; b < nrow; b += nrows_interleaved) {
+        for (int64_t x = 0; x < nblocks; x++) {
+            for (int i = 0; i < nrows_interleaved; i++) {
+                dst_tmp[i] = src[x + i * nblocks];
+            }
+            *dst++ = make_block_q4_1x16(dst_tmp, interleave_block);
+        }
+        src += nrows_interleaved * nblocks;
+    }
+    return 0;
+
+    GGML_UNUSED(data_size);
+}
+
+static inline void get_scale_min_k4(int                           j,
+                                    const uint8_t * GGML_RESTRICT q,
+                                    uint8_t * GGML_RESTRICT       d,
+                                    uint8_t * GGML_RESTRICT       m) {
+    if (j < 4) {
+        *d = q[j] & 63;
+        *m = q[j + 4] & 63;
+    } else {
+        *d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
+        *m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4);
+    }
+}
+
+static int repack_q4_k_to_q4_1_16_bl(struct ggml_tensor *       t,
+                                     int                        interleave_block,
+                                     const void * GGML_RESTRICT data,
+                                     size_t                     data_size) {
+    GGML_ASSERT(t->type == GGML_TYPE_Q4_K);
+    GGML_ASSERT(interleave_block == 16);
+    GGML_ASSERT(QK_K / QK4_1 == 8);
+
+    constexpr int nrows_interleaved = 16;
+
+    block_q4_1x16 *    dst = (block_q4_1x16 *) t->data;
+    const block_q4_K * src = (const block_q4_K *) data;
+    block_q4_1         dst_tmp[16];
+    int                nrow    = ggml_nrows(t);
+    int                nblocks = t->ne[0] / QK_K;
+
+    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK_K != 0) {
+        return -1;
+    }
+
+    for (int b = 0; b < nrow; b += nrows_interleaved) {
+        for (int64_t x = 0; x < nblocks; x++) {
+            for (int j = 0; j < 8; j++) {
+                for (int i = 0; i < nrows_interleaved; i++) {
+                    uint8_t     sc, m;
+                    const float d = GGML_FP16_TO_FP32(src[x + i * nblocks].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d);
+                    const float min =
+                        GGML_FP16_TO_FP32(src[x + i * nblocks].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin);
+                    get_scale_min_k4(j, src[x + i * nblocks].scales, &sc, &m);
+                    const float d1 = d * sc;
+                    const float m1 = min * m;
+
+                    dst_tmp[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d = GGML_FP32_TO_FP16(d1);
+                    dst_tmp[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.m = GGML_FP32_TO_FP16(-m1);
+                    // src -> [b0, b32] [b1, b33] ... [b31, b63]
+                    // dst -> [b0, b16] [b1, b17] ... [b15, b31] [b32, b48] [b33, b49] ... [b47, b63]
+                    const uint8_t * q                                  = src[x + i * nblocks].qs + (j / 2) * QK4_1;
+                    if (j % 2 == 0) {
+                        for (int ii = 0; ii < 16; ii++) {
+                            dst_tmp[i].qs[ii] = (q[ii] & 0x0F) | ((q[ii + 16] & 0x0F) << 4);
+                        }
+                    } else {
+                        for (int ii = 0; ii < 16; ii++) {
+                            dst_tmp[i].qs[ii] = ((q[ii] & 0xF0) >> 4) | (q[ii + 16] & 0xF0);
+                        }
+                    }
+                }
+                *dst++ = make_block_q4_1x16(dst_tmp, interleave_block);
+            }
+        }
+        src += nrows_interleaved * nblocks;
+    }
+    return 0;
+
+    GGML_UNUSED(data_size);
+}
+
+namespace ggml::cpu::riscv64_spacemit {
+
+template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS>
+int repack(struct ggml_tensor *, const void *, size_t);
+
+template <> int repack<block_q4_0, 8, 16>(struct ggml_tensor * t, const void * data, size_t data_size) {
+    return repack_q4_0_to_q4_0_16_bl(t, 16, data, data_size);
+}
+
+template <> int repack<block_q4_1, 8, 16>(struct ggml_tensor * t, const void * data, size_t data_size) {
+    return repack_q4_1_to_q4_1_16_bl(t, 16, data, data_size);
+}
+
+template <> int repack<block_q4_K, 8, 16>(struct ggml_tensor * t, const void * data, size_t data_size) {
+    return repack_q4_k_to_q4_1_16_bl(t, 16, data, data_size);
+}
+
+class tensor_traits_base : public ggml::cpu::tensor_traits {
+  public:
+    virtual int repack(struct ggml_tensor * t, const void * data, size_t data_size) = 0;
+};
+
+template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_traits : public tensor_traits_base {
+    bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
+        switch (op->op) {
+            case GGML_OP_MUL_MAT:
+                size = ggml_row_size(GGML_TYPE_Q8_0, ggml_nelements(op->src[1])) * 4;
+                size = ((size + QK4_0 - 1) / QK4_0) * (QK4_0 * sizeof(float) + sizeof(float));
+                return true;
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+        return false;
+    }
+
+    bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override {
+        switch (op->op) {
+            case GGML_OP_MUL_MAT:
+                if (op->src[0]->type == GGML_TYPE_Q4_0 ||  //
+                    op->src[0]->type == GGML_TYPE_Q4_1 ||  //
+                    op->src[0]->type == GGML_TYPE_Q4_K) {
+                    forward_mul_mat_q4(params, op);
+                    return true;
+                }
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+        return false;
+    }
+
+    void forward_mul_mat_q4(ggml_compute_params * params, ggml_tensor * op) {
+        const ggml_tensor * src0 = op->src[0];
+        const ggml_tensor * src1 = op->src[1];
+        ggml_tensor *       dst  = op;
+
+        GGML_TENSOR_BINARY_OP_LOCALS
+
+        int ith = params->ith;
+        int nth = params->nth;
+
+        [[maybe_unused]] const enum ggml_type type = src0->type;
+
+        void *        w_data  = (void *) src0->data;
+        const float * feature = (const float *) src1->data;
+        float *       output  = (float *) dst->data;
+
+        const size_t                  batch_feature = ne12 * ne13;
+        [[maybe_unused]] const size_t batch_weight  = ne02 * ne03;
+        const size_t                  gemm_m        = ne11;
+        const size_t                  gemm_k        = ne10;
+        const size_t                  gemm_n        = ne01;
+
+        GGML_ASSERT(batch_weight == 1);
+
+        const size_t block_count_k           = div_round_up(gemm_k, QK4_0);
+        const size_t per_gemm_workspace_size = gemm_m * block_count_k * q8_blk_size(QK4_0);
+        const size_t per_gemm_workspace_stride =
+            div_round_up(per_gemm_workspace_size, alignof(uint64_t)) * alignof(uint64_t);
+        const size_t gemm_workspace_size = batch_feature * per_gemm_workspace_stride;
+        const size_t desired_wsize       = gemm_workspace_size + alignof(uint64_t) - 1;
+
+        if (ith == 0 && params->wsize < desired_wsize) {
+            throw std::runtime_error("wsize less than desired_wsize");
+        }
+
+        std::vector<qnbitgemm_spacemit_ime_args> qnbitgemm_args(batch_feature);
+
+        for (size_t i = 0; i < batch_feature; i++) {
+            qnbitgemm_args[i].a_ptr               = feature + gemm_m * gemm_k * i;
+            qnbitgemm_args[i].lda                 = gemm_k;
+            qnbitgemm_args[i].packed_quant_b_data = (const std::byte *) w_data;
+            qnbitgemm_args[i].quant_b_scale       = nullptr;
+
+            if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0>) {
+                qnbitgemm_args[i].quant_b_zp = nullptr;
+            } else {
+                qnbitgemm_args[i].quant_b_zp = w_data;
+            }
+
+            qnbitgemm_args[i].bias  = nullptr;
+            qnbitgemm_args[i].c_ptr = output + gemm_m * gemm_n * i;
+            qnbitgemm_args[i].ldc   = gemm_n;
+        }
+
+        const uintptr_t ws_ptr = reinterpret_cast<uintptr_t>(params->wdata);
+        void *          ws = reinterpret_cast<void *>((ws_ptr + alignof(uint64_t) - 1) & (~(alignof(uint64_t) - 1)));
+        const size_t    quant_a_stride = block_count_k * q8_blk_size(QK4_0);
+
+        {
+            constexpr size_t block_size_m           = 4;
+            size_t           per_gemm_block_count_m = div_round_up(gemm_m, block_size_m);
+            int32_t          task_count             = batch_feature * per_gemm_block_count_m;
+            int32_t          task_per_thread        = (task_count + nth - 1) / nth;
+            int32_t          start                  = ith * task_per_thread;
+            int32_t          end                    = std::min((ith + 1) * task_per_thread, task_count);
+            for (int32_t compute_idx = start; compute_idx < end; compute_idx++) {
+                int32_t                             gemm_idx = compute_idx / per_gemm_block_count_m;
+                int32_t                             block_idx_in_gemm = compute_idx % per_gemm_block_count_m;
+                int32_t                             m_idx    = block_idx_in_gemm * block_size_m;
+                const qnbitgemm_spacemit_ime_args & data     = qnbitgemm_args[gemm_idx];
+                int32_t rows_tobe_handled = (gemm_m - m_idx) > block_size_m ? block_size_m : (gemm_m - m_idx);
+
+                if (rows_tobe_handled == block_size_m) {
+                    const float * a_row_ptr = data.a_ptr + m_idx * data.lda;
+                    std::byte *   quant_a_row_ptr =
+                        static_cast<std::byte *>(ws) + gemm_idx * per_gemm_workspace_stride + m_idx * quant_a_stride;
+                    sqnbitgemm_spacemit_ime::ime1::quantize_a_4row_i8(QK4_0, a_row_ptr, gemm_k, quant_a_row_ptr);
+                } else {
+                    while (rows_tobe_handled) {
+                        const float * a_row_ptr       = data.a_ptr + m_idx * data.lda;
+                        std::byte *   quant_a_row_ptr = static_cast<std::byte *>(ws) +
+                                                      gemm_idx * per_gemm_workspace_stride + m_idx * quant_a_stride;
+                        sqnbitgemm_spacemit_ime::ime1::quantize_a_row_i8(QK4_0, a_row_ptr, gemm_k, quant_a_row_ptr);
+                        rows_tobe_handled -= 1;
+                        m_idx += 1;
+                    }
+                }
+            }
+        }
+
+        ggml_barrier(params->threadpool);
+
+        if (ith >= ggml::cpu::riscv64_spacemit::num_ai_cores) {
+            return;
+        }
+        nth = std::min(nth, int{ ggml::cpu::riscv64_spacemit::num_ai_cores });
+
+        size_t           threads_per_gemm = nth / batch_feature;
+        constexpr size_t gemm_m_stride    = 128;
+        size_t           nc               = gemm_n;
+        const size_t     gemm_m_blocked   = div_round_up(gemm_m, gemm_m_stride);
+        const size_t     max_nc           = div_round_up(gemm_n * gemm_m_blocked, threads_per_gemm);
+        if (max_nc < nc) {
+            nc = std::min(nc, div_round_up(max_nc, QGEMM_STRIDEN_THREAD_ALIGN) * QGEMM_STRIDEN_THREAD_ALIGN);
+        }
+        const size_t gemm_n_stride  = nc;
+        const size_t thread_count_m = div_round_up(gemm_m, gemm_m_stride);
+        const size_t thread_count_n = div_round_up(gemm_n, gemm_n_stride);
+        threads_per_gemm            = thread_count_m * thread_count_n;
+
+        {
+            int task_count      = batch_feature * threads_per_gemm;
+            int task_per_thread = (task_count + nth - 1) / nth;
+            int start           = ith * task_per_thread;
+            int end             = std::min((ith + 1) * task_per_thread, task_count);
+            for (int compute_idx = start; compute_idx < end; compute_idx++) {
+                const auto   gemm_i = compute_idx / threads_per_gemm;
+                const auto   blk_i  = compute_idx % threads_per_gemm;
+                const auto * data   = &qnbitgemm_args[gemm_i];
+
+                const auto tid_n = blk_i / thread_count_m;
+                const auto tid_m = blk_i % thread_count_m;
+
+                const size_t m_start = tid_m * gemm_m_stride;
+                const size_t m_count = std::min(gemm_m - m_start, (size_t) gemm_m_stride);
+
+                const size_t n_start = tid_n * gemm_n_stride;
+                const size_t n_count = std::min(gemm_n - n_start, (size_t) gemm_n_stride);
+
+                void * per_gemm_ws = reinterpret_cast<std::byte *>(ws) + gemm_i * per_gemm_workspace_stride;
+
+                sqnbitgemm_spacemit_ime_i8i4(QK4_0, gemm_k, data, per_gemm_ws, m_start, m_count, n_start, n_count);
+            }
+        }
+    }
+
+    int repack(struct ggml_tensor * t, const void * data, size_t data_size) override {
+        GGML_LOG_DEBUG("%s: repack tensor %s with %s_%dx%d\n", __func__, t->name, ggml_type_name(t->type),
+                       (int) NB_COLS, (int) INTER_SIZE);
+        return ggml::cpu::riscv64_spacemit::repack<BLOC_TYPE, INTER_SIZE, NB_COLS>(t, data, data_size);
+    }
+};
+
+class tensor_traits_common : public tensor_traits_base {
+    bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
+        switch (op->op) {
+            case GGML_OP_NORM:
+            case GGML_OP_RMS_NORM:
+                size = 0;
+                return true;
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+        return false;
+    }
+
+    bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override {
+        switch (op->op) {
+            case GGML_OP_NORM:
+                forward_norm_f32(params, op);
+                return true;
+            case GGML_OP_RMS_NORM:
+                forward_rms_norm_f32(params, op);
+                return true;
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+        return false;
+    }
+
+    void forward_norm_f32(ggml_compute_params * params, ggml_tensor * op) {
+        const ggml_tensor * src0 = op->src[0];
+        ggml_tensor *       dst  = op;
+        GGML_ASSERT(ggml_are_same_shape(src0, dst));
+        GGML_ASSERT(src0->nb[0] == sizeof(float));
+
+        const int ith = params->ith;
+        const int nth = params->nth;
+
+        GGML_TENSOR_UNARY_OP_LOCALS
+
+        float epsilon;
+        memcpy(&epsilon, dst->op_params, sizeof(float));
+
+        GGML_ASSERT(epsilon > 0.0f);
+
+        auto * input  = (float *) src0->data;
+        auto * output = (float *) dst->data;
+
+        const auto hidden_size     = ne00;
+        const auto task_count      = ne01 * ne02 * ne03;
+        const auto task_per_thread = (task_count + nth - 1) / nth;
+
+        const auto task_begin = ith * task_per_thread;
+        const auto task_end   = std::min((ith + 1) * task_per_thread, task_count);
+
+        for (auto task_idx = task_begin; task_idx < task_end; task_idx++) {
+            auto   offset  = task_idx * hidden_size;
+            auto * p_input = const_cast<float *>(input + offset);
+
+            auto *       p_output      = output + offset;
+            auto *       p_temp_output = p_output;
+            auto *       p_gamma_data  = (const float *) nullptr;
+            auto *       p_beta_data   = (const float *) nullptr;
+            size_t       gvl           = __riscv_vsetvlmax_e32m4();
+            vfloat32m4_t sum           = __riscv_vfmv_v_f_f32m4(0.f, gvl);
+            vfloat32m4_t sum_sq        = __riscv_vfmv_v_f_f32m4(0.f, gvl);
+            int64_t      length        = hidden_size;
+            while (length > 0) {
+                gvl                   = __riscv_vsetvl_e32m4(length);
+                // load data
+                vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_input, gvl);
+
+                sum    = __riscv_vfadd_vv_f32m4(sum, src_data, gvl);
+                sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl);
+
+                __riscv_vse32_v_f32m4(p_temp_output, src_data, gvl);
+
+                p_input += gvl;
+                p_temp_output += gvl;
+                length -= gvl;
+            }
+
+            gvl = __riscv_vsetvlmax_e32m1();
+
+            float        mean   = 0.f;
+            vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.f, gvl);
+            vfloat32m1_t mean_v =
+                __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum, 0), __riscv_vget_v_f32m4_f32m1(sum, 1), gvl);
+            mean_v = __riscv_vfadd_vv_f32m1(mean_v, __riscv_vget_v_f32m4_f32m1(sum, 2), gvl);
+            mean_v = __riscv_vfadd_vv_f32m1(mean_v, __riscv_vget_v_f32m4_f32m1(sum, 3), gvl);
+            mean_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_v, zero_v, gvl);
+            mean   = __riscv_vfmv_f_s_f32m1_f32(mean_v);
+            mean /= hidden_size;
+
+            vfloat32m1_t mean_square_v = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum_sq, 0),
+                                                                __riscv_vget_v_f32m4_f32m1(sum_sq, 1), gvl);
+            mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 2), gvl);
+            mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 3), gvl);
+            mean_square_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_square_v, zero_v, gvl);
+
+            float mean_square = __riscv_vfmv_f_s_f32m1_f32(mean_square_v);
+            mean_square /= hidden_size;
+            mean_square = sqrt(mean_square - mean * mean + epsilon);
+
+            mean_square   = 1.0f / mean_square;
+            length        = hidden_size;
+            p_temp_output = p_output;
+
+            if (p_gamma_data == nullptr && p_beta_data == nullptr) {
+                while (length > 0) {
+                    gvl                   = __riscv_vsetvl_e32m4(length);
+                    vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
+                    src_data              = __riscv_vfsub_vf_f32m4(src_data, mean, gvl);
+                    src_data              = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
+                    __riscv_vse32_v_f32m4(p_output, src_data, gvl);
+                    p_temp_output += gvl;
+                    p_output += gvl;
+                    length -= gvl;
+                }
+            } else if (p_beta_data == nullptr) {
+                while (length > 0) {
+                    gvl                       = __riscv_vsetvl_e32m4(length);
+                    vfloat32m4_t src_data     = __riscv_vle32_v_f32m4(p_temp_output, gvl);
+                    vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
+                    src_data                  = __riscv_vfsub_vf_f32m4(src_data, mean, gvl);
+                    src_data                  = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
+                    src_data                  = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
+                    __riscv_vse32_v_f32m4(p_output, src_data, gvl);
+                    p_temp_output += gvl;
+                    p_output += gvl;
+                    p_gamma_data += gvl;
+                    length -= gvl;
+                }
+            } else if (p_gamma_data != nullptr) {
+                while (length > 0) {
+                    gvl                       = __riscv_vsetvl_e32m4(length);
+                    vfloat32m4_t src_data     = __riscv_vle32_v_f32m4(p_temp_output, gvl);
+                    vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
+                    src_data                  = __riscv_vfsub_vf_f32m4(src_data, mean, gvl);
+                    src_data                  = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
+                    src_data                  = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
+                    vfloat32m4_t beta_data_v  = __riscv_vle32_v_f32m4(p_beta_data, gvl);
+                    src_data                  = __riscv_vfadd_vv_f32m4(src_data, beta_data_v, gvl);
+                    p_beta_data += gvl;
+                    __riscv_vse32_v_f32m4(p_output, src_data, gvl);
+                    p_temp_output += gvl;
+                    p_output += gvl;
+                    p_gamma_data += gvl;
+                    length -= gvl;
+                }
+            }
+        }
+    }
+
+    void forward_rms_norm_f32(ggml_compute_params * params, ggml_tensor * op) {
+        const ggml_tensor * src0 = op->src[0];
+        ggml_tensor *       dst  = op;
+        GGML_ASSERT(ggml_are_same_shape(src0, dst));
+        GGML_ASSERT(src0->nb[0] == sizeof(float));
+
+        const int ith = params->ith;
+        const int nth = params->nth;
+
+        GGML_TENSOR_UNARY_OP_LOCALS
+
+        float epsilon;
+        memcpy(&epsilon, dst->op_params, sizeof(float));
+
+        GGML_ASSERT(epsilon > 0.0f);
+
+        auto * input  = (float *) src0->data;
+        auto * output = (float *) dst->data;
+
+        const auto hidden_size     = ne00;
+        const auto task_count      = ne01 * ne02 * ne03;
+        const auto task_per_thread = (task_count + nth - 1) / nth;
+
+        const auto task_begin = ith * task_per_thread;
+        const auto task_end   = std::min((ith + 1) * task_per_thread, task_count);
+
+        for (auto task_idx = task_begin; task_idx < task_end; task_idx++) {
+            auto   offset        = task_idx * hidden_size;
+            auto * p_input       = const_cast<float *>(input + offset);
+            auto * p_output      = output + offset;
+            auto * p_temp_output = p_output;
+            auto * p_gamma_data  = (const float *) nullptr;
+            auto * p_beta_data   = (const float *) nullptr;
+
+            size_t       gvl    = __riscv_vsetvlmax_e32m4();
+            // vfloat32m4_t sum = __riscv_vfmv_v_f_f32m4(0.f, gvl);
+            vfloat32m4_t sum_sq = __riscv_vfmv_v_f_f32m4(0.f, gvl);
+            int64_t      length = hidden_size;
+            while (length > 0) {
+                gvl                   = __riscv_vsetvl_e32m4(length);
+                // load data
+                vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_input, gvl);
+
+                sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl);
+
+                __riscv_vse32_v_f32m4(p_temp_output, src_data, gvl);
+
+                p_input += gvl;
+                p_temp_output += gvl;
+                length -= gvl;
+            }
+
+            gvl = __riscv_vsetvlmax_e32m1();
+
+            // float mean = 0.f;
+            vfloat32m1_t zero_v = __riscv_vfmv_v_f_f32m1(0.f, gvl);
+
+            vfloat32m1_t mean_square_v = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m4_f32m1(sum_sq, 0),
+                                                                __riscv_vget_v_f32m4_f32m1(sum_sq, 1), gvl);
+            mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 2), gvl);
+            mean_square_v = __riscv_vfadd_vv_f32m1(mean_square_v, __riscv_vget_v_f32m4_f32m1(sum_sq, 3), gvl);
+            mean_square_v = __riscv_vfredusum_vs_f32m1_f32m1(mean_square_v, zero_v, gvl);
+
+            float mean_square = __riscv_vfmv_f_s_f32m1_f32(mean_square_v);
+            mean_square /= hidden_size;
+
+            mean_square = sqrt(mean_square + epsilon);
+
+            mean_square   = 1.0f / mean_square;
+            length        = hidden_size;
+            p_temp_output = p_output;
+
+            if (p_gamma_data == nullptr && p_beta_data == nullptr) {
+                while (length > 0) {
+                    gvl                   = __riscv_vsetvl_e32m4(length);
+                    vfloat32m4_t src_data = __riscv_vle32_v_f32m4(p_temp_output, gvl);
+                    src_data              = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
+                    __riscv_vse32_v_f32m4(p_output, src_data, gvl);
+                    p_temp_output += gvl;
+                    p_output += gvl;
+                    length -= gvl;
+                }
+            } else if (p_beta_data == nullptr) {
+                while (length > 0) {
+                    gvl                       = __riscv_vsetvl_e32m4(length);
+                    vfloat32m4_t src_data     = __riscv_vle32_v_f32m4(p_temp_output, gvl);
+                    vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
+                    src_data                  = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
+                    src_data                  = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
+                    __riscv_vse32_v_f32m4(p_output, src_data, gvl);
+                    p_temp_output += gvl;
+                    p_output += gvl;
+                    p_gamma_data += gvl;
+                    length -= gvl;
+                }
+            } else if (p_gamma_data != nullptr) {
+                while (length > 0) {
+                    gvl                       = __riscv_vsetvl_e32m4(length);
+                    vfloat32m4_t src_data     = __riscv_vle32_v_f32m4(p_temp_output, gvl);
+                    vfloat32m4_t gamma_data_v = __riscv_vle32_v_f32m4(p_gamma_data, gvl);
+                    src_data                  = __riscv_vfmul_vf_f32m4(src_data, mean_square, gvl);
+                    src_data                  = __riscv_vfmul_vv_f32m4(src_data, gamma_data_v, gvl);
+                    vfloat32m4_t beta_data_v  = __riscv_vle32_v_f32m4(p_beta_data, gvl);
+                    src_data                  = __riscv_vfadd_vv_f32m4(src_data, beta_data_v, gvl);
+                    p_beta_data += gvl;
+                    __riscv_vse32_v_f32m4(p_output, src_data, gvl);
+                    p_temp_output += gvl;
+                    p_output += gvl;
+                    p_gamma_data += gvl;
+                    length -= gvl;
+                }
+            }
+        }
+    }
+
+    int repack(struct ggml_tensor * t, const void * data, size_t data_size) override {
+        memcpy(t->data, data, data_size);
+        return 0;
+    }
+};
+
+static const tensor_traits<block_q4_0, 8, 16> q4_0_16x8_q8_0;
+static const tensor_traits<block_q4_1, 8, 16> q4_1_16x8_q8_0;
+static const tensor_traits<block_q4_K, 8, 16> q4_k_16x8_q8_0;
+static const tensor_traits_common             rvv_impl;
+
+}  // namespace ggml::cpu::riscv64_spacemit
+
+static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack_type(const struct ggml_tensor * cur) {
+    if (cur->type == GGML_TYPE_Q4_0) {
+        if (cur->ne[1] % 16 == 0) {
+            return &ggml::cpu::riscv64_spacemit::q4_0_16x8_q8_0;
+        }
+    } else if (cur->type == GGML_TYPE_Q4_1) {
+        if (cur->ne[1] % 16 == 0) {
+            return &ggml::cpu::riscv64_spacemit::q4_1_16x8_q8_0;
+        }
+    } else if (cur->type == GGML_TYPE_Q4_K) {
+        if (cur->ne[1] % 16 == 0) {
+            return &ggml::cpu::riscv64_spacemit::q4_k_16x8_q8_0;
+        }
+    } else if (cur->type == GGML_TYPE_F32) {
+        return &ggml::cpu::riscv64_spacemit::rvv_impl;
+    }
+
+    return nullptr;
+}
+
+static enum ggml_status ggml_backend_riscv64_spacemit_buffer_init_tensor(ggml_backend_buffer_t buffer,
+                                                                         struct ggml_tensor *  tensor) {
+    tensor->extra =
+        (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_riscv64_spacemit_get_optimal_repack_type(tensor));
+
+    GGML_UNUSED(buffer);
+
+    return GGML_STATUS_SUCCESS;
+}
+
+static void ggml_backend_riscv64_spacemit_buffer_set_tensor(ggml_backend_buffer_t buffer,
+                                                            struct ggml_tensor *  tensor,
+                                                            const void *          data,
+                                                            size_t                offset,
+                                                            size_t                size) {
+    GGML_ASSERT(offset == 0);
+    GGML_ASSERT(size == ggml_nbytes(tensor));
+
+    auto tensor_traits = (ggml::cpu::riscv64_spacemit::tensor_traits_base *) tensor->extra;
+    if (tensor_traits) {
+        auto OK = tensor_traits->repack(tensor, data, size);
+        GGML_ASSERT(OK == 0);
+    }
+
+    GGML_UNUSED(buffer);
+}
+
+static const char * ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
+    return "CPU_RISCV64_SPACEMIT";
+
+    GGML_UNUSED(buft);
+}
+
+static ggml_backend_buffer_t ggml_backend_cpu_riscv64_spacemit_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft,
+                                                                                        size_t size) {
+    ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size);
+
+    if (buffer == nullptr) {
+        return nullptr;
+    }
+
+    buffer->buft              = buft;
+    buffer->iface.init_tensor = ggml_backend_riscv64_spacemit_buffer_init_tensor;
+    buffer->iface.set_tensor  = ggml_backend_riscv64_spacemit_buffer_set_tensor;
+    buffer->iface.get_tensor  = nullptr;
+    buffer->iface.cpy_tensor  = nullptr;
+    return buffer;
+}
+
+static size_t ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
+    return 64;
+
+    GGML_UNUSED(buft);
+}
+
+static size_t ggml_backend_cpu_riscv64_spacemit_nbytes(ggml_backend_buffer_type_t buft,
+                                                       const struct ggml_tensor * tensor) {
+    for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+        if (tensor->ne[i] <= 0) {
+            return 0;
+        }
+    }
+
+    size_t       nbytes;
+    const size_t blck_size = ggml_blck_size(tensor->type);
+    if (blck_size == 1) {
+        nbytes = ggml_type_size(tensor->type);
+        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+            nbytes += (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    } else {
+        nbytes = tensor->ne[0] * tensor->nb[0] / blck_size;
+        if (tensor->type == GGML_TYPE_Q4_K) {
+            GGML_ASSERT(nbytes % sizeof(block_q4_K) == 0);
+            nbytes = (nbytes / sizeof(block_q4_K)) * sizeof(block_q4_1) * 8;
+            for (int i = 1; i < GGML_MAX_DIMS; ++i) {
+                nbytes += (tensor->ne[i] - 1) * (tensor->nb[i] / sizeof(block_q4_K)) * sizeof(block_q4_1) * 8;
+            }
+        } else {
+            for (int i = 1; i < GGML_MAX_DIMS; ++i) {
+                nbytes += (tensor->ne[i] - 1) * tensor->nb[i];
+            }
+        }
+    }
+
+    GGML_UNUSED(buft);
+    return nbytes;
+}
+
+namespace ggml::cpu::riscv64_spacemit {
+
+class extra_buffer_type : ggml::cpu::extra_buffer_type {
+    bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override {
+        switch (op->op) {
+            case GGML_OP_MUL_MAT:
+                if (op->src[0]->buffer && (ggml_n_dims(op->src[0]) == 2) &&
+                    op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type() &&
+                    ggml_riscv64_spacemit_get_optimal_repack_type(op->src[0])) {
+                    if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
+                        return false;
+                    }
+                    if (op->src[1]->type == GGML_TYPE_F32) {
+                        return true;
+                    }
+                }
+                break;
+            case GGML_OP_NORM:
+            case GGML_OP_RMS_NORM:
+                if (op->src[0]->type == GGML_TYPE_F32) {
+                    return true;
+                }
+                break;
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+        return false;
+    }
+
+    ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {
+        switch (op->op) {
+            case GGML_OP_MUL_MAT:
+                if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_riscv64_spacemit_buffer_type()) {
+                    return (ggml::cpu::tensor_traits *) op->src[0]->extra;
+                }
+                break;
+            case GGML_OP_NORM:
+            case GGML_OP_RMS_NORM:
+                return (ggml::cpu::tensor_traits *) (&ggml::cpu::riscv64_spacemit::rvv_impl);
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+
+        return nullptr;
+    }
+};
+
+}  // namespace ggml::cpu::riscv64_spacemit
+
+ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void) {
+    static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type_riscv64_spacemit = {
+  /* .iface    = */
+        {
+         /* .get_name         = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name,
+         /* .alloc_buffer     = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_alloc_buffer,
+         /* .get_alignment    = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment,
+         /* .get_max_size     = */ nullptr,
+         /* .get_alloc_size   = */ ggml_backend_cpu_riscv64_spacemit_nbytes,
+         /* .is_host          = */ nullptr,
+         },
+ /* .device  = */
+        ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0),
+ /* .context = */
+        new ggml::cpu::riscv64_spacemit::extra_buffer_type(),
+    };
+
+    return &ggml_backend_cpu_buffer_type_riscv64_spacemit;
+}