whispercpp 1.3.6 → 1.3.7

data/ext/sources/ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -3,19 +3,32 @@
 
 #include "ime.h"
 
+#include "binary-ops.h"
+#include "common.h"
 #include "ggml-backend-impl.h"
 #include "ggml-common.h"
 #include "ggml-cpu.h"
+#include "ime_env.h"
 #include "ime_kernels.h"
+#include "ops.h"
+#include "repack.h"
+#include "rvv_kernels.h"
+#include "spine_mem_pool.h"
 #include "traits.h"
+#include "vec.h"
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <unistd.h>
 
 #include <algorithm>
+#include <atomic>
 #include <cassert>
+#include <cerrno>
 #include <cmath>
 #include <cstdio>  // for GGML_ASSERT
 #include <stdexcept>
 #include <thread>
-
 // clang-format off
 #if defined(__riscv)
 
@@ -25,13 +38,17 @@
 #include <riscv_vector.h>
 #endif
 
-#if !defined(__riscv_zfh)
-#error "riscv zfh extension not enabled"
+#if !defined(__riscv_zfh) || !defined(__riscv_zvfh)
+#error "riscv zfh extension not enabled, GGML_RV_ZFH and GGML_RV_ZVFH must be defined to 1"
 #endif
 
-#if defined(RISCV64_SPACEMIT_IME1)
+#if !defined(__riscv_zba)
+#error "riscv zba extension not enabled, GGML_RV_ZBA must be defined to 1"
+#endif
+
+#if defined(RISCV64_SPACEMIT_IME1) || defined(RISCV64_SPACEMIT_IME2)
 #else
-#error "RISCV64_SPACEMIT_IME1 not defined"
+#error "RISCV64_SPACEMIT_IME1 or RISCV64_SPACEMIT_IME2 not defined"
 #endif
 
 #else
@@ -46,382 +63,490 @@
 #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;
+extern "C" {
+extern void ggml_threadpool_chunk_set(struct ggml_threadpool * tp, int value);
+extern int  ggml_threadpool_chunk_add(struct ggml_threadpool * tp, int value);
 }
 
 namespace ggml::cpu::riscv64_spacemit {
 
-const int num_ai_cores = std::thread::hardware_concurrency() / 2;
-
-}  // namespace ggml::cpu::riscv64_spacemit
+struct TLSContext {
+    int       cpu_id{ -1 };
+    cpu_set_t cpuset;
+    void *    tcm_buffer{ nullptr };
+    size_t    tcm_buffer_size{ 0 };
+};
 
-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;
+thread_local TLSContext tls_context;
+
+template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> constexpr size_t get_repacked_block_type_size() {
+    if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0>) {
+        return sizeof(block_q8_0);
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0>) {
+        return sizeof(block_q4_0) * INTER_SIZE / QK4_0;
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_1> || std::is_same_v<BLOC_TYPE, block_q4_K>) {
+        return (sizeof(block_q4_0) + sizeof(uint8_t)) * INTER_SIZE / QK4_1;
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K>) {
+        return sizeof(spacemit_kernels::nrow_block_q2_k<1>);
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q3_K>) {
+        return sizeof(spacemit_kernels::nrow_block_q3_k<1>);
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_mxfp4>) {
+        return sizeof(spacemit_kernels::nrow_block_mxfp4<1>);
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_1> || std::is_same_v<BLOC_TYPE, block_q5_K>) {
+        return sizeof(spacemit_kernels::nrow_block_q5_1<1>);
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_0>) {
+        return sizeof(spacemit_kernels::nrow_block_q5_0<1>);
+    } else {
+        assert(false);
+        return 0;
+    }
+}
 
-    const size_t k_blks = div_round_up(gemm_k, blk_len);
+template <typename BLOC_TYPE> constexpr bool block_type_has_zp() {
+    if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0> ||
+                  std::is_same_v<BLOC_TYPE, block_q3_K> || std::is_same_v<BLOC_TYPE, block_q4_0> ||
+                  std::is_same_v<BLOC_TYPE, block_mxfp4> || std::is_same_v<BLOC_TYPE, block_q5_0>) {
+        return false;
+    } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_1> || std::is_same_v<BLOC_TYPE, block_q4_K> ||
+                         std::is_same_v<BLOC_TYPE, block_q2_K> || std::is_same_v<BLOC_TYPE, block_q5_1> ||
+                         std::is_same_v<BLOC_TYPE, block_q5_K>) {
+        return true;
+    } else {
+        assert(false);
+        return false;
+    }
+}
 
-    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;
+class tensor_traits_base : public ggml::cpu::tensor_traits {
+  public:
+    virtual int repack(ggml_tensor * t, const void * data, size_t data_size) = 0;
+};
 
-    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;
+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 ggml_tensor * op, size_t & size) override {
+        switch (op->op) {
+            case GGML_OP_MUL_MAT:
+                {
+                    int64_t src1_nelements = ggml_nelements(op->src[1]);
+
+                    if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K> || std::is_same_v<BLOC_TYPE, block_q3_K>) {
+                        size =
+                            spacemit_kernels::div_round_up(src1_nelements, QK_K) * spacemit_kernels::q8k_blk_size(QK_K);
+                    } else if constexpr (INTER_SIZE == QK4_0) {
+                        size = spacemit_kernels::div_round_up(src1_nelements, QK4_0) *
+                               spacemit_kernels::q8_blk_size(QK4_0, true);
+                    } else if constexpr (INTER_SIZE == 256) {
+                        size = spacemit_kernels::div_round_up(src1_nelements, 256) *
+                               spacemit_kernels::q8_hp_blk_size(256, true, true);
+                    } else {
+                        GGML_ABORT("unsupported block type");
+                    }
 
-    float * c_ptr = gemm_args->c_ptr + m_start * ldc + n_start;
+                    size = GGML_PAD(size, sizeof(int64_t));
 
-    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);
+                    return true;
+                }
+            case GGML_OP_MUL_MAT_ID:
+                {
+                    int64_t src1_nelements = ggml_nelements(op->src[1]);
+
+                    if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K> || std::is_same_v<BLOC_TYPE, block_q3_K>) {
+                        size =
+                            spacemit_kernels::div_round_up(src1_nelements, QK_K) * spacemit_kernels::q8k_blk_size(QK_K);
+                    } else if constexpr (INTER_SIZE == QK4_0) {
+                        size = spacemit_kernels::div_round_up(src1_nelements, QK4_0) *
+                               spacemit_kernels::q8_blk_size(QK4_0, true);
+                    } else if constexpr (INTER_SIZE == 256) {
+                        size = spacemit_kernels::div_round_up(src1_nelements, 256) *
+                               spacemit_kernels::q8_hp_blk_size(256, true, true);
+                    } else {
+                        GGML_ABORT("unsupported block type");
+                    }
 
-        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;
+                    size = GGML_PAD(size, sizeof(int64_t));
 
-        int32_t rows_remaining = m_count;
+                    const int64_t ne02 = op->src[0]->ne[2];  // n_as, n_expert
+                    const int64_t ne12 = op->src[1]->ne[2];  // n_tokens
 
-        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);
+                    const size_t sizeof_mmid_row_mapping = sizeof(int64_t);
+                    size += sizeof_mmid_row_mapping * ne02 * (ne12 + 1) + (ne02 + 1) * sizeof(int64_t);
 
-            c_blk += rows_handled * ldc;
-            a_row += rows_handled * lda;
+                    size = GGML_PAD(size, sizeof(int64_t));
 
-            rows_remaining -= rows_handled;
+                    return true;
+                }
+            default:
+                // GGML_ABORT("fatal error");
+                break;
         }
+        return false;
     }
-}
 
-template <int K> constexpr int QK_0() {
-    if constexpr (K == 4) {
-        return QK4_0;
-    }
-    if constexpr (K == 8) {
-        return QK8_0;
+    bool compute_forward(ggml_compute_params * params, ggml_tensor * op) override {
+        switch (op->op) {
+            case GGML_OP_MUL_MAT:
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_Q2_K:
+                    case GGML_TYPE_Q3_K:
+                    case GGML_TYPE_Q4_0:
+                    case GGML_TYPE_Q4_1:
+                    case GGML_TYPE_Q4_K:
+                    case GGML_TYPE_Q6_K:
+                    case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_Q5_1:
+                    case GGML_TYPE_Q5_K:
+                        //case GGML_TYPE_MXFP4:
+                        forward_mul_mat(params, op);
+                        return true;
+                    default:
+                        // GGML_ABORT("fatal error: unsupported type for src0 in MUL_MAT");
+                        return false;
+                }
+                break;
+            case GGML_OP_MUL_MAT_ID:
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_Q2_K:
+                    case GGML_TYPE_Q3_K:
+                    case GGML_TYPE_Q4_0:
+                    case GGML_TYPE_Q4_1:
+                    case GGML_TYPE_Q4_K:
+                    case GGML_TYPE_Q6_K:
+                    case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_Q5_1:
+                    case GGML_TYPE_Q5_K:
+                        //case GGML_TYPE_MXFP4:
+                        forward_mul_mat_id(params, op);
+                        return true;
+                    default:
+                        // GGML_ABORT("fatal error: unsupported type for src0 in MUL_MAT_ID");
+                        return false;
+                }
+                break;
+            default:
+                // GGML_ABORT("fatal error");
+                break;
+        }
+        return false;
     }
-    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
-};
+    void forward_mul_mat(ggml_compute_params * params, ggml_tensor * op) {
+        constexpr size_t a_blk_len = INTER_SIZE;
+        constexpr size_t b_blk_len = INTER_SIZE;
 
-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
-};
+        const ggml_tensor * src0 = op->src[0];
+        const ggml_tensor * src1 = op->src[1];
+        ggml_tensor *       dst  = op;
 
-// 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");
+        GGML_TENSOR_BINARY_OP_LOCALS
 
-using block_q4_0x16 = block<4, 16>;
-using block_q4_1x16 = block_with_zp<4, 16>;
-using block_q8_0x16 = block<8, 16>;
+        int ith = params->ith;
+        int nth = params->nth;
 
-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);
+        [[maybe_unused]] const enum ggml_type type = src0->type;
 
-    for (int i = 0; i < 16; i++) {
-        out.d[i] = in[i].d;
-    }
+        void *        w_data  = (void *) src0->data;
+        const float * feature = (const float *) src1->data;
+        float *       output  = (float *) dst->data;
 
-    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);
+        const int64_t gemm_m = ne11 * ne12 * ne13;
+        const int64_t gemm_k = ne10;
+        const int64_t gemm_n = ne01;
+
+        spacemit_kernels::quantize_a_row_def       quantize_a_row_i8;
+        spacemit_kernels::quantize_a_row_def       quantize_a_4row_i8;
+        spacemit_kernels::gemm_kernel_quantize_def gemm_kernel;
+        bool                                       set_kernel_impl = false;
+
+        int64_t block_stride_a = spacemit_kernels::q8_blk_size(a_blk_len);
+
+#if defined(RISCV64_SPACEMIT_IME2)
+        if (!set_kernel_impl && (global_spine_env_info.use_ime2)) {
+            quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8;
+            quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8;
+            block_stride_a     = spacemit_kernels::q8_blk_size(a_blk_len, true);
+
+            if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0>) {
+                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i8;
+                set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
+                                 std::is_same_v<BLOC_TYPE, block_q4_K>) {
+                if constexpr (INTER_SIZE == 256) {
+                    gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i4_hp;
+                    quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8_hp;
+                    quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8_hp;
+                    block_stride_a     = spacemit_kernels::q8_hp_blk_size(a_blk_len, true, true);
+                    set_kernel_impl    = true;
+                } else {
+                    gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i4;
+                    quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8;
+                    quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8;
+                    block_stride_a     = spacemit_kernels::q8_blk_size(a_blk_len, true);
+                    set_kernel_impl    = true;
+                }
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K>) {
+                quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8k;
+                quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8k;
+                block_stride_a     = spacemit_kernels::q8k_blk_size(a_blk_len);
+
+                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i2k;
+                set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q3_K>) {
+                quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8k;
+                quantize_a_4row_i8 = spacemit_kernels::rvv::quantize_a_4row_i8k;
+                block_stride_a     = spacemit_kernels::q8k_blk_size(a_blk_len);
+
+                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i3k;
+                set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_mxfp4>) {
+                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8mxfp4;
+                set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_1> || std::is_same_v<BLOC_TYPE, block_q5_K> ||
+                                 std::is_same_v<BLOC_TYPE, block_q5_0>) {
+                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i5;
+                set_kernel_impl = true;
+            }
         }
-    }
+#endif
 
-    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);
+#if defined(RISCV64_SPACEMIT_IME1)
+        if (!set_kernel_impl && (global_spine_env_info.use_ime1)) {
+            quantize_a_row_i8  = spacemit_kernels::ime1::quantize_a_row_i8;
+            quantize_a_4row_i8 = spacemit_kernels::ime1::quantize_a_4row_i8;
+
+            if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
+                          std::is_same_v<BLOC_TYPE, block_q4_K>) {
+                gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i4;
+                set_kernel_impl = true;
+            }
+        }
+#endif
+        if (!set_kernel_impl) {
+            GGML_ABORT("no kernel implementation found for the block type");
         }
-    }
 
-    return out;
-}
+        const int64_t a_k_blks = spacemit_kernels::div_round_up(gemm_k, a_blk_len);
+        const int64_t b_k_blks = spacemit_kernels::div_round_up(gemm_k, b_blk_len);
 
-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);
-    }
+        const int64_t row_stride_a        = a_k_blks * block_stride_a;
+        const int64_t gemm_workspace_size = GGML_PAD(gemm_m * row_stride_a, alignof(int64_t));
 
-    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);
+        if (ith == 0 && params->wsize < gemm_workspace_size) {
+            GGML_ABORT("wsize less than gemm_workspace_size");
         }
-    }
 
-    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);
-        }
-    }
+        uintptr_t ws_ptr = reinterpret_cast<uintptr_t>(params->wdata);
 
-    return out;
-}
+        void *        tcm_buffer      = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer;
+        const int64_t tcm_buffer_size = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size;
 
-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);
+        auto * quant_a_buffer = reinterpret_cast<uint8_t *>(ws_ptr);
 
-    constexpr int nrows_interleaved = 16;
+        constexpr int64_t row_align = 4;
+        const int64_t     row_blks  = spacemit_kernels::div_round_up(gemm_m, row_align);
 
-    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;
+        const int64_t row_stride_b      = b_k_blks * get_repacked_block_type_size<BLOC_TYPE, INTER_SIZE, NB_COLS>();
+        const int64_t per_mb_rows_wsize = row_align * row_stride_a;
+        const int64_t per_nb_cols_wsize = NB_COLS * row_stride_b;
 
-    GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0));
+        const int64_t barrier_idx = static_cast<int64_t>(ith / 2);
 
-    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_0 != 0) {
-        return -1;
-    }
+        GGML_ASSERT(global_spine_env_info.init_barrier != nullptr);
+        GGML_ASSERT(barrier_idx < spine_init_barrier_count);
+        spine_barrier_t * cur_barrier = &global_spine_env_info.init_barrier[barrier_idx];
 
-    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];
+        if (gemm_m == 1) {
+            int task_per_thread = spacemit_kernels::div_round_up(a_k_blks, nth);
+            int a_blk_start     = ith * task_per_thread;
+            int a_blk_end       = std::min(a_blk_start + task_per_thread, (int) a_k_blks);
+            if (a_blk_start < a_blk_end) {
+                quantize_a_row_i8(a_blk_len, feature + a_blk_start * a_blk_len, (a_blk_end - a_blk_start) * a_blk_len,
+                                  quant_a_buffer + a_blk_start * block_stride_a);
+            }
+        } else {
+            int task_per_thread = spacemit_kernels::div_round_up(row_blks, nth);
+            int m_row_blk_start = ith * task_per_thread;
+            int m_row_blk_end   = std::min(m_row_blk_start + task_per_thread, (int) row_blks);
+            for (int m_row_blk = m_row_blk_start; m_row_blk < m_row_blk_end; m_row_blk++) {
+                int m_idx             = m_row_blk * row_align;
+                int rows_tobe_handled = (gemm_m - m_idx) > row_align ? row_align : (gemm_m - m_idx);
+
+                if (rows_tobe_handled == row_align && quantize_a_4row_i8 != nullptr) {
+                    const float * a_row_ptr       = feature + m_idx * gemm_k;
+                    auto *        quant_a_row_ptr = quant_a_buffer + m_idx * row_stride_a;
+                    quantize_a_4row_i8(a_blk_len, a_row_ptr, gemm_k, quant_a_row_ptr);
+                } else {
+                    while (rows_tobe_handled) {
+                        const float * a_row_ptr       = feature + m_idx * gemm_k;
+                        auto *        quant_a_row_ptr = quant_a_buffer + m_idx * row_stride_a;
+                        quantize_a_row_i8(a_blk_len, a_row_ptr, gemm_k, quant_a_row_ptr);
+                        rows_tobe_handled -= 1;
+                        m_idx += 1;
+                    }
+                }
             }
-            *dst++ = make_block_q4_0x16(dst_tmp, interleave_block);
         }
-        src += nrows_interleaved * nblocks;
-    }
-    return 0;
 
-    GGML_UNUSED(data_size);
-}
+        ggml_barrier(params->threadpool);
 
-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);
+        const int64_t gemm_m_stride     = gemm_n / gemm_m > 64 ? gemm_m : 16;
+        const int64_t gemm_m_blocked    = spacemit_kernels::div_round_up(gemm_m, gemm_m_stride);
+        const int64_t max_gemm_n_stride = spacemit_kernels::div_round_up(gemm_n * gemm_m_blocked, nth);
 
-    constexpr int nrows_interleaved = 16;
+        int64_t gemm_n_stride = gemm_n;
+        if (max_gemm_n_stride < gemm_n) {
+            gemm_n_stride =
+                std::min(gemm_n_stride, spacemit_kernels::div_round_up(max_gemm_n_stride, NB_COLS) * NB_COLS);
+        }
 
-    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;
+        if (gemm_n_stride == gemm_n && tcm_buffer != nullptr && per_mb_rows_wsize <= tcm_buffer_size) {
+            for (int64_t m_start = ith * row_align; m_start < gemm_m; m_start += row_align * nth) {
+                uint8_t * b_col    = reinterpret_cast<uint8_t *>(w_data);
+                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
 
-    GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_1));
+                int64_t m_row_real = std::min(gemm_m - m_start, row_align);
 
-    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK4_1 != 0) {
-        return -1;
-    }
+                spacemit_kernels::rvv::memcpy1d(tcm_buffer, quant_a_buffer + m_start * row_stride_a,
+                                                m_row_real * row_stride_a);
 
-    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];
+                int64_t n_blk_real = 0;
+                for (int64_t ni = 0; ni < gemm_n; ni += n_blk_real, b_col += n_blk_real * row_stride_b) {
+                    n_blk_real = std::min(gemm_n - ni, (int64_t) NB_COLS);
+
+                    uint8_t * a_row_ptr = (uint8_t *) tcm_buffer;
+                    float *   c_blk     = output + m_start * gemm_n + ni;
+
+                    int32_t rows_remaining = m_row_real;
+
+                    while (rows_remaining > 0) {
+                        auto rows_handled = gemm_kernel(b_blk_len, a_row_ptr, b_col, b_col_zp, c_blk, rows_remaining,
+                                                        n_blk_real, b_k_blks, gemm_n);
+
+                        c_blk += rows_handled * gemm_n;
+                        a_row_ptr += rows_handled * row_stride_a;
+
+                        rows_remaining -= rows_handled;
+                    }
+                }
             }
-            *dst++ = make_block_q4_1x16(dst_tmp, interleave_block);
-        }
-        src += nrows_interleaved * nblocks;
-    }
-    return 0;
+        } else if (tcm_buffer != nullptr && per_nb_cols_wsize <= tcm_buffer_size) {
+            uint8_t * a_row = quant_a_buffer;
+            uint8_t * b_col = reinterpret_cast<uint8_t *>(tcm_buffer);
+            if ((gemm_workspace_size + per_nb_cols_wsize) <= tcm_buffer_size) {
+                a_row = (uint8_t *) tcm_buffer;
+                b_col = reinterpret_cast<uint8_t *>(tcm_buffer) + gemm_workspace_size;
+            }
+            uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
 
-    GGML_UNUSED(data_size);
-}
+            int64_t ni      = ith * NB_COLS;
+            int64_t nb_real = std::min(gemm_n - ni, NB_COLS);
 
-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);
-    }
-}
+            if (ith % 2 == 0 && nb_real > 0) {
+                spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(w_data) + ni * row_stride_b,
+                                                nb_real * row_stride_b);
+                if (a_row != quant_a_buffer) {
+                    spacemit_kernels::rvv::memcpy1d(a_row, quant_a_buffer, gemm_workspace_size);
+                }
+            }
 
-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);
+            spine_barrier_wait(cur_barrier);
 
-    constexpr int nrows_interleaved = 16;
+            if (ith % 2 != 0 && nb_real > 0) {
+                if (a_row != quant_a_buffer) {
+                    spacemit_kernels::rvv::memcpy1d(a_row, quant_a_buffer, gemm_workspace_size);
+                }
+                spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(w_data) + ni * row_stride_b,
+                                                nb_real * row_stride_b);
+            }
 
-    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;
+            for (; ni < gemm_n; ni += NB_COLS * nth) {
+                int64_t rows_remaining = gemm_m;
+                float * c_blk          = output + ni;
+                auto *  a_row_cur      = a_row;
 
-    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK_K != 0) {
-        return -1;
-    }
+                if (ith % 2 != 0) {
+                    spine_barrier_wait(cur_barrier);
+                }
 
-    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);
-                        }
-                    }
+                while (rows_remaining > 0) {
+                    auto rows_handled = gemm_kernel(b_blk_len, a_row_cur, b_col, b_col_zp, c_blk, rows_remaining,
+                                                    nb_real, b_k_blks, gemm_n);
+
+                    c_blk += rows_handled * gemm_n;
+                    a_row_cur += rows_handled * row_stride_a;
+
+                    rows_remaining -= rows_handled;
+                }
+
+                if (ith % 2 == 0) {
+                    spine_barrier_wait(cur_barrier);
+                }
+
+                const int64_t next_ni = ni + NB_COLS * nth;
+                if (next_ni < gemm_n) {
+                    nb_real = std::min(gemm_n - next_ni, NB_COLS);
+                    spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(w_data) + next_ni * row_stride_b,
+                                                    nb_real * row_stride_b);
                 }
-                *dst++ = make_block_q4_1x16(dst_tmp, interleave_block);
             }
-        }
-        src += nrows_interleaved * nblocks;
-    }
-    return 0;
+        } else {
+            const int64_t task_count_m = spacemit_kernels::div_round_up(gemm_m, gemm_m_stride);
+            const int64_t task_count_n = spacemit_kernels::div_round_up(gemm_n, gemm_n_stride);
 
-    GGML_UNUSED(data_size);
-}
+            int64_t task_count      = task_count_m * task_count_n;
+            int64_t task_per_thread = (task_count + nth - 1) / nth;
+            int64_t start           = ith * task_per_thread;
+            int64_t end             = std::min((ith + 1) * task_per_thread, task_count);
+            for (int64_t compute_idx = start; compute_idx < end; compute_idx++) {
+                const auto tid_n = compute_idx / task_count_m;
+                const auto tid_m = compute_idx % task_count_m;
 
-namespace ggml::cpu::riscv64_spacemit {
+                const int64_t m_start = tid_m * gemm_m_stride;
+                const int64_t m_count = std::min(gemm_m - m_start, (int64_t) gemm_m_stride);
 
-template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS>
-int repack(struct ggml_tensor *, const void *, size_t);
+                const int64_t n_start = tid_n * gemm_n_stride;
+                const int64_t n_count = std::min(gemm_n - n_start, (int64_t) gemm_n_stride);
 
-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);
-}
+                const int64_t n_blk = m_count == 1 ? n_count : NB_COLS;
 
-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);
-}
+                uint8_t * b_col    = reinterpret_cast<uint8_t *>(w_data) + n_start * row_stride_b;
+                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
 
-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);
-}
+                int64_t n_blk_real = 0;
+                for (int64_t ni = 0; ni < n_count; ni += n_blk_real, b_col += n_blk_real * row_stride_b) {
+                    n_blk_real = std::min(n_count - ni, n_blk);
 
-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;
-};
+                    uint8_t * a_row = quant_a_buffer + m_start * row_stride_a;
 
-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;
-    }
+                    float * c_blk = output + m_start * gemm_n + n_start + ni;
 
-    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;
+                    int64_t rows_remaining = m_count;
+
+                    uint8_t * b_col_cur    = b_col;
+                    uint8_t * b_col_zp_cur = b_col_zp;
+
+                    while (rows_remaining > 0) {
+                        auto rows_handled = gemm_kernel(b_blk_len, a_row, b_col_cur, b_col_zp_cur, c_blk,
+                                                        rows_remaining, n_blk_real, b_k_blks, gemm_n);
+
+                        c_blk += rows_handled * gemm_n;
+                        a_row += rows_handled * row_stride_a;
+
+                        rows_remaining -= rows_handled;
+                    }
                 }
-            default:
-                // GGML_ABORT("fatal error");
-                break;
+            }
         }
-        return false;
     }
 
-    void forward_mul_mat_q4(ggml_compute_params * params, ggml_tensor * op) {
+    void forward_mul_mat_id(ggml_compute_params * params, ggml_tensor * op) {
+        constexpr size_t a_blk_len = INTER_SIZE;
+        constexpr size_t b_blk_len = INTER_SIZE;
+
         const ggml_tensor * src0 = op->src[0];
         const ggml_tensor * src1 = op->src[1];
+        const ggml_tensor * ids  = op->src[2];
         ggml_tensor *       dst  = op;
 
         GGML_TENSOR_BINARY_OP_LOCALS
@@ -429,133 +554,381 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
         int ith = params->ith;
         int nth = params->nth;
 
-        [[maybe_unused]] const enum ggml_type type = src0->type;
+        // row groups
+        const int n_ids = ids->ne[0];  // n_expert_used
+        const int n_as  = ne02;        // n_expert
+
+        struct mmid_row_mapping {
+            int32_t i1;
+            int32_t i2;
+        };
+
+        spacemit_kernels::quantize_a_row_def           quantize_a_row_i8;
+        spacemit_kernels::gemm_kernel_quantize_def     gemm_kernel;
+        spacemit_kernels::moe_gemm_kernel_quantize_def moe_gemm_kernel_m2;
+        bool                                           set_kernel_impl = false;
+        size_t                                         block_stride_a  = spacemit_kernels::q8_blk_size(QK4_0);
+
+#if defined(RISCV64_SPACEMIT_IME2)
+        if (!set_kernel_impl && (global_spine_env_info.use_ime2)) {
+            quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8;
+            block_stride_a    = spacemit_kernels::q8_blk_size(QK4_0, true);
+
+            if constexpr (std::is_same_v<BLOC_TYPE, block_q6_K> || std::is_same_v<BLOC_TYPE, block_q8_0>) {
+                gemm_kernel     = spacemit_kernels::ime2::gemm_kernel_i8i8;
+                set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
+                                 std::is_same_v<BLOC_TYPE, block_q4_K>) {
+                if constexpr (INTER_SIZE == 256) {
+                    gemm_kernel       = spacemit_kernels::ime2::gemm_kernel_i8i4_hp;
+                    quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8_hp;
+                    block_stride_a    = spacemit_kernels::q8_hp_blk_size(a_blk_len, true, true);
+                    set_kernel_impl   = true;
+                } else {
+                    gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i4;
+                    moe_gemm_kernel_m2 = spacemit_kernels::ime2::moe_m2_gemm_kernel_i8i4;
+                    quantize_a_row_i8  = spacemit_kernels::rvv::quantize_a_row_i8;
+                    block_stride_a     = spacemit_kernels::q8_blk_size(a_blk_len, true);
+                    set_kernel_impl    = true;
+                }
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q2_K>) {
+                quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8k;
+                block_stride_a    = spacemit_kernels::q8k_blk_size(a_blk_len);
+                gemm_kernel       = spacemit_kernels::ime2::gemm_kernel_i8i2k;
+                set_kernel_impl   = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q3_K>) {
+                quantize_a_row_i8 = spacemit_kernels::rvv::quantize_a_row_i8k;
+                block_stride_a    = spacemit_kernels::q8k_blk_size(a_blk_len);
+                gemm_kernel       = spacemit_kernels::ime2::gemm_kernel_i8i3k;
+                set_kernel_impl   = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_mxfp4>) {
+                gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8mxfp4;
+                moe_gemm_kernel_m2 = spacemit_kernels::ime2::moe_m2_gemm_kernel_i8mxfp4;
+                set_kernel_impl    = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q5_1> || std::is_same_v<BLOC_TYPE, block_q5_K> ||
+                                 std::is_same_v<BLOC_TYPE, block_q5_0>) {
+                gemm_kernel        = spacemit_kernels::ime2::gemm_kernel_i8i5;
+                moe_gemm_kernel_m2 = spacemit_kernels::ime2::moe_m2_gemm_kernel_i8i5;
+                set_kernel_impl    = true;
+            }
+        }
+#endif
 
-        void *        w_data  = (void *) src0->data;
-        const float * feature = (const float *) src1->data;
-        float *       output  = (float *) dst->data;
+#if defined(RISCV64_SPACEMIT_IME1)
+        if (!set_kernel_impl && (global_spine_env_info.use_ime1)) {
+            quantize_a_row_i8 = spacemit_kernels::ime1::quantize_a_row_i8;
+
+            if constexpr (std::is_same_v<BLOC_TYPE, block_q4_0> || std::is_same_v<BLOC_TYPE, block_q4_1> ||
+                          std::is_same_v<BLOC_TYPE, block_q4_K>) {
+                gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i4;
+                set_kernel_impl = true;
+            }
+        }
+#endif
+        if (!set_kernel_impl) {
+            GGML_ABORT("no kernel implementation found for the block type");
+        }
 
-        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;
+        const size_t a_k_blks = spacemit_kernels::div_round_up(ne10, a_blk_len);
+        const size_t b_k_blks = spacemit_kernels::div_round_up(ne10, b_blk_len);
 
-        GGML_ASSERT(batch_weight == 1);
+        const size_t nbw1                = a_k_blks * block_stride_a;
+        const size_t nbw2                = ne11 * nbw1;
+        const size_t nbw3                = nbw2 * ne12;
+        const size_t gemm_workspace_size = GGML_PAD(nbw3, alignof(int64_t));
 
-        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;
+        const uintptr_t ws_ptr         = reinterpret_cast<uintptr_t>(params->wdata);
+        auto *          quant_a_buffer = reinterpret_cast<uint8_t *>(ws_ptr);
 
-        if (ith == 0 && params->wsize < desired_wsize) {
-            throw std::runtime_error("wsize less than desired_wsize");
+        if (ne11 == 1) {
+            for (int64_t ii = ith; ii < ne12 * a_k_blks; ii += nth) {
+                int64_t i12       = ii / a_k_blks;
+                int64_t ak_blk_id = ii % a_k_blks;
+                quantize_a_row_i8(a_blk_len, (float *) ((char *) src1->data + i12 * nb12) + ak_blk_id * a_blk_len,
+                                  a_blk_len, quant_a_buffer + i12 * nbw2 + ak_blk_id * block_stride_a);
+            }
+        } else {
+            for (int64_t ii = ith; ii < ne12 * ne11; ii += nth) {
+                int64_t i12 = ii / ne11;
+                int64_t i11 = ii % ne11;
+                quantize_a_row_i8(a_blk_len, (float *) ((char *) src1->data + i12 * nb12 + i11 * nb11), ne10,
+                                  quant_a_buffer + i12 * nbw2 + i11 * nbw1);
+            }
         }
 
-        std::vector<qnbitgemm_spacemit_ime_args> qnbitgemm_args(batch_feature);
+#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) *ne12 + (i1)]
 
-        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;
+        int64_t *          matrix_row_counts       = (int64_t *) (ws_ptr + gemm_workspace_size);
+        int32_t *          valid_ep_count          = (int32_t *) (matrix_row_counts + n_as);
+        int32_t *          valid_act_count         = (int32_t *) (valid_ep_count + 1);
+        int64_t *          valid_matrix_row_counts = (int64_t *) (valid_act_count + 1);
+        mmid_row_mapping * matrix_rows             = (mmid_row_mapping *) (valid_matrix_row_counts + n_as);
 
-            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;
+        if (ith == 0) {
+            // initialize matrix_row_counts
+            memset(matrix_row_counts, 0, n_as * sizeof(int64_t));
+
+            // group rows by src0 matrix
+            for (int32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+                for (int32_t id = 0; id < n_ids; ++id) {
+                    const int32_t i02 =
+                        *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
+
+                    GGML_ASSERT(i02 >= 0 && i02 < n_as);
+
+                    MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
+                    matrix_row_counts[i02] += 1;
+                }
             }
 
-            qnbitgemm_args[i].bias  = nullptr;
-            qnbitgemm_args[i].c_ptr = output + gemm_m * gemm_n * i;
-            qnbitgemm_args[i].ldc   = gemm_n;
+            int32_t valid_ep_count_t  = 0;
+            int32_t valid_act_count_t = 0;
+            for (int cur_a = 0; cur_a < n_as; ++cur_a) {
+                const int64_t cne1 = matrix_row_counts[cur_a];
+                if (cne1 == 0) {
+                    continue;
+                }
+                valid_matrix_row_counts[valid_ep_count_t] = cur_a;
+                valid_act_count_t += cne1;
+                valid_ep_count_t += 1;
+            }
+            valid_ep_count[0]  = valid_ep_count_t;
+            valid_act_count[0] = valid_act_count_t;
         }
 
-        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);
+        const int64_t barrier_idx = static_cast<int64_t>(ith / 2);
 
-        {
-            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_ASSERT(global_spine_env_info.init_barrier != nullptr);
+        GGML_ASSERT(barrier_idx < spine_init_barrier_count);
+        spine_barrier_t * cur_barrier = &global_spine_env_info.init_barrier[barrier_idx];
+
+        ggml_barrier(params->threadpool);
+
+        const size_t row_stride_b      = b_k_blks * get_repacked_block_type_size<BLOC_TYPE, INTER_SIZE, NB_COLS>();
+        const size_t expert_b_stride   = ne01 * row_stride_b;
+        const size_t per_nb_cols_wsize = NB_COLS * row_stride_b;
+
+        std::array<const uint8_t *, 2> src_workspaces;
+        std::array<float *, 2>         dst_workspaces;
+
+        auto *     tcm_buffer      = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer;
+        const auto tcm_buffer_size = ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size;
+
+        const auto valid_ep_count_t  = valid_ep_count[0];
+        const auto valid_act_count_t = valid_act_count[0];
+
+        int nth_es = 1;
+        int nth_n  = nth;
+
+        int ith_es = ith % nth_es;
+        int ith_n  = (ith / nth_es) % nth_n;
+
+        if (valid_ep_count_t % nth == 0 && tcm_buffer != nullptr && valid_ep_count_t == n_as &&
+            valid_act_count_t == n_as && per_nb_cols_wsize <= tcm_buffer_size) {
+            for (int64_t valid_id = ith; valid_id < valid_ep_count_t; valid_id += nth) {
+                const int64_t cur_a = valid_matrix_row_counts[valid_id];
+
+                auto * src0_cur = (uint8_t *) src0->data + cur_a * expert_b_stride;
+
+                mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, 0);
+                const int        id          = row_mapping.i1;
+                const int64_t    i11         = id % ne11;
+                const int64_t    i12         = row_mapping.i2;
+                const int64_t    i1          = id;
+                const int64_t    i2          = i12;
+
+                auto *  src1_col = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
+                float * c_blk    = (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2));
+
+                uint8_t * a_row = src1_col;
+                uint8_t * b_col = reinterpret_cast<uint8_t *>(tcm_buffer);
+                if ((nbw1 + per_nb_cols_wsize) <= tcm_buffer_size) {
+                    a_row = (uint8_t *) tcm_buffer;
+                    b_col = reinterpret_cast<uint8_t *>(tcm_buffer) + nbw1;
+                }
+                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? b_col : nullptr;
+
+                if (ith % 2 == 0) {
+                    spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(src0_cur), per_nb_cols_wsize);
+
+                    if (a_row != src1_col) {
+                        spacemit_kernels::rvv::memcpy1d(a_row, src1_col, nbw1);
+                    }
+                }
+
+                spine_barrier_wait(cur_barrier);
+
+                if (ith % 2 != 0) {
+                    if (a_row != src1_col) {
+                        spacemit_kernels::rvv::memcpy1d(a_row, src1_col, nbw1);
+                    }
+
+                    spacemit_kernels::rvv::memcpy1d(b_col, reinterpret_cast<uint8_t *>(src0_cur), per_nb_cols_wsize);
+                }
+
+                int64_t nb_real = std::min(ne01, NB_COLS);
+                for (int64_t ni = 0; ni < ne01; ni += NB_COLS) {
+                    if (ith % 2 != 0) {
+                        spine_barrier_wait(cur_barrier);
+                    }
+
+                    gemm_kernel(b_blk_len, a_row, b_col, b_col_zp, c_blk + ni, 1, nb_real, b_k_blks, ne01);
+
+                    if (ith % 2 == 0) {
+                        spine_barrier_wait(cur_barrier);
+                    }
+
+                    const int64_t next_ni = ni + NB_COLS;
+                    if (next_ni < ne01) {
+                        nb_real = std::min(ne01 - next_ni, NB_COLS);
+                        spacemit_kernels::rvv::memcpy1d(
+                            b_col, reinterpret_cast<uint8_t *>(src0_cur) + next_ni * row_stride_b, per_nb_cols_wsize);
                     }
                 }
             }
-        }
+        } else {
+            for (int64_t valid_id = ith_es; valid_id < valid_ep_count_t; valid_id += nth_es) {
+                const int64_t cur_a = valid_matrix_row_counts[valid_id];
+                const int64_t cne1  = matrix_row_counts[cur_a];
 
-        ggml_barrier(params->threadpool);
+                int64_t src1_cur_start = 0;
+                int64_t src1_cur_end   = cne1;
 
-        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;
+                int64_t src0_cur_start = (ith_n * ne01) / nth_n;
+                int64_t src0_cur_end   = MIN(((ith_n + 1) * ne01) / nth_n, ne01);
 
-        {
-            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];
+                if (src1_cur_start >= src1_cur_end || src0_cur_start >= src0_cur_end) {
+                    continue;
+                }
+
+                src0_cur_start =
+                    (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start;
+                src0_cur_end =
+                    (src0_cur_end % NB_COLS) ? src0_cur_end + NB_COLS - (src0_cur_end % NB_COLS) : src0_cur_end;
+
+                auto *    src0_cur = (uint8_t *) src0->data + cur_a * expert_b_stride + src0_cur_start * row_stride_b;
+                uint8_t * b_col_zp = block_type_has_zp<BLOC_TYPE>() ? src0_cur : nullptr;
+
+                size_t extra_tcm_buffer_size = tcm_buffer_size;
+                void * extra_tcm_buffer      = tcm_buffer;
+                if (tcm_buffer != nullptr && (src1_cur_end - src1_cur_start) >= 4 &&
+                    (src0_cur_end - src0_cur_start) * row_stride_b <= tcm_buffer_size) {
+                    spacemit_kernels::rvv::memcpy1d(tcm_buffer, src0_cur,
+                                                    (src0_cur_end - src0_cur_start) * row_stride_b);
+                    src0_cur = reinterpret_cast<uint8_t *>(tcm_buffer);
+                    b_col_zp = block_type_has_zp<BLOC_TYPE>() ? src0_cur : nullptr;
+                    extra_tcm_buffer_size -= (src0_cur_end - src0_cur_start) * row_stride_b;
+                    extra_tcm_buffer = reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(tcm_buffer) +
+                                                                (src0_cur_end - src0_cur_start) * row_stride_b);
+                }
 
-                const auto tid_n = blk_i / thread_count_m;
-                const auto tid_m = blk_i % thread_count_m;
+                int ir1 = src1_cur_start;
 
-                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);
+                if (extra_tcm_buffer_size >= nbw1 && extra_tcm_buffer != nullptr) {
+                    int64_t quant_a_tile_size = extra_tcm_buffer_size / nbw1;
+                    do {
+                        quant_a_tile_size = MIN(quant_a_tile_size, src1_cur_end - ir1);
 
-                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);
+                        uint8_t * quant_a_tile_buffer = reinterpret_cast<uint8_t *>(extra_tcm_buffer);
 
-                void * per_gemm_ws = reinterpret_cast<std::byte *>(ws) + gemm_i * per_gemm_workspace_stride;
+                        int iir1 = ir1;
+                        for (; iir1 < (ir1 + quant_a_tile_size); ++iir1) {
+                            mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, iir1);
 
-                sqnbitgemm_spacemit_ime_i8i4(QK4_0, gemm_k, data, per_gemm_ws, m_start, m_count, n_start, n_count);
+                            const int id = row_mapping.i1;  // selected expert index
+
+                            const int64_t i11 = id % ne11;
+                            const int64_t i12 = row_mapping.i2;  // row index in src1
+
+                            auto * src1_col = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
+                            spacemit_kernels::rvv::memcpy1d(quant_a_tile_buffer, src1_col, nbw1);
+                            quant_a_tile_buffer = quant_a_tile_buffer + nbw1;
+                        }
+
+                        quant_a_tile_buffer = reinterpret_cast<uint8_t *>(extra_tcm_buffer);
+                        iir1                = ir1;
+
+                        if (moe_gemm_kernel_m2 != nullptr) {
+                            for (; iir1 < (ir1 + quant_a_tile_size - 1); iir1 += 2, quant_a_tile_buffer += 2 * nbw1) {
+                                mmid_row_mapping row_mapping_0 = MMID_MATRIX_ROW(cur_a, iir1);
+                                mmid_row_mapping row_mapping_1 = MMID_MATRIX_ROW(cur_a, iir1 + 1);
+
+                                src_workspaces[0] = quant_a_tile_buffer;
+                                src_workspaces[1] = quant_a_tile_buffer + nbw1;
+
+                                dst_workspaces[0] =
+                                    (float *) ((char *) dst->data + (row_mapping_0.i1 * nb1 + row_mapping_0.i2 * nb2)) +
+                                    src0_cur_start;
+                                dst_workspaces[1] = (float *) ((char *) dst->data +
+                                                               ((row_mapping_1.i1) * nb1 + (row_mapping_1.i2) * nb2)) +
+                                                    src0_cur_start;
+                                moe_gemm_kernel_m2(b_blk_len, src_workspaces.data(), src0_cur, b_col_zp,
+                                                   dst_workspaces.data(), 1, src0_cur_end - src0_cur_start, b_k_blks,
+                                                   ne01);
+                            }
+                        }
+
+                        for (; iir1 < (ir1 + quant_a_tile_size); iir1++, quant_a_tile_buffer += nbw1) {
+                            mmid_row_mapping row_mapping_0 = MMID_MATRIX_ROW(cur_a, iir1);
+
+                            gemm_kernel(
+                                b_blk_len, quant_a_tile_buffer, src0_cur, b_col_zp,
+                                (float *) ((char *) dst->data + (row_mapping_0.i1 * nb1 + row_mapping_0.i2 * nb2)) +
+                                    src0_cur_start,
+                                1, src0_cur_end - src0_cur_start, b_k_blks, ne01);
+                        }
+
+                        ir1 += quant_a_tile_size;
+                    } while (ir1 < src1_cur_end);
+                } else {
+                    if (moe_gemm_kernel_m2 != nullptr) {
+                        for (; ir1 < src1_cur_end - 1; ir1 += 2) {
+                            for (int iir1 = 0; iir1 < 2; ++iir1) {
+                                mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1 + iir1);
+
+                                const int id = row_mapping.i1;  // selected expert index
+
+                                const int64_t i11 = id % ne11;
+                                const int64_t i12 = row_mapping.i2;  // row index in src1
+
+                                const int64_t i1 = id;               // selected expert index
+                                const int64_t i2 = i12;              // row
+
+                                src_workspaces[iir1] = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
+
+                                dst_workspaces[iir1] =
+                                    (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start;
+                            }
+
+                            moe_gemm_kernel_m2(b_blk_len, src_workspaces.data(), src0_cur, b_col_zp,
+                                               dst_workspaces.data(), 1, src0_cur_end - src0_cur_start, b_k_blks, ne01);
+                        }
+                    }
+
+                    for (; ir1 < src1_cur_end; ir1++) {
+                        mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1);
+
+                        const int id = row_mapping.i1;  // selected expert index
+
+                        const int64_t i11 = id % ne11;
+                        const int64_t i12 = row_mapping.i2;  // row index in src1
+
+                        const int64_t i1 = id;               // selected expert index
+                        const int64_t i2 = i12;              // row
+
+                        auto * src1_col = quant_a_buffer + (i11 * nbw1 + i12 * nbw2);
+
+                        gemm_kernel(b_blk_len, src1_col, src0_cur, b_col_zp,
+                                    (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start, 1,
+                                    src0_cur_end - src0_cur_start, b_k_blks, ne01);
+                    }
+                }
             }
         }
+#undef MMID_MATRIX_ROW
     }
 
-    int repack(struct ggml_tensor * t, const void * data, size_t data_size) override {
+    int repack(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);
@@ -563,309 +936,464 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
 };
 
 class tensor_traits_common : public tensor_traits_base {
-    bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
+    bool work_size(int n_threads, const ggml_tensor * op, size_t & size) override {
         switch (op->op) {
-            case GGML_OP_NORM:
-            case GGML_OP_RMS_NORM:
-                size = 0;
+            case GGML_OP_FLASH_ATTN_EXT:
+                {
+                    const int     n_tasks = n_threads;
+                    const int64_t neq2    = op->src[0]->ne[2];  // number of query heads
+                    const int64_t DK      = op->src[1]->ne[0];
+                    const int64_t DV      = op->src[2]->ne[0];  // DV
+
+                    // Tiled flash attention scratch (tile sizes defined in common.h)
+                    // Per-thread: Q_q + KQ + mask + VKQ32 + V32 + K_f32 + padding
+                    size_t prefill = sizeof(float) *
+                                     (GGML_FA_TILE_Q * DK + 2 * GGML_FA_TILE_Q * GGML_FA_TILE_KV + GGML_FA_TILE_Q * DV +
+                                      GGML_FA_TILE_KV * DV + GGML_FA_TILE_KV * DK) *
+                                     n_tasks;
+
+                    // Decode path: n_kv_chunks = n_tasks (one chunk per thread)
+                    // Per-thread: VKQ accmulator (DV), partial M, partial S + intra-thread scratch for V, Q and VKQ
+                    size_t n_chunks = n_tasks;
+                    size_t decode   = sizeof(float) * (neq2 * n_chunks * (2 + DV) + n_tasks * (DK + 2 * DV));
+
+                    size = MAX(prefill, decode);
+                }
                 return true;
             default:
-                // GGML_ABORT("fatal error");
                 break;
         }
         return false;
     }
 
-    bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override {
+    bool compute_forward(ggml_compute_params * params, ggml_tensor * op) override {
         switch (op->op) {
             case GGML_OP_NORM:
-                forward_norm_f32(params, op);
-                return true;
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_norm_f32(params, op);
+                        return true;
+                    default:
+                        GGML_ABORT("fatal error");
+                }
             case GGML_OP_RMS_NORM:
-                forward_rms_norm_f32(params, op);
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_rms_norm_f32(params, op);
+                        return true;
+                    default:
+                        GGML_ABORT("fatal error");
+                }
+            case GGML_OP_ADD:
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_ADD, float>(params, op);
+                        return true;
+                    case GGML_TYPE_F16:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_ADD, _Float16>(params, op);
+                        return true;
+                    default:
+                        ggml_compute_forward_add(params, op);
+                        return true;
+                }
+            case GGML_OP_SUB:
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_SUB, float>(params, op);
+                        return true;
+                    case GGML_TYPE_F16:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_SUB, _Float16>(params, op);
+                        return true;
+                    default:
+                        ggml_compute_forward_sub(params, op);
+                        return true;
+                }
+            case GGML_OP_MUL:
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_MUL, float>(params, op);
+                        return true;
+                    case GGML_TYPE_F16:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_MUL, _Float16>(params, op);
+                        return true;
+                    default:
+                        ggml_compute_forward_mul(params, op);
+                        return true;
+                }
+            case GGML_OP_DIV:
+                switch (op->src[0]->type) {
+                    case GGML_TYPE_F32:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_DIV, float>(params, op);
+                        return true;
+                    case GGML_TYPE_F16:
+                        spacemit_kernels::rvv::forward_binary<GGML_OP_DIV, _Float16>(params, op);
+                        return true;
+                    default:
+                        ggml_compute_forward_div(params, op);
+                        return true;
+                }
+            case GGML_OP_FLASH_ATTN_EXT:
+                forward_flash_attn_ext_f16(params, op);
+                return true;
+            case GGML_OP_CONT:
+                {
+                    const ggml_tensor * src0 = op->src[0];
+                    if (op->type == src0->type && op->nb[0] != src0->nb[0] && op->nb[0] == src0->nb[1] &&
+                        op->ne[3] * op->ne[2] * op->nb[2] == src0->ne[3] * src0->ne[2] * src0->nb[2]) {
+                        spacemit_kernels::rvv::forward_cont_with_permute(params, op);
+                    } else {
+                        ggml_compute_forward_cont(params, op);
+                    }
+                    return true;
+                }
+            case GGML_OP_CPY:
+                {
+                    const ggml_tensor * src0 = op->src[0];
+                    if (op->type == src0->type && op->nb[0] == src0->nb[1] && src0->nb[0] != src0->nb[1] &&
+                        ggml_nelements(src0) == ggml_nelements(op)) {
+                        spacemit_kernels::rvv::forward_cpy_with_permute(params, op);
+                    } else {
+                        ggml_compute_forward_cpy(params, op);
+                    }
+                    return true;
+                }
+            case GGML_OP_REPEAT:
+                {
+                    const bool rows_equal         = ggml_nrows(op->src[0]) == ggml_nrows(op);
+                    const bool broadcast_or_equal = op->src[0]->ne[0] == 1 || op->src[0]->ne[0] == op->ne[0];
+
+                    if (rows_equal && broadcast_or_equal) {
+                        switch (op->src[0]->type) {
+                            case GGML_TYPE_F32:
+                                spacemit_kernels::rvv::forward_repeat_nrows<int32_t>(params, op);
+                                return true;
+                            case GGML_TYPE_F16:
+                                spacemit_kernels::rvv::forward_repeat_nrows<int16_t>(params, op);
+                                return true;
+                            default:
+                                break;
+                        }
+                    }
+
+                    if (op->src[0]->ne[1] == 1 && op->src[0]->ne[0] == op->ne[0]) {
+                        switch (op->src[0]->type) {
+                            case GGML_TYPE_F32:
+                                spacemit_kernels::rvv::forward_repeat_dim1<int32_t>(params, op);
+                                return true;
+                            case GGML_TYPE_F16:
+                                spacemit_kernels::rvv::forward_repeat_dim1<int16_t>(params, op);
+                                return true;
+                            default:
+                                break;
+                        }
+                    }
+
+                    ggml_compute_forward_repeat(params, op);
+                }
+                return true;
+            case GGML_OP_SUM_ROWS:
+                {
+                    if (op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) {
+                        spacemit_kernels::rvv::forward_sum_rows<float>(params, op);
+                    } else {
+                        ggml_compute_forward_sum_rows(params, op);
+                    }
+                }
+                return true;
+            case GGML_OP_GET_ROWS:
+                {
+                    if (op->src[0]->type == op->type) {
+                        switch (op->src[0]->type) {
+                            case GGML_TYPE_F32:
+                                spacemit_kernels::rvv::forward_get_rows<int32_t>(params, op);
+                                return true;
+                            case GGML_TYPE_F16:
+                                spacemit_kernels::rvv::forward_get_rows<int16_t>(params, op);
+                                return true;
+                            default:
+                                break;
+                        }
+                    }
+
+                    ggml_compute_forward_get_rows(params, op);
+                }
                 return true;
+            case GGML_OP_CONCAT:
+                {
+                    const int32_t dim = ggml_get_op_params_i32(op, 0);
+                    if (dim == 0 && op->type == op->src[0]->type) {
+                        switch (op->src[0]->type) {
+                            case GGML_TYPE_F32:
+                                spacemit_kernels::rvv::forward_concat<int32_t>(params, op);
+                                return true;
+                            case GGML_TYPE_F16:
+                                spacemit_kernels::rvv::forward_concat<int16_t>(params, op);
+                                return true;
+                            default:
+                                break;
+                        }
+                    }
+
+                    ggml_compute_forward_concat(params, op);
+                }
+                return true;
+            // TODO For GGML_OP_GATED_DELTA_NET
+            // case GGML_OP_GATED_DELTA_NET:
+            //     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));
+    void forward_flash_attn_ext_f16(const ggml_compute_params * params, ggml_tensor * dst) {
+        const ggml_tensor * q = dst->src[0];
+        const ggml_tensor * k = dst->src[1];
+        const ggml_tensor * v = dst->src[2];
+
+        GGML_TENSOR_LOCALS(int64_t, neq, q, ne)
+        GGML_TENSOR_LOCALS(size_t, nbq, q, nb)
+        GGML_TENSOR_LOCALS(int64_t, nek, k, ne)
+        GGML_TENSOR_LOCALS(size_t, nbk, k, nb)
+        GGML_TENSOR_LOCALS(int64_t, nev, v, ne)
+        GGML_TENSOR_LOCALS(size_t, nbv, v, nb)
+        GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
+        GGML_TENSOR_LOCALS(size_t, nb, dst, nb)
+
+        const int64_t DK = nek0;
+        const int64_t DV = nev0;
+
+        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_DEFAULT);
+        const bool supported_types = (q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
+        const bool supported_shape = (DK > 0 && DK <= 128 && DV > 0 && DV <= 128);
+        const bool supported_vlen  = (__riscv_vlenb() == 128);
+
+        if (!(supported_prec && supported_types && supported_shape && supported_vlen)) {
+            ggml_compute_forward_flash_attn_ext(params, dst);
+            return;
+        }
+
+        // total rows in q
+        const int64_t nr = neq1 * neq2 * neq3;
 
+        // rows per thread
         const int ith = params->ith;
         const int nth = params->nth;
 
-        GGML_TENSOR_UNARY_OP_LOCALS
+        static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q;
+        const bool               use_tiled = !params->use_ref && (neq1 >= Q_TILE_SZ);
 
-        float epsilon;
-        memcpy(&epsilon, dst->op_params, sizeof(float));
+        // 4x chunks per thread
+        // int     nth_scaled = nth * 4;
+        // int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled;
+        // int64_t nchunk     = (nr + chunk_size - 1) / chunk_size;
 
-        GGML_ASSERT(epsilon > 0.0f);
+        // if (nth == 1 || nchunk < nth) {
+        //     nchunk = nth;
+        // }
 
-        auto * input  = (float *) src0->data;
-        auto * output = (float *) dst->data;
+        int64_t nchunk = nth;
 
-        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);
+        if (ith == 0) {
+            // Every thread starts at ith, so the first unprocessed chunk is nth.  This save a bit of coordination right at the start.
+            ggml_threadpool_chunk_set(params->threadpool, nth);
+        }
 
-        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);
+        ggml_barrier(params->threadpool);
 
-            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);
+        // The number of elements in each chunk
+        const int64_t dr = (nr + nchunk - 1) / nchunk;
 
-                sum    = __riscv_vfadd_vv_f32m4(sum, src_data, gvl);
-                sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl);
+        // The first chunk comes from our thread_id, the rest will get auto-assigned.
+        int current_chunk = ith;
 
-                __riscv_vse32_v_f32m4(p_temp_output, src_data, gvl);
+        while (current_chunk < nchunk) {
+            const int64_t ir0 = dr * current_chunk;
+            const int64_t ir1 = MIN(ir0 + dr, nr);
 
-                p_input += gvl;
-                p_temp_output += gvl;
-                length -= gvl;
+            if (use_tiled) {
+                spacemit_kernels::rvv::forward_flash_attn_ext_f16_tiled_vlen1024_vf16(
+                    params, dst, ir0, ir1, ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer,
+                    ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size);
+            } else {
+                spacemit_kernels::rvv::forward_flash_attn_ext_f16_one_chunk_vlen1024_vf16(
+                    params, dst, ir0, ir1, ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer,
+                    ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size);
             }
 
-            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;
-                }
-            }
+            current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);
         }
     }
 
-    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);
+    int repack(ggml_tensor * t, const void * data, size_t data_size) override {
+        memcpy(t->data, data, data_size);
+        return 0;
+    }
+};
 
-                sum_sq = __riscv_vfmacc_vv_f32m4(sum_sq, src_data, src_data, gvl);
+// Impl By IME1
+static const tensor_traits<block_q4_0, 32, 16>  q4_0_16x32_q8_0;
+static const tensor_traits<block_q4_1, 32, 16>  q4_1_16x32_q8_0;
+static const tensor_traits<block_q4_K, 32, 16>  q4_k_16x32_q8_0;
+// Impl By IME2
+static const tensor_traits<block_q2_K, 256, 32> q2_k_32x256_q8_0;
+static const tensor_traits<block_q3_K, 256, 32> q3_k_32x256_q8_0;
+static const tensor_traits<block_q4_0, 32, 32>  q4_0_32x32_q8_0;
+static const tensor_traits<block_q4_1, 32, 32>  q4_1_32x32_q8_0;
+static const tensor_traits<block_q4_0, 256, 32> q4_0_32x256_q8_0;
+static const tensor_traits<block_q4_1, 256, 32> q4_1_32x256_q8_0;
+static const tensor_traits<block_q4_K, 32, 32>  q4_k_32x32_q8_0;
+static const tensor_traits<block_q6_K, 32, 32>  q6_k_32x32_q8_0;
+static const tensor_traits<block_q8_0, 32, 32>  q8_0_32x32_q8_0;
+static const tensor_traits<block_mxfp4, 32, 32> mxfp4_32x32_q8_0;
+static const tensor_traits<block_q5_K, 32, 32>  q5_k_32x32_q8_0;
+static const tensor_traits<block_q5_1, 32, 32>  q5_1_32x32_q8_0;
+static const tensor_traits<block_q5_0, 32, 32>  q5_0_32x32_q8_0;
+// Impl By RVV
+static const tensor_traits_common               rvv_impl;
 
-                __riscv_vse32_v_f32m4(p_temp_output, src_data, gvl);
+}  // namespace ggml::cpu::riscv64_spacemit
 
-                p_input += gvl;
-                p_temp_output += gvl;
-                length -= gvl;
+static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack_type(const ggml_tensor * cur) {
+    switch (cur->type) {
+        case GGML_TYPE_Q2_K:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q2_k_32x256_q8_0;
+                }
+#endif
             }
+            break;
+        case GGML_TYPE_Q3_K:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q3_k_32x256_q8_0;
+                }
+#endif
+            }
+            break;
+        case GGML_TYPE_Q4_0:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && cur->ne[0] % 256 == 0 &&
+                    (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_0_32x256_q8_0;
+                }
 
-            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;
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_0_32x32_q8_0;
+                }
+#endif
 
-            mean_square = sqrt(mean_square + epsilon);
+#if defined(RISCV64_SPACEMIT_IME1)
+                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_0_16x32_q8_0;
+                }
+#endif
+            }
+            break;
+        case GGML_TYPE_Q4_1:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                // TODO
+                // if (cur->ne[1] % 32 == 0 && cur->ne[0] % 256 == 0 &&
+                //     (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                //     return &ggml::cpu::riscv64_spacemit::q4_1_32x256_q8_0;
+                // }
+
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_1_32x32_q8_0;
+                }
+#endif
 
-            mean_square   = 1.0f / mean_square;
-            length        = hidden_size;
-            p_temp_output = p_output;
+#if defined(RISCV64_SPACEMIT_IME1)
+                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_1_16x32_q8_0;
+                }
+#endif
+            }
+            break;
+        case GGML_TYPE_Q4_K:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_k_32x32_q8_0;
+                }
+#endif
 
-            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;
+#if defined(RISCV64_SPACEMIT_IME1)
+                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
+                    return &ggml::cpu::riscv64_spacemit::q4_k_16x32_q8_0;
                 }
-            } 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;
+#endif
+            }
+            break;
+        case GGML_TYPE_Q6_K:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if ((ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q6_k_32x32_q8_0;
                 }
-            } 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;
+#endif
+            }
+            break;
+        case GGML_TYPE_Q8_0:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if ((ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q8_0_32x32_q8_0;
                 }
+#endif
             }
-        }
-    }
-
-    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;
+            break;
+        case GGML_TYPE_MXFP4:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                // TODO
+                // if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                //     return &ggml::cpu::riscv64_spacemit::mxfp4_32x32_q8_0;
+                // }
+#endif
+            }
+            break;
+        case GGML_TYPE_Q5_K:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q5_k_32x32_q8_0;
+                }
+#endif
+            }
+            break;
+        case GGML_TYPE_Q5_1:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q5_1_32x32_q8_0;
+                }
+#endif
+            }
+            break;
+        case GGML_TYPE_Q5_0:
+            {
+#if defined(RISCV64_SPACEMIT_IME2)
+                if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::q5_0_32x32_q8_0;
+                }
+#endif
+            }
+            break;
+        default:
+            break;
     }
 
     return nullptr;
 }
 
 static enum ggml_status ggml_backend_riscv64_spacemit_buffer_init_tensor(ggml_backend_buffer_t buffer,
-                                                                         struct ggml_tensor *  tensor) {
+                                                                         ggml_tensor *         tensor) {
     tensor->extra =
         (void *) const_cast<ggml::cpu::tensor_traits *>(ggml_riscv64_spacemit_get_optimal_repack_type(tensor));
 
@@ -874,8 +1402,46 @@ static enum ggml_status ggml_backend_riscv64_spacemit_buffer_init_tensor(ggml_ba
     return GGML_STATUS_SUCCESS;
 }
 
+static void ggml_backend_riscv64_spacemit_buffer_free_buffer(ggml_backend_buffer_t buffer) {
+    GGML_ASSERT(buffer);
+
+    void * base = buffer->context;
+    if (base == nullptr) {
+        return;
+    }
+
+    ggml::cpu::riscv64_spacemit::spine_mem_pool_free(base);
+}
+
+static void * ggml_backend_riscv64_spacemit_buffer_get_base(ggml_backend_buffer_t buffer) {
+    GGML_ASSERT(buffer);
+
+    void * base = buffer->context;
+    GGML_ASSERT(base != nullptr);
+    return base;
+}
+
+static void ggml_backend_riscv64_spacemit_buffer_memset_tensor(ggml_backend_buffer_t buffer,
+                                                               ggml_tensor *         tensor,
+                                                               uint8_t               value,
+                                                               size_t                offset,
+                                                               size_t                size) {
+    GGML_ASSERT(tensor);
+    memset((char *) tensor->data + offset, value, size);
+
+    GGML_UNUSED(buffer);
+}
+
+static void ggml_backend_riscv64_spacemit_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
+    GGML_ASSERT(buffer);
+
+    void * base = buffer->context;
+    GGML_ASSERT(base != nullptr);
+    memset(base, value, buffer->size);
+}
+
 static void ggml_backend_riscv64_spacemit_buffer_set_tensor(ggml_backend_buffer_t buffer,
-                                                            struct ggml_tensor *  tensor,
+                                                            ggml_tensor *         tensor,
                                                             const void *          data,
                                                             size_t                offset,
                                                             size_t                size) {
@@ -891,6 +1457,20 @@ static void ggml_backend_riscv64_spacemit_buffer_set_tensor(ggml_backend_buffer_
     GGML_UNUSED(buffer);
 }
 
+static const ggml_backend_buffer_i ggml_backend_riscv64_spacemit_buffer_i = {
+    /* .free_buffer     = */ ggml_backend_riscv64_spacemit_buffer_free_buffer,
+    /* .get_base        = */ ggml_backend_riscv64_spacemit_buffer_get_base,
+    /* .init_tensor     = */ ggml_backend_riscv64_spacemit_buffer_init_tensor,
+    /* .memset_tensor   = */ ggml_backend_riscv64_spacemit_buffer_memset_tensor,
+    /* .set_tensor      = */ ggml_backend_riscv64_spacemit_buffer_set_tensor,
+    /* .get_tensor      = */ nullptr,
+    /* .set_tensor_2d   = */ nullptr,
+    /* .get_tensor_2d   = */ nullptr,
+    /* .cpy_tensor      = */ nullptr,
+    /* .clear           = */ ggml_backend_riscv64_spacemit_buffer_clear,
+    /* .reset           = */ nullptr,
+};
+
 static const char * ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
     return "CPU_RISCV64_SPACEMIT";
 
@@ -899,18 +1479,12 @@ static const char * ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name(ggml_
 
 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) {
+    void * base = ggml::cpu::riscv64_spacemit::spine_mem_pool_alloc(size, 64);
+    if (base == 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;
+    return ggml_backend_buffer_init(buft, ggml_backend_riscv64_spacemit_buffer_i, base, size);
 }
 
 static size_t ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
@@ -919,44 +1493,91 @@ static size_t ggml_backend_cpu_riscv64_spacemit_buffer_type_get_alignment(ggml_b
     GGML_UNUSED(buft);
 }
 
-static size_t ggml_backend_cpu_riscv64_spacemit_nbytes(ggml_backend_buffer_type_t buft,
-                                                       const struct ggml_tensor * tensor) {
+static size_t ggml_backend_cpu_riscv64_spacemit_nbytes(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
     for (int i = 0; i < GGML_MAX_DIMS; ++i) {
         if (tensor->ne[i] <= 0) {
             return 0;
         }
     }
 
-    size_t       nbytes;
+    GGML_UNUSED(buft);
+
+    const auto plain_nbytes = [&]() {
+        size_t total = ggml_type_size(tensor->type);
+        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+            total += (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+        return total;
+    };
+
     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];
+        return plain_nbytes();
+    }
+
+    const size_t row_nbytes = tensor->ne[0] * tensor->nb[0] / blck_size;
+
+    const auto add_strided_nbytes = [&](size_t total, size_t src_block_size, size_t dst_block_size) {
+        for (int i = 1; i < GGML_MAX_DIMS; ++i) {
+            total += (tensor->ne[i] - 1) * (tensor->nb[i] / src_block_size) * dst_block_size;
         }
-    } 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];
-            }
+        return total;
+    };
+
+    const auto remap_block_nbytes = [&](size_t src_block_size, size_t dst_block_size, int64_t padded_rows = 0) {
+        GGML_ASSERT(row_nbytes % src_block_size == 0);
+
+        size_t total =
+            add_strided_nbytes((row_nbytes / src_block_size) * dst_block_size, src_block_size, dst_block_size);
+
+        if (padded_rows > 0 && tensor->ne[1] % padded_rows != 0) {
+            total += (padded_rows - tensor->ne[1] % padded_rows) * (tensor->nb[1] / src_block_size) * dst_block_size;
         }
+
+        return total;
+    };
+
+    size_t nbytes = row_nbytes;
+    switch (tensor->type) {
+        case GGML_TYPE_Q4_K:
+            nbytes = remap_block_nbytes(sizeof(block_q4_K), sizeof(block_q4_1) * 8);
+            break;
+        case GGML_TYPE_Q6_K:
+            nbytes = remap_block_nbytes(sizeof(block_q6_K), sizeof(block_q8_0) * 8, 32);
+            break;
+        case GGML_TYPE_Q8_0:
+            nbytes = remap_block_nbytes(sizeof(block_q8_0), sizeof(block_q8_0), 32);
+            break;
+        case GGML_TYPE_Q2_K:
+            nbytes = remap_block_nbytes(sizeof(block_q2_K), sizeof(spacemit_kernels::nrow_block_q2_k<1>));
+            break;
+        case GGML_TYPE_Q3_K:
+            nbytes = remap_block_nbytes(sizeof(block_q3_K), sizeof(spacemit_kernels::nrow_block_q3_k<1>));
+            break;
+        case GGML_TYPE_MXFP4:
+            nbytes = remap_block_nbytes(sizeof(block_mxfp4), sizeof(spacemit_kernels::nrow_block_mxfp4<1>));
+            break;
+        case GGML_TYPE_Q5_K:
+            nbytes = remap_block_nbytes(sizeof(block_q5_K), sizeof(spacemit_kernels::nrow_block_q5_1<1>) * 8);
+            break;
+        case GGML_TYPE_Q5_1:
+            nbytes = remap_block_nbytes(sizeof(block_q5_1), sizeof(spacemit_kernels::nrow_block_q5_1<1>));
+            break;
+        case GGML_TYPE_Q5_0:
+            nbytes = remap_block_nbytes(sizeof(block_q5_0), sizeof(spacemit_kernels::nrow_block_q5_0<1>));
+            break;
+        default:
+            nbytes = add_strided_nbytes(row_nbytes, 1, 1);
+            break;
     }
 
-    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 {
+    bool supports_op(ggml_backend_dev_t, const ggml_tensor * op) override {
         switch (op->op) {
             case GGML_OP_MUL_MAT:
                 if (op->src[0]->buffer && (ggml_n_dims(op->src[0]) == 2) &&
@@ -970,10 +1591,16 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
                     }
                 }
                 break;
-            case GGML_OP_NORM:
-            case GGML_OP_RMS_NORM:
-                if (op->src[0]->type == GGML_TYPE_F32) {
-                    return true;
+            case GGML_OP_MUL_MAT_ID:
+                if (op->src[0]->buffer && (ggml_n_dims(op->src[0]) == 3) &&
+                    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;
             default:
@@ -983,15 +1610,28 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
         return false;
     }
 
-    ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {
+    ggml::cpu::tensor_traits * get_tensor_traits(const ggml_tensor * op) override {
         switch (op->op) {
             case GGML_OP_MUL_MAT:
+            case GGML_OP_MUL_MAT_ID:
                 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:
+            case GGML_OP_ADD:
+            case GGML_OP_SUB:
+            case GGML_OP_MUL:
+            case GGML_OP_DIV:
+            case GGML_OP_FLASH_ATTN_EXT:
+            case GGML_OP_CONT:
+            case GGML_OP_CPY:
+            case GGML_OP_REPEAT:
+            case GGML_OP_SUM_ROWS:
+            case GGML_OP_GET_ROWS:
+            case GGML_OP_CONCAT:
+                // case GGML_OP_GATED_DELTA_NET:
                 return (ggml::cpu::tensor_traits *) (&ggml::cpu::riscv64_spacemit::rvv_impl);
             default:
                 // GGML_ABORT("fatal error");
@@ -1005,7 +1645,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
 }  // 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 = {
+    static ggml_backend_buffer_type ggml_backend_cpu_buffer_type_riscv64_spacemit = {
   /* .iface    = */
         {
          /* .get_name         = */ ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name,
@@ -1023,3 +1663,78 @@ ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void) {
 
     return &ggml_backend_cpu_buffer_type_riscv64_spacemit;
 }
+
+extern "C" {
+static int bind_ai_thread() {
+    int  fd, bytes;
+    char str[32];
+
+    fd = open("/proc/set_ai_thread", O_WRONLY);
+    if (fd < 0) {
+        GGML_LOG_ERROR("try open /proc/set_ai_thread failed\n");
+        return -1;
+    }
+
+    snprintf(str, 16, "%d", 0);
+    bytes = write(fd, str, strlen(str));
+    if (bytes < 0) {
+        GGML_LOG_ERROR("try write /proc/set_ai_thread failed\n");
+        close(fd);
+        return -1;
+    }
+
+    close(fd);
+    return 0;
+}
+
+void ggml_backend_cpu_riscv64_spacemit_set_numa_thread_affinity(int thread_n) {
+    int cpu_id = sched_getcpu();
+    if (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2 &&
+        !((1 << cpu_id) & ggml::cpu::riscv64_spacemit::global_spine_env_info.cpu_mask)) {
+        GGML_PRINT_DEBUG("bind_ai_thread for thread %d, pid %d\n", thread_n, getpid());
+        bind_ai_thread();
+    }
+
+    if (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_tcm &&
+        ggml::cpu::riscv64_spacemit::tls_context.cpu_id == -1) {
+        CPU_ZERO(&(ggml::cpu::riscv64_spacemit::tls_context.cpuset));
+        pthread_t    main_thread     = pthread_self();
+        const auto & perfer_core_ids = ggml::cpu::riscv64_spacemit::global_spine_env_info.perfer_core_ids;
+        if (thread_n < 0 || static_cast<size_t>(thread_n) >= perfer_core_ids.size()) {
+            GGML_ABORT("thread_n %d exceeds perfer_core_ids size %zu\n", thread_n, perfer_core_ids.size());
+        }
+        auto perfer_cpu_id = perfer_core_ids[static_cast<size_t>(thread_n)];
+        CPU_SET(perfer_cpu_id, &(ggml::cpu::riscv64_spacemit::tls_context.cpuset));
+        int s =
+            pthread_setaffinity_np(main_thread, sizeof(cpu_set_t), &(ggml::cpu::riscv64_spacemit::tls_context.cpuset));
+        if (s != 0) {
+            GGML_ABORT("set thread affinity error for thread_n %d, cpu_id %d\n", thread_n, perfer_cpu_id);
+        }
+
+        int ai_cpu_id = perfer_cpu_id - ggml::cpu::riscv64_spacemit::global_spine_env_info.aicpu_id_offset;
+        ggml::cpu::riscv64_spacemit::tls_context.cpu_id = ai_cpu_id;
+        ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer =
+            ggml::cpu::riscv64_spacemit::spine_mem_pool_tcm_mem_get(ai_cpu_id);
+        ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer_size =
+            ggml::cpu::riscv64_spacemit::global_spine_env_info.tcm_blk_size;
+    }
+
+    if (ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer != nullptr) {
+        void * rt =
+            ggml::cpu::riscv64_spacemit::spine_mem_pool_tcm_mem_wait(ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
+        if (rt == nullptr) {
+            GGML_ABORT("wait tcm buffer failed for cpu_id: %d", ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
+        }
+    }
+}
+
+void ggml_backend_cpu_riscv64_spacemit_clear_numa_thread_affinity_threaded(int thread_n) {
+    if (ggml::cpu::riscv64_spacemit::tls_context.tcm_buffer != nullptr) {
+        auto rt = ggml::cpu::riscv64_spacemit::spine_mem_pool_tcm_mem_release(
+            ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
+        if (rt != 0) {
+            GGML_ABORT("release tcm buffer failed for cpu_id: %d", ggml::cpu::riscv64_spacemit::tls_context.cpu_id);
+        }
+    }
+}
+}