chore: import upstream snapshot with attribution

kt-kernel/operators/moe_kernel/api/common.h

@@ -0,0 +1,63 @@
+// BOOST_STRONG_TYPEDEF(int8_t, int4_2_t);
+#pragma once
+#include <cstdint>
+
+#include "llama.cpp/ggml.h"
+#if !defined(CPUINFER_HAS_FLOAT16_T)
+using float16_t = ggml_fp16_t;
+#define CPUINFER_HAS_FLOAT16_T 1
+#endif
+
+#if !defined(CPUINFER_HAS_BFLOAT16_T)
+using bfloat16_t = ggml_bf16_t;
+#define CPUINFER_HAS_BFLOAT16_T 1
+#endif  // CPUINFER_HAS_BFLOAT16_T
+const bool PACKED = true;
+#if defined(__aarch64__) || defined(__arm__) || defined(CPU_USE_KML)
+#ifndef CPU_USE_KML
+#define CPU_USE_KML
+#endif
+#endif  // USE_MOE_KERNEL_AMD or CPU_USE_KML
+
+#define STRONG_TYPEDEF(T, D)                                   \
+  struct D {                                                   \
+    T t;                                                       \
+    explicit D(const T &v) : t(v) {}                           \
+    D() = default;                                             \
+    D(const D &) = default;                                    \
+    D &operator=(const D &) = default;                         \
+    D &operator=(const T &rhs) {                               \
+      t = rhs;                                                 \
+      return *this;                                            \
+    }                                                          \
+    operator const T &() const { return t; }                   \
+    operator T &() { return t; }                               \
+    bool operator==(const D &rhs) const { return t == rhs.t; } \
+    bool operator!=(const D &rhs) const { return t != rhs.t; } \
+    bool operator<(const D &rhs) const { return t < rhs.t; }   \
+  };
+STRONG_TYPEDEF(int8_t, int4_2_t)
+typedef int8_t BLASINT8;
+
+/* matrix transpose or conjugate transpose */
+typedef enum KERNEL_CBLAS_TRANSPOSE {
+  KernelCblasNoTrans = 111,
+  KernelCblasTrans = 112,
+  KernelCblasConjTrans = 113,
+  KernelCblasConjNoTrans = 114
+} KERNEL_CBLAS_TRANSPOSE;
+/* matrix stored in rows or cols */
+typedef enum KERNEL_CBLAS_ORDER { KernelCblasRowMajor = 101, KernelCblasColMajor = 102 } KERNEL_CBLAS_ORDER;
+/* matrix position is left or right */
+typedef enum KERNEL_CBLAS_SIDE { KernelCblasLeft = 141, KernelCblasRight = 142 } KERNEL_CBLAS_SIDE;
+typedef KERNEL_CBLAS_ORDER KERNEL_CBLAS_LAYOUT;
+typedef enum KERNEL_CBLAS_OFFSET {
+  KernelCblasRowOffset = 171,
+  KernelCblasColOffset = 172,
+  KernelCblasFixOffset = 173
+} KERNEL_CBLAS_OFFSET;
+
+enum class MatKernelVariant {
+  Decode,
+  Prefill,
+};

kt-kernel/operators/moe_kernel/api/mat_kernel.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include "common.h"
+
+using GemmFn = void (*)(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                        const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                        const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                        const int8_t oa, const void* b, const size_t ldb, const int8_t ob, const float beta, int32_t* c,
+                        const size_t ldc, const int32_t* oc);
+
+struct MatKernelSelection {
+  GemmFn fn;
+  int divide_elements_size;
+};
+
+MatKernelSelection select_kernel_for_int4(MatKernelVariant variant);
+MatKernelSelection select_kernel_for_int8(MatKernelVariant variant);
+
+template <typename T>
+MatKernelSelection select_mat_kernel(MatKernelVariant variant) {
+  if constexpr (std::is_same_v<typename T::dt, int4_2_t>) {
+    return select_kernel_for_int4(variant);
+  } else {
+    return select_kernel_for_int8(variant);
+  }
+}

kt-kernel/operators/moe_kernel/la/mat_kernel.cpp

@@ -0,0 +1,54 @@
+#include "../api/mat_kernel.h"
+
+#include <cassert>
+
+namespace {
+constexpr int kInt4ElementDivisor = 2;
+constexpr int kInt8ElementDivisor = 1;
+}  // namespace
+extern "C" {
+void decode_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                               const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                               const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                               const int8_t oa, const void* b, const size_t ldb, const int8_t ob, const float beta,
+                               int32_t* c, const size_t ldc, const int32_t* oc);
+
+void prefill_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                                const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                                const int8_t oa, const void* b, const size_t ldb, const int8_t ob, const float beta,
+                                int32_t* c, const size_t ldc, const int32_t* oc);
+
+void decode_int4_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                    const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc,
+                                    const size_t m, const size_t n, const size_t k, const float alpha, const void* a,
+                                    const size_t lda, const int8_t oa, const void* b, const size_t ldb, const int8_t ob,
+                                    const float beta, int32_t* c, const size_t ldc, const int32_t* oc);
+
+void prefill_int4_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                     const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc,
+                                     const size_t m, const size_t n, const size_t k, const float alpha, const void* a,
+                                     const size_t lda, const int8_t oa, const void* b, const size_t ldb,
+                                     const int8_t ob, const float beta, int32_t* c, const size_t ldc,
+                                     const int32_t* oc);
+}
+
+MatKernelSelection select_kernel_for_int4(MatKernelVariant variant) {
+  switch (variant) {
+    case MatKernelVariant::Decode:
+      return {decode_int4_cblas_gemm_s8s8s32, kInt4ElementDivisor};
+    case MatKernelVariant::Prefill:
+      return {prefill_int4_cblas_gemm_s8s8s32, kInt4ElementDivisor};
+  }
+  return {nullptr, 0};
+}
+
+MatKernelSelection select_kernel_for_int8(MatKernelVariant variant) {
+  switch (variant) {
+    case MatKernelVariant::Decode:
+      return {decode_cblas_gemm_s8s8s32, kInt8ElementDivisor};
+    case MatKernelVariant::Prefill:
+      return {prefill_cblas_gemm_s8s8s32, kInt8ElementDivisor};
+  }
+  return {nullptr, 0};
+}

kt-kernel/operators/moe_kernel/la/utils.hpp

@@ -0,0 +1,29 @@
+#pragma once
+// #include <arm_sve.h>
+#include <cstdint>
+#include <cstring>
+
+// 简单截断模式：直接丢弃低 16 位
+static inline uint16_t float_to_bf16_trunc(float f) {
+  uint32_t u;
+  // 按位拷贝，避免 strict‑aliasing UB
+  memcpy(&u, &f, sizeof(u));   // :contentReference[oaicite:3]{index=3}
+  return (uint16_t)(u >> 16);  // 截断得到高 16 位 :contentReference[oaicite:4]{index=4}
+}
+
+static inline void convert_32fp32_to_32bf16_pure_c(const float* src, uint16_t* dst) {
+  // src 已偏移至 token_nth * hidden_size
+  for (int e = 0; e < 32; e++) {  // 共 32 个元素
+    // 选择截断或四舍五入
+    dst[e] = float_to_bf16_trunc(src[e]);
+  }
+}
+
+// 把 32 个 bf16 元素转换成 32 个 fp32 元素
+
+static inline void convert_32bf16_to_32fp32_pure_c(const uint16_t* src, float* dst) {
+  for (int e = 0; e < 32; e++) {
+    uint32_t temp = ((uint32_t)src[e]) << 16;  // 将 BF16 左移 16 位
+    memcpy(&dst[e], &temp, sizeof(float));     // 将结果复制到 FP32 变量中
+  }
+}

kt-kernel/operators/moe_kernel/mat_kernel/aocl_kernel/kernel.cpp

@@ -0,0 +1,100 @@
+#include <stdexcept>
+
+#include "../batch_gemm_api.hpp"
+#include "blis.h"
+
+namespace {
+
+char ToAoclOrder(KERNEL_CBLAS_LAYOUT layout) {
+  switch (layout) {
+    case KernelCblasRowMajor:
+      return 'r';
+    case KernelCblasColMajor:
+      return 'c';
+  }
+  throw std::invalid_argument("Unsupported KERNEL_CBLAS_LAYOUT value");
+}
+
+char ToAoclTranspose(KERNEL_CBLAS_TRANSPOSE transpose) {
+  switch (transpose) {
+    case KernelCblasNoTrans:
+      return 'n';
+    case KernelCblasTrans:
+      return 't';
+    case KernelCblasConjTrans:
+    case KernelCblasConjNoTrans:
+      break;
+  }
+  throw std::invalid_argument("Unsupported KERNEL_CBLAS_TRANSPOSE value");
+}
+
+}  // namespace
+
+// 映射表，layout 从KERNEL_CBLAS_ORDER 映射到'r'或者'c',以及将KERNEL_CBLAS_TRANSPOSE映射到'n'或者't'
+#ifdef __cplusplus
+extern "C" {
+#endif
+void decode_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                               const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                               const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                               const BLASINT8 oa, const void* b, const size_t ldb, const BLASINT8 ob, const float beta,
+                               int32_t* c, const size_t ldc, const int32_t* oc) {
+  const char order = ToAoclOrder(layout);
+  const char op_a = ToAoclTranspose(transa);
+  const char op_b = ToAoclTranspose(transb);
+  (void)offsetc;
+  aocl_gemm_s8s8s32os32(order, op_a, op_b, static_cast<dim_t>(m), static_cast<dim_t>(n), static_cast<dim_t>(k),
+                        static_cast<int32_t>(alpha), static_cast<const int8_t*>(a), static_cast<dim_t>(lda), 'n',
+                        static_cast<const int8_t*>(b), static_cast<dim_t>(ldb), 'r', static_cast<int32_t>(beta), c,
+                        static_cast<dim_t>(ldc), nullptr);
+}
+
+void prefill_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                                const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                                const BLASINT8 oa, const void* b, const size_t ldb, const BLASINT8 ob, const float beta,
+                                int32_t* c, const size_t ldc, const int32_t* oc) {
+  const char order = ToAoclOrder(layout);
+  const char op_a = ToAoclTranspose(transa);
+  const char op_b = ToAoclTranspose(transb);
+  (void)offsetc;
+  aocl_gemm_s8s8s32os32(order, op_a, op_b, static_cast<dim_t>(m), static_cast<dim_t>(n), static_cast<dim_t>(k),
+                        static_cast<int32_t>(alpha), static_cast<const int8_t*>(a), static_cast<dim_t>(lda), 'n',
+                        static_cast<const int8_t*>(b), static_cast<dim_t>(ldb), 'r', static_cast<int32_t>(beta), c,
+                        static_cast<dim_t>(ldc), nullptr);
+}
+
+void prefill_int4_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                     const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc,
+                                     const size_t m, const size_t n, const size_t k, const float alpha, const void* a,
+                                     const size_t lda, const BLASINT8 oa, const void* b, const size_t ldb,
+                                     const BLASINT8 ob, const float beta, int32_t* c, const size_t ldc,
+                                     const int32_t* oc) {
+  throw std::runtime_error("int4 not support prefill");
+}
+
+void decode_int4_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                    const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc,
+                                    const size_t m, const size_t n, const size_t k, const float alpha, const void* a,
+                                    const size_t lda, const BLASINT8 oa, const void* b, const size_t ldb,
+                                    const BLASINT8 ob, const float beta, int32_t* c, const size_t ldc,
+                                    const int32_t* oc) {
+  throw std::runtime_error("int4 not support decode");
+}
+
+void reorder_B_gemm(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transb, const size_t k,
+                    const size_t n, const size_t ldb, const void* b, void* b_reordered) {
+  const char order = ToAoclOrder(layout);
+  const char op_b = ToAoclTranspose(transb);
+  aocl_reorder_s8s8s32os32(order, op_b, 'B', static_cast<const int8_t*>(b), static_cast<int8_t*>(b_reordered), k, n,
+                           ldb);
+}
+
+size_t get_reorder_B_size(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transb, const size_t k,
+                          const size_t n) {
+  return aocl_get_reorder_buf_size_s8s8s32os32(ToAoclOrder(layout), ToAoclTranspose(transb), 'B', k, n);
+}
+
+#ifdef __cplusplus
+}
+#endif

kt-kernel/operators/moe_kernel/mat_kernel/batch_gemm_api.hpp

@@ -0,0 +1,42 @@
+#pragma once
+#include <cstddef>
+#ifndef _BATCH_GEMM_KERNEL_API_
+#define _BATCH_GEMM_KERNEL_API_
+#include "../api/common.h"
+#ifdef __cplusplus
+extern "C" {
+#endif
+void decode_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                               const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                               const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                               const BLASINT8 oa, const void* b, const size_t ldb, const BLASINT8 ob, const float beta,
+                               int32_t* c, const size_t ldc, const int32_t* oc);
+
+void prefill_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc, const size_t m,
+                                const size_t n, const size_t k, const float alpha, const void* a, const size_t lda,
+                                const BLASINT8 oa, const void* b, const size_t ldb, const BLASINT8 ob, const float beta,
+                                int32_t* c, const size_t ldc, const int32_t* oc);
+
+void decode_int4_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                    const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc,
+                                    const size_t m, const size_t n, const size_t k, const float alpha, const void* a,
+                                    const size_t lda, const BLASINT8 oa, const void* b, const size_t ldb,
+                                    const BLASINT8 ob, const float beta, int32_t* c, const size_t ldc,
+                                    const int32_t* oc);
+
+void prefill_int4_cblas_gemm_s8s8s32(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transa,
+                                     const KERNEL_CBLAS_TRANSPOSE transb, const KERNEL_CBLAS_OFFSET offsetc,
+                                     const size_t m, const size_t n, const size_t k, const float alpha, const void* a,
+                                     const size_t lda, const BLASINT8 oa, const void* b, const size_t ldb,
+                                     const BLASINT8 ob, const float beta, int32_t* c, const size_t ldc,
+                                     const int32_t* oc);
+void reorder_B_gemm(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transb, const size_t k,
+                    const size_t n, const size_t ldb, const void* b, void* b_reordered);
+size_t get_reorder_B_size(const KERNEL_CBLAS_LAYOUT layout, const KERNEL_CBLAS_TRANSPOSE transb, const size_t k,
+                          const size_t n);
+
+#ifdef __cplusplus
+}
+#endif
+#endif /*** _BATCH_GEMM_KERNEL_API_ ***/

kt-kernel/operators/moe_kernel/moe.hpp

@@ -0,0 +1,800 @@
+#ifndef MOE_KERNEL_HPP
+#define MOE_KERNEL_HPP
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <vector>
+
+#include "../../cpu_backend/shared_mem_buffer.h"
+#include "../common.hpp"
+#include "../moe-tp.hpp"
+#include "api/common.h"
+#include "api/mat_kernel.h"
+#include "llama.cpp/ggml.h"
+template <class T, bool PLAIN = true>
+class MOE_KERNEL_TP
+#ifdef FORWARD_TIME_PROFILE
+    : protected TimePerf
+#endif
+{
+ private:
+  int tp_part_idx;
+  std::filesystem::path prefix;
+
+  void* gate_proj_;  // [expert_num * intermediate_size * hidden_size ( /32 if
+                     // quantized)]
+  void* up_proj_;    // [expert_num * intermediate_size * hidden_size ( /32 if
+                     // quantized)]
+  void* down_proj_;  // [expert_num * hidden_size * intermediate_size ( /32 if
+                     // quantized)]
+
+  ggml_bf16_t* m_local_input_;  // [routed_expert_num * max_len * hidden_size]
+  float* m_local_gate_output_;  // [routed_expert_num * max_len * intermediate_size]
+  float* m_local_up_output_;    // [routed_expert_num * max_len * intermediate_size]
+  float* m_local_down_output_;  // [routed_expert_num * max_len * hidden_size]
+
+  std::vector<std::vector<int>> m_local_pos_;    // [max_len, routed_expert_num]
+  std::vector<int> m_local_num_;                 // [expert_num]
+  std::vector<int> m_expert_id_map_;             // [expert_num]
+  std::vector<ggml_bf16_t*> m_local_input_ptr_;  // [expert_num]
+  std::vector<float*> m_local_gate_output_ptr_;  // [expert_num]
+  std::vector<float*> m_local_up_output_ptr_;    // [expert_num]
+  std::vector<float*> m_local_down_output_ptr_;  // [expert_num]
+
+  std::vector<std::shared_ptr<typename T::BufferA>> gate_up_ba_;
+  std::vector<std::shared_ptr<typename T::BufferB>> gate_bb_;
+  std::vector<std::shared_ptr<typename T::BufferC>> gate_bc_;
+  std::vector<std::shared_ptr<typename T::BufferB>> up_bb_;
+  std::vector<std::shared_ptr<typename T::BufferC>> up_bc_;
+  std::vector<std::shared_ptr<typename T::BufferA>> down_ba_;
+  std::vector<std::shared_ptr<typename T::BufferB>> down_bb_;
+  std::vector<std::shared_ptr<typename T::BufferC>> down_bc_;
+
+  std::vector<void*> gate_up_owner_ptr_;
+  std::vector<void*> down_owner_ptr_;
+
+  inline void write_weights(std::filesystem::path prefix, std::string mat_class, char* bb, int expert_idx, size_t size,
+                            size_t scale_size) {
+    // printf("expert %d, size %ld, scale size %ld\n", expert_idx, size, scale_size);
+    // std::ofstream of(prefix / (T::name() + mat_class + std::to_string(expert_idx)  + "_quant_" + ".kt"));
+    std::ofstream of(prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" +
+                               std::to_string(size - scale_size) + "Byte" + "_quant_" + ".kt"));
+    if (of.is_open() == false) {
+      printf("no such file: %s", (prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" +
+                                            std::to_string(size - scale_size) + "Byte" + "_quant_" + ".kt"))
+                                     .c_str());
+      // throw std::runtime_error("No such file");
+    }
+    of.write((char*)bb, size - scale_size);
+    of.close();
+    // of.open(prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_scale_" + ".kt"));
+    of.open(prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" + std::to_string(scale_size) + "Byte" +
+                      "_scale_" + ".kt"));
+    if (of.is_open() == false) {
+      printf("no such file\n");
+      // throw std::runtime_error("No such file");
+    }
+    of.write(((char*)bb) + size - scale_size, scale_size);
+  }
+
+  inline void read_weights(std::filesystem::path prefix, std::string mat_class, char* bb, int expert_idx, size_t size,
+                           size_t scale_size, uint8_t mat_split, uint8_t mat_split_idex) {
+    // std::ifstream f(prefix / (T::name() + mat_class + std::to_string(expert_idx)  + "_quant_" + ".kt"));
+    std::ifstream f(prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" +
+                              std::to_string(size - scale_size) + "Byte" + "_quant_" + ".kt"));
+    if (f.is_open() == false) {
+      printf("no such file: %s\n", (prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" +
+                                              std::to_string(size - scale_size) + "Byte" + "_quant_" + ".kt"))
+                                       .c_str());
+      // throw std::runtime_error("No such file");
+    }
+    f.seekg(mat_split_idex * (size - scale_size) / mat_split);
+    f.read(((char*)bb) + mat_split_idex * (size - scale_size) / mat_split, (size - scale_size) / mat_split);
+    f.close();
+    // f.open(prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_scale_" + ".kt"));
+    f.open(prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" + std::to_string(scale_size) + "Byte" +
+                     "_scale_" + ".kt"));
+    if (f.is_open() == false) {
+      printf("no such file: %s\n", (prefix / (T::name() + mat_class + std::to_string(expert_idx) + "_" +
+                                              std::to_string(scale_size) + "Byte" + "_scale_" + ".kt"))
+                                       .c_str());
+      // throw std::runtime_error("No such file");
+    }
+    f.seekg(mat_split_idex * scale_size / mat_split);
+    f.read((((char*)bb) + size - scale_size) + mat_split_idex * scale_size / mat_split, scale_size / mat_split);
+  }
+
+ public:
+  using input_t = ggml_bf16_t;
+  using output_t = float;
+
+  GeneralMOEConfig config_;
+  static constexpr double ELEMENT_SIZE = T::ELEMENT_SIZE;
+
+  MOE_KERNEL_TP(GeneralMOEConfig config, int tp_part_idx) {
+    printf("  Creating AMD_MOE_TP %d at numa %d\n", tp_part_idx, numa_node_of_cpu(sched_getcpu()));
+    auto& load = config.load;
+    auto& save = config.save;
+    if (load && config.path == "") {
+      load = false;
+    }
+
+    prefix = config.path;
+    prefix = prefix / ("_layer_" + std::to_string(config.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx));
+    if (save) {
+      std::cout << "Creating " << prefix << std::endl;
+      std::filesystem::create_directories(prefix);
+    }
+    if (load) {
+      if (std::filesystem::exists(prefix)) {
+        std::cout << "Loading from " << prefix << std::endl;
+      } else {
+        throw std::runtime_error("Path not found: " + prefix.string());
+      }
+    }
+
+    this->tp_part_idx = tp_part_idx;
+    config_ = config;
+    gate_proj_ = config_.gate_proj;
+    up_proj_ = config_.up_proj;
+    down_proj_ = config_.down_proj;
+
+    MemoryRequest mem_requests;
+    mem_requests.append_pointer(&m_local_input_,
+                                sizeof(input_t) * config_.num_experts_per_tok * config_.max_len * config_.hidden_size);
+
+    mem_requests.append_pointer(&m_local_gate_output_, sizeof(float) * config_.num_experts_per_tok * config_.max_len *
+                                                           config_.intermediate_size);
+    mem_requests.append_pointer(
+        &m_local_up_output_, sizeof(float) * config_.num_experts_per_tok * config_.max_len * config_.intermediate_size);
+    mem_requests.append_pointer(&m_local_down_output_,
+                                sizeof(float) * config_.num_experts_per_tok * config_.max_len * config_.hidden_size);
+
+    m_local_pos_.resize(config_.max_len);
+    for (int i = 0; i < config_.max_len; i++) {
+      m_local_pos_[i].resize(config_.num_experts_per_tok);
+    }
+    m_expert_id_map_.resize(config_.expert_num);
+    m_local_num_.resize(config_.expert_num);
+    m_local_input_ptr_.resize(config_.expert_num);
+    m_local_gate_output_ptr_.resize(config_.expert_num);
+    m_local_up_output_ptr_.resize(config_.expert_num);
+    m_local_down_output_ptr_.resize(config_.expert_num);
+
+    // printf("tp part %d alloc layer %d, %f GB, on numa %d\n", tp_part_idx, config_.layer_idx,
+    //        1e-9 * config_.expert_num *
+    //            (T::BufferB::required_size(config_.intermediate_size, config_.hidden_size) * 2 +
+    //             T::BufferB::required_size(config_.hidden_size, config_.intermediate_size)),
+    //        numa_node_of_cpu(sched_getcpu()));
+    // 统一分配一块巨大的内存用于权重：
+    size_t gate_up_exp_size =
+        T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN) +
+        T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+
+    for (uint64_t i = 0; i < config_.expert_num; i++) {
+      gate_up_ba_.push_back(std::make_shared<typename T::BufferA>(config_.max_len, config_.hidden_size, nullptr));
+      gate_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.intermediate_size, nullptr));
+      up_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.intermediate_size, nullptr));
+      down_ba_.push_back(std::make_shared<typename T::BufferA>(config_.max_len, config_.intermediate_size, nullptr));
+      down_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.hidden_size, nullptr));
+      void* gate_up_down_per_exp_ptr = std::aligned_alloc(64, gate_up_exp_size);
+      gate_up_owner_ptr_.push_back(gate_up_down_per_exp_ptr);
+
+      gate_bb_.push_back(std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size,
+                                                               gate_up_down_per_exp_ptr, PACKED, 'u', PLAIN));
+      up_bb_.push_back(std::make_shared<typename T::BufferB>(
+          config_.intermediate_size, config_.hidden_size,
+          offset_pointer(gate_up_down_per_exp_ptr,
+                         T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN)),
+          PACKED, 'u', PLAIN));
+
+      void* down_bb_ptr = std::aligned_alloc(
+          64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size, PACKED, 'd', PLAIN));
+      down_owner_ptr_.push_back(down_bb_ptr);
+      down_bb_.push_back(std::make_shared<typename T::BufferB>(config_.hidden_size, config_.intermediate_size,
+                                                               down_bb_ptr, PACKED, 'd', PLAIN));
+    }
+
+    for (int i = 0; i < config_.expert_num; i++) {
+      mem_requests.append_function([this, i](void* new_ptr) { gate_up_ba_[i]->set_data(new_ptr); },
+                                   T::BufferA::required_size(config_.max_len, config_.hidden_size));
+      mem_requests.append_function([this, i](void* new_ptr) { gate_bc_[i]->set_data(new_ptr); },
+                                   T::BufferC::required_size(config_.max_len, config_.intermediate_size));
+      mem_requests.append_function([this, i](void* new_ptr) { up_bc_[i]->set_data(new_ptr); },
+                                   T::BufferC::required_size(config_.max_len, config_.intermediate_size));
+      mem_requests.append_function([this, i](void* new_ptr) { down_ba_[i]->set_data(new_ptr); },
+                                   T::BufferA::required_size(config_.max_len, config_.intermediate_size));
+      mem_requests.append_function([this, i](void* new_ptr) { down_bc_[i]->set_data(new_ptr); },
+                                   T::BufferC::required_size(config_.max_len, config_.hidden_size));
+    }
+
+    shared_mem_buffer_numa.alloc(tp_part_idx, this, mem_requests);
+  }
+
+  MOE_KERNEL_TP(const MOE_KERNEL_TP&) = delete;
+  MOE_KERNEL_TP& operator=(const MOE_KERNEL_TP&) = delete;
+  MOE_KERNEL_TP(MOE_KERNEL_TP&&) = delete;
+  MOE_KERNEL_TP& operator=(MOE_KERNEL_TP&&) = delete;
+
+  ~MOE_KERNEL_TP() {
+    // printf("  Destroying KML_MOE_TP %lx\n", (intptr_t)(this));
+    for (void* ptr : gate_up_owner_ptr_) {
+      std::free(ptr);
+    }
+    for (void* ptr : down_owner_ptr_) {
+      std::free(ptr);
+    }
+  }
+
+  void load_weights() {
+    auto pool = config_.pool->get_subpool(tp_part_idx);
+    const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map;
+    if (config_.gate_projs.size()) {
+      printf("load from safetensor");
+      pool->do_work_stealing_job(
+          config_.expert_num, nullptr,
+          [this, physical_to_logical_map](int expert_id) {
+            uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_id);
+            {
+              size_t scale_size = config_.intermediate_size * sizeof(float);
+              size_t whole_size_ =
+                  T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+              size_t size = whole_size_ - scale_size;
+              void* dst_ = PLAIN ? gate_bb_[expert_id]->b : gate_bb_[expert_id]->b_pack[0];
+
+              memcpy(dst_, config_.gate_projs[tp_part_idx][logical_expert_id], size);
+
+              if constexpr (T::BufferB::SCALE) {
+                memcpy(gate_bb_[expert_id]->d, config_.gate_scales[tp_part_idx][logical_expert_id], scale_size);
+              }
+
+              whole_size_ =
+                  T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+              size = whole_size_ - scale_size;
+              dst_ = PLAIN ? up_bb_[expert_id]->b : up_bb_[expert_id]->b_pack[0];
+              memcpy(dst_, config_.up_projs[tp_part_idx][logical_expert_id], size);
+
+              if constexpr (T::BufferB::SCALE) {
+                memcpy(up_bb_[expert_id]->d, config_.up_scales[tp_part_idx][logical_expert_id], scale_size);
+              }
+            }
+
+            {
+              size_t scale_size = config_.hidden_size * sizeof(float);
+              size_t whole_size_ =
+                  T::BufferB::required_size(config_.hidden_size, config_.intermediate_size, PACKED, 'd', PLAIN);
+              size_t size = whole_size_ - scale_size;
+              void* dst_ = PLAIN ? down_bb_[expert_id]->b : down_bb_[expert_id]->b_pack[0];
+              memcpy(dst_, config_.down_projs[tp_part_idx][logical_expert_id], size);
+
+              if constexpr (T::BufferB::SCALE) {
+                memcpy(down_bb_[expert_id]->d, config_.down_scales[tp_part_idx][logical_expert_id], scale_size);
+              }
+            }
+          },
+          nullptr);
+
+    } else {
+      static uint8_t mat_type_all = 3, mat_split = 1;
+      if (config_.load) {
+        std::cout << "Loading from " << prefix << std::endl;
+        for (int task_id = 0; task_id < config_.expert_num * mat_type_all * mat_split; task_id++) {
+          int64_t expert_idx = task_id / (mat_type_all * mat_split);
+          uint8_t mat_class = (task_id % (mat_type_all * mat_split)) / mat_split;
+          uint8_t mat_split_idex = task_id % mat_split;
+          uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
+          void* src_;
+          if (mat_class == 0) {  // the up matrix
+            src_ = PLAIN ? up_bb_[expert_idx]->b : up_bb_[expert_idx]->b_pack[0];
+            size_t size = T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+            size_t scale_size = config_.intermediate_size * sizeof(float);
+            read_weights(prefix, "_up_", (char*)src_, logical_expert_id, size, scale_size, mat_split, mat_split_idex);
+          } else if (mat_class == 1) {
+            void* src_ = PLAIN ? gate_bb_[expert_idx]->b : gate_bb_[expert_idx]->b_pack[0];
+            size_t size = T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+            size_t scale_size = config_.intermediate_size * sizeof(float);
+            read_weights(prefix, "_gate_", (char*)src_, logical_expert_id, size, scale_size, mat_split, mat_split_idex);
+          } else {
+            void* src_ = PLAIN ? down_bb_[expert_idx]->b : down_bb_[expert_idx]->b_pack[0];
+            size_t size = T::BufferB::required_size(config_.hidden_size, config_.intermediate_size, PACKED, 'd', PLAIN);
+            size_t scale_size = config_.hidden_size * sizeof(float);
+            read_weights(prefix, "_down_", (char*)src_, logical_expert_id, size, scale_size, mat_split, mat_split_idex);
+          }
+        }
+      }
+// check process, store down matrix to check
+#ifdef CHECK
+      load_check();
+#endif
+#ifndef CHECK
+      else
+#endif
+      {
+        if (tp_part_idx == 0) {
+          std::cout << "  online quant from bf16" << std::endl;
+        }
+        int nth = T::recommended_nth_up_gate(config_.intermediate_size);
+        pool->do_work_stealing_job(
+            nth * config_.expert_num, nullptr,
+            [this, nth, physical_to_logical_map](int task_id) {
+              int64_t expert_idx = task_id / nth;
+              uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
+              int ith = task_id % nth;
+              // gate part
+              gate_bb_[logical_expert_id]->from_mat(
+                  (ggml_bf16_t*)config_.gate_proj + logical_expert_id * config_.intermediate_size * config_.hidden_size,
+                  ith, nth, -1, PACKED, PLAIN);
+              // up part
+              up_bb_[logical_expert_id]->from_mat(
+                  (ggml_bf16_t*)config_.up_proj + logical_expert_id * config_.intermediate_size * config_.hidden_size,
+                  ith, nth, -1, PACKED, PLAIN);
+            },
+            nullptr);
+
+        nth = T::recommended_nth_down(config_.hidden_size);
+        pool->do_work_stealing_job(
+            nth * config_.expert_num, nullptr,
+            [this, nth, physical_to_logical_map](int task_id) {
+              int64_t expert_idx = task_id / nth;
+              int ith = task_id % nth;
+              uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
+              // down part
+              down_bb_[logical_expert_id]->from_mat(
+                  (ggml_bf16_t*)config_.down_proj + logical_expert_id * config_.hidden_size * config_.intermediate_size,
+                  ith, nth, -1, PACKED, PLAIN);
+            },
+            nullptr);
+      }
+#ifdef CHECK
+      verify_load_right();
+#endif
+      // save process
+      if (config_.save) {
+        pool->do_work_stealing_job(
+            config_.expert_num * mat_type_all, nullptr,
+            [this, physical_to_logical_map](int task_id) {
+              int64_t expert_idx = task_id / mat_type_all;
+              expert_idx = expert_map(physical_to_logical_map, expert_idx);
+              uint8_t mat_class = task_id % mat_type_all;
+              if (mat_class == 0) {  // the up matrix
+                size_t size =
+                    T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+                size_t scale_size = config_.intermediate_size * sizeof(float);
+                write_weights(prefix, "_up_", (char*)up_bb_[expert_idx]->b_pack[0], expert_idx, size, scale_size);
+              } else if (mat_class == 1) {
+                size_t size =
+                    T::BufferB::required_size(config_.intermediate_size, config_.hidden_size, PACKED, 'u', PLAIN);
+                size_t scale_size = config_.intermediate_size * sizeof(float);
+                write_weights(prefix, "_gate_", (char*)gate_bb_[expert_idx]->b_pack[0], expert_idx, size, scale_size);
+              } else if (mat_class == 2) {
+                size_t size =
+                    T::BufferB::required_size(config_.hidden_size, config_.intermediate_size, PACKED, 'd', PLAIN);
+                size_t scale_size = config_.hidden_size * sizeof(float);
+                write_weights(prefix, "_down_", (char*)down_bb_[expert_idx]->b_pack[0], expert_idx, size, scale_size);
+              }
+            },
+            nullptr);
+      }
+    }
+  }
+
+  void warm_up() {
+    int qlen = config_.max_len;
+    std::vector<uint8_t> input(sizeof(input_t) * qlen * config_.hidden_size);
+    std::vector<uint8_t> output(sizeof(output_t) * qlen * config_.hidden_size);
+    std::vector<int64_t> expert_ids(qlen * config_.num_experts_per_tok);
+    std::vector<float> weights(qlen * config_.num_experts_per_tok);
+    for (int i = 0; i < qlen * config_.num_experts_per_tok; i++) {
+      expert_ids[i] = i % config_.expert_num;
+      weights[i] = 0.01;
+    }
+    forward(qlen, config_.num_experts_per_tok, expert_ids.data(), weights.data(), input.data(), output.data());
+  }
+
+#define MOE_DIRECT_OR_POOL_BY_VAR(var, fn)                       \
+  do {                                                           \
+    if (var < 5) {                                               \
+      for (int i = 0; i < (var); i++) {                          \
+        (fn)(i);                                                 \
+      }                                                          \
+    } else {                                                     \
+      pool->do_work_stealing_job((var), nullptr, (fn), nullptr); \
+    }                                                            \
+  } while (0)
+  static float act_fn(float x) { return x / (1.0f + expf(-x)); }
+
+  void forward(int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input, void* output) {
+    // Unified forward path: 'd' for decode (qlen<=1), 'p' for prefill (qlen>1)
+    char mode = (qlen <= 1) ? 'd' : 'p';
+    forward_unified(mode, qlen, k, expert_ids, weights, input, output);
+  }
+
+  // Helper to select B pointer for up or gate mat based on packing
+  inline int8_t* select_up_or_gate_B_ptr_(bool do_up, int expert_idx, int ith, int devide_elements_size) {
+    if constexpr (PLAIN) {
+      int8_t* base = do_up ? (int8_t*)up_bb_[expert_idx]->b : (int8_t*)gate_bb_[expert_idx]->b;
+      return base + ith * config_.hidden_size * T::N_BLOCK_UP_GATE / devide_elements_size;
+    } else {
+      return do_up ? (int8_t*)up_bb_[expert_idx]->b_pack[ith] : (int8_t*)gate_bb_[expert_idx]->b_pack[ith];
+    }
+  }
+
+  // Helper to select B pointer for down mat based on packing
+  inline int8_t* select_down_B_ptr_(int expert_idx, int ith, int devide_elements_size) {
+    if constexpr (PLAIN) {
+      return ((int8_t*)down_bb_[expert_idx]->b) +
+             ith * config_.intermediate_size * T::N_BLOCK_DOWN / devide_elements_size;
+    } else {
+      return (int8_t*)down_bb_[expert_idx]->b_pack[ith];
+    }
+  }
+
+  // Unified implementation for decode/prefill using mode 'd' or 'p'
+  void forward_unified(char mode, int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input,
+                       void* output) {
+    MatKernelVariant var = (mode == 'p') ? MatKernelVariant::Prefill : MatKernelVariant::Decode;
+    MatKernelSelection kernel = select_mat_kernel<T>(var);
+    GemmFn cblas_gemm_s8s8s32 = kernel.fn;
+    int devide_elements_size = kernel.divide_elements_size;
+
+#ifdef FORWARD_TIME_PROFILE
+    forward_perf_start();
+#endif
+    int max_local_num = 0;
+
+    auto pool = config_.pool->get_subpool(tp_part_idx);
+
+    int activated_expert = 0;
+    for (int i = 0; i < config_.expert_num; i++) {
+      m_local_num_[i] = 0;
+    }
+    for (int i = 0; i < qlen; i++) {
+      for (int j = 0; j < k; j++) {
+        if (config_.should_skip_expert(expert_ids[i * k + j])) {
+          continue;
+        }
+        m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++;
+      }
+    }
+
+    for (int i = 0; i < config_.expert_num; i++) {
+      if (m_local_num_[i] > 0) {
+        max_local_num = std::max(max_local_num, m_local_num_[i]);
+        m_expert_id_map_[activated_expert] = i;
+        activated_expert++;
+      }
+    }
+
+    uint64_t offset = 0;
+    for (int i = 0; i < config_.expert_num; i++) {
+      m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size;
+      m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size;
+      m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size;
+      m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size;
+      offset += m_local_num_[i];
+    }
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("prepare");
+#endif
+
+    // Copy inputs into expert-local buffers
+    MOE_DIRECT_OR_POOL_BY_VAR(qlen, [&](int i) {
+      for (int j = 0; j < k; j++) {
+        if (config_.should_skip_expert(expert_ids[i * k + j])) {
+          continue;
+        }
+        memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size,
+               (input_t*)input + i * config_.hidden_size, sizeof(input_t) * config_.hidden_size);
+      }
+    });
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("copy_input");
+#endif
+
+    // Quantize expert inputs (row-wise)
+    {
+      size_t mth = T::recommended_mth(max_local_num);
+      MOE_DIRECT_OR_POOL_BY_VAR(activated_expert * mth, [&](int task_id) {
+        int task_id_expert = task_id / mth;
+        int ith = task_id % mth;
+        int expert_idx = m_expert_id_map_[task_id_expert];
+        if (ith * T::M_BLOCK >= m_local_num_[expert_idx]) return;
+        gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], ith, mth);
+      });
+    }
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("quant_input");
+#endif
+
+    int nth_up = T::recommended_nth_up_gate(config_.intermediate_size, mode);
+    int mth = T::recommended_mth(max_local_num);
+    int32_t oc = 0;
+
+    // Up and Gate GEMMs + dequant scale
+    pool->do_work_stealing_job(
+        mth * nth_up * activated_expert * 2, nullptr,
+        [this, qlen, nth_up, oc, &cblas_gemm_s8s8s32, devide_elements_size, mth](int task_id2) {
+          int task_id = task_id2 / 2;
+          bool do_up = task_id2 % 2;
+          int expert_idx = m_expert_id_map_[task_id / (nth_up * mth)];
+          task_id = task_id % (nth_up * mth);
+          int ith = task_id % nth_up;
+          int jth = task_id / nth_up;
+          if (jth * T::M_BLOCK >= m_local_num_[expert_idx]) return;
+          int m_block = T::M_BLOCK;
+          if ((jth + 1) * T::M_BLOCK > m_local_num_[expert_idx]) {
+            m_block = m_local_num_[expert_idx] - jth * T::M_BLOCK;
+          }
+          int8_t* a_ptr = (int8_t*)gate_up_ba_[expert_idx]->a + jth * T::M_BLOCK * config_.hidden_size;
+          int8_t* b_ptr = select_up_or_gate_B_ptr_(do_up, expert_idx, ith, devide_elements_size);
+          int32_t* c_ptr = (do_up ? (int32_t*)up_bc_[expert_idx]->c : (int32_t*)gate_bc_[expert_idx]->c) +
+                           ith * T::N_BLOCK_UP_GATE + jth * T::M_BLOCK * config_.intermediate_size;
+
+          cblas_gemm_s8s8s32(KernelCblasRowMajor, KernelCblasNoTrans, KernelCblasTrans, KernelCblasFixOffset, m_block,
+                             T::N_BLOCK_UP_GATE, config_.hidden_size, 1.0, a_ptr, config_.hidden_size, 0, b_ptr,
+                             config_.hidden_size, 0, 0.0, c_ptr, config_.intermediate_size, &oc);
+
+          if (do_up) {
+            T::apply_scale(m_local_num_[expert_idx], config_.intermediate_size, m_local_up_output_ptr_[expert_idx],
+                           gate_up_ba_[expert_idx].get(), up_bb_[expert_idx].get(), up_bc_[expert_idx].get(), ith,
+                           nth_up, T::N_BLOCK_UP_GATE, jth);
+          } else {
+            T::apply_scale(m_local_num_[expert_idx], config_.intermediate_size, m_local_gate_output_ptr_[expert_idx],
+                           gate_up_ba_[expert_idx].get(), gate_bb_[expert_idx].get(), gate_bc_[expert_idx].get(), ith,
+                           nth_up, T::N_BLOCK_UP_GATE, jth);
+          }
+        },
+        nullptr);
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("up_gate");
+#endif
+
+    // Activate gate and multiply by up
+    {
+      int nth = T::recommended_nth(config_.intermediate_size);
+      auto up_gate_fn = [this, nth](int task_id) {
+        int expert_idx = m_expert_id_map_[task_id / nth];
+        int ith = task_id % nth;
+        auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth);
+        for (int i = 0; i < m_local_num_[expert_idx]; i++) {
+          float* gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size];
+          float* up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size];
+          for (int j = n_start; j < n_end; j++) {
+            gate_output_ptr[j] = act_fn(gate_output_ptr[j]) * up_output_ptr[j];
+          }
+        }
+      };
+      MOE_DIRECT_OR_POOL_BY_VAR(nth * activated_expert, up_gate_fn);
+    }
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("act");
+#endif
+
+    pool->do_work_stealing_job(
+        activated_expert, nullptr,
+        [this](int task_id) {
+          int expert_idx = m_expert_id_map_[task_id];
+          down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx]);
+        },
+        nullptr);
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("quant_down_input");
+#endif
+
+    int nth_down = T::recommended_nth_down(config_.hidden_size, mode);
+    pool->do_work_stealing_job(
+        mth * nth_down * activated_expert, nullptr,
+        [this, qlen, nth_down, oc, &cblas_gemm_s8s8s32, devide_elements_size, mth](int task_id) {
+          int expert_idx = m_expert_id_map_[task_id / (nth_down * mth)];
+          task_id = task_id % (nth_down * mth);
+          int ith = task_id % nth_down;
+          int jth = task_id / nth_down;
+          if (jth * T::M_BLOCK >= m_local_num_[expert_idx]) return;
+          int m_block = T::M_BLOCK;
+          if ((jth + 1) * T::M_BLOCK > m_local_num_[expert_idx]) {
+            m_block = m_local_num_[expert_idx] - jth * T::M_BLOCK;
+          }
+          int8_t* a_ptr = ((int8_t*)down_ba_[expert_idx]->a) + jth * T::M_BLOCK * config_.intermediate_size;
+          int8_t* b_ptr = select_down_B_ptr_(expert_idx, ith, devide_elements_size);
+          int32_t* c_ptr =
+              ((int32_t*)down_bc_[expert_idx]->c) + ith * T::N_BLOCK_DOWN + jth * T::M_BLOCK * config_.hidden_size;
+          cblas_gemm_s8s8s32(KernelCblasRowMajor, KernelCblasNoTrans, KernelCblasTrans, KernelCblasFixOffset, m_block,
+                             T::N_BLOCK_DOWN, config_.intermediate_size, 1.0, a_ptr, config_.intermediate_size, 0,
+                             b_ptr, config_.intermediate_size, 0, 0.0, c_ptr, config_.hidden_size, &oc);
+
+          T::apply_scale(m_local_num_[expert_idx], config_.hidden_size, m_local_down_output_ptr_[expert_idx],
+                         down_ba_[expert_idx].get(), down_bb_[expert_idx].get(), down_bc_[expert_idx].get(), ith,
+                         nth_down, T::N_BLOCK_DOWN, jth);
+        },
+        nullptr);
+
+#ifdef FORWARD_TIME_PROFILE
+    PROFILE_RECORD_TIME_STAMP("down");
+#endif
+
+    // Merge k experts per token with weights
+    size_t block_dim = 512;
+    size_t block_num = (config_.hidden_size + block_dim - 1) / block_dim;
+    pool->do_work_stealing_job(
+        qlen * block_num, nullptr,
+        [this, k, expert_ids, weights, output, block_dim, block_num](int i) {
+          int q_idx = i / block_num;
+          int block_idx = i % block_num;
+          int e_start = block_idx * block_dim;
+          int e_end =
+              ((block_idx + 1) * block_dim) < config_.hidden_size ? ((block_idx + 1) * block_dim) : config_.hidden_size;
+          for (int e = e_start; e < e_end; e++) {
+            float sum = 0;
+            for (int j = 0; j < k; j++) {
+              if (config_.should_skip_expert(expert_ids[q_idx * k + j])) {
+                continue;
+              }
+              sum += weights[q_idx * k + j] * ((float*)m_local_down_output_ptr_[expert_ids[q_idx * k + j]])
+                                                  [m_local_pos_[q_idx][j] * config_.hidden_size + e];
+            }
+            ((float*)output)[q_idx * config_.hidden_size + e] = sum;
+          }
+        },
+        nullptr);
+
+#ifdef FORWARD_TIME_PROFILE
+    time_perf_name = std::string("[moe] ") + ((mode == 'p') ? "layer prefill" : "decode layer ") +
+                     std::to_string(config_.layer_idx) + " tp_part_idx: " + std::to_string(tp_part_idx);
+    perf_report();
+#endif
+  }
+
+  /* merged into forward_unified */
+  void forward_decode(int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input,
+                      void* output) {
+    forward_unified('d', qlen, k, expert_ids, weights, input, output);
+  }
+
+  void forward_prefill(int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input,
+                       void* output) {
+    forward_unified('p', qlen, k, expert_ids, weights, input, output);
+  }
+};
+
+template <typename K, bool T>
+class TP_MOE<MOE_KERNEL_TP<K, T>> : public TP_MOE_Common<MOE_KERNEL_TP<K, T>> {
+ public:
+  using TP_MOE_Common<MOE_KERNEL_TP<K, T>>::TP_MOE_Common;
+
+  void load_weights() {
+    auto& config = this->config;
+    auto& tps = this->tps;
+    auto& tp_count = this->tp_count;
+    auto pool = config.pool;
+    const uint64_t* physical_to_logical_map = (const uint64_t*)config.physical_to_logical_map;
+    if (config.gate_projs.empty() == false) {
+      printf("TP Load from loader\n");
+      pool->dispense_backend()->do_numa_job([this, pool](int numa_id) { this->tps[numa_id]->load_weights(); });
+      this->weights_loaded = true;
+    } else if (config.gate_proj != nullptr) {
+      printf("From BF16\n");
+      for (auto i = 0; i < tp_count; i++) {
+        auto& tpc = tps[i]->config_;
+        size_t gate_up_elcount = tpc.intermediate_size * tpc.hidden_size;
+        tpc.gate_proj = new ggml_bf16_t[tpc.expert_num * gate_up_elcount];
+        tpc.up_proj = new ggml_bf16_t[tpc.expert_num * gate_up_elcount];
+        tpc.down_proj = new ggml_bf16_t[tpc.expert_num * gate_up_elcount];
+        if (tps[i]->config_.load == false) {
+          pool->get_subpool(i)->do_work_stealing_job(
+              tpc.expert_num, nullptr,
+              [&](int expert_id_) {
+                size_t expert_id = expert_map(physical_to_logical_map, expert_id_);
+                memcpy((ggml_bf16_t*)tpc.gate_proj + expert_id * gate_up_elcount,
+                       (ggml_bf16_t*)config.gate_proj + expert_id * config.intermediate_size * config.hidden_size +
+                           i * gate_up_elcount,
+                       sizeof(ggml_bf16_t) * gate_up_elcount);
+                memcpy((ggml_bf16_t*)tpc.up_proj + expert_id * gate_up_elcount,
+                       (ggml_bf16_t*)config.up_proj + expert_id * config.intermediate_size * config.hidden_size +
+                           i * gate_up_elcount,
+                       sizeof(ggml_bf16_t) * gate_up_elcount);
+                for (size_t col = 0; col < config.hidden_size; col++) {
+                  memcpy((ggml_bf16_t*)tpc.down_proj + expert_id * tpc.hidden_size * tpc.intermediate_size +
+                             col * tpc.intermediate_size,
+                         (ggml_bf16_t*)config.down_proj + expert_id * config.intermediate_size * config.hidden_size +
+                             col * config.intermediate_size + i * tpc.intermediate_size,
+                         sizeof(ggml_bf16_t) * tpc.intermediate_size);
+                }
+              },
+              nullptr);
+        }
+      }
+
+      pool->dispense_backend()->do_numa_job([this, pool](int numa_id) { this->tps[numa_id]->load_weights(); });
+
+      for (auto i = 0; i < tp_count; i++) {
+        auto& tpc = tps[i]->config_;
+        delete[] (ggml_bf16_t*)(tpc.gate_proj);
+        delete[] (ggml_bf16_t*)(tpc.up_proj);
+        delete[] (ggml_bf16_t*)(tpc.down_proj);
+      }
+      if (config.save) {
+        // free the bf16 weights after saving
+        tps.clear();
+      }
+
+      this->weights_loaded = true;
+    } else if (config.path != "") {
+      printf("TP Load from file\n");
+      pool->dispense_backend()->do_numa_job([this, pool](int numa_id) { this->tps[numa_id]->load_weights(); });
+      this->weights_loaded = true;
+    } else {
+      throw std::runtime_error("no weight source");
+    }
+  }
+
+  void merge_results(int qlen, void* output, bool incremental) {
+    // #ifdef FORWARD_TIME_PROFILE
+    //     forward_perf_start();
+    // #endif
+    auto pool = this->config.pool;
+    auto merge_fn = [this, output, incremental](int token_nth) {
+      auto& local_output_numa = this->local_output_numa;
+      auto& tp_configs = this->tp_configs;
+      auto& tp_count = this->tp_count;
+      auto& config = this->config;
+      float* merge_to = local_output_numa[0] + token_nth * tp_configs[0].hidden_size;
+      if (incremental) {
+        for (int e = 0; e < config.hidden_size; e++) {
+          merge_to[e] += ggml_bf16_to_fp32(((ggml_bf16_t*)output + token_nth * config.hidden_size)[e]);
+        }
+      }
+
+      for (int i = 1; i < tp_count; i++) {
+        float* merge_from = local_output_numa[i] + token_nth * tp_configs[i].hidden_size;
+        // TODO: 后续用 SVE 来加速
+        // for (int e = 0; e < tp_configs[i].hidden_size; e += 16) {
+        //   *((__m512 *)(merge_to + e)) = _mm512_add_ps(*((__m512 *)(merge_to + e)), *((__m512 *)(merge_from + e)));
+        // }
+        // CHECK: 目前用普通的纯 C++ 来实现
+        for (int e = 0; e < tp_configs[i].hidden_size; e++) {
+          merge_to[e] += merge_from[e];
+        }
+      }
+
+      convert_or_copy((ggml_bf16_t*)output + token_nth * config.hidden_size, merge_to, config.hidden_size);
+
+      // for (int e = 0; e < config.hidden_size; e += 32) {
+      // TODO: 这里需要用 SVE 来加速，实现 fp32 到 bf16 的转换
+      // __m512 x0 = *(__m512 *)(merge_to + e);
+      // __m512 x1 = *(__m512 *)(merge_to + e + 16);
+      // avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + token_nth * config.hidden_size + e));
+
+      // CHECK: 目前用普通的纯 C++ 来实现 fp32 到 bf16 的转换
+
+      // convert_32fp32_to_32bf16_pure_c(merge_to + e,
+      // (uint16_t *)((ggml_bf16_t *)output + token_nth * config.hidden_size + e));
+
+      // }
+    };
+    MOE_DIRECT_OR_POOL_BY_VAR(qlen, merge_fn);
+    // #ifdef FORWARD_TIME_PROFILE
+    //     PROFILE_RECORD_TIME_STAMP("moe merge done");
+    // #endif
+    // #ifdef FORWARD_TIME_PROFILE
+    //     time_perf_name = "[moe merge] decode layer " + std::to_string(this->config.layer_idx);
+    //     perf_report();
+    // #endif
+  }
+
+  void merge_results(int qlen, void* output) { merge_results(qlen, output, false); }
+};
+
+#endif

kt-kernel/operators/moe_kernel/test/convert-test.cpp

@@ -0,0 +1,88 @@
+#include <arm_sve.h>
+
+#include <cmath>
+#include <cstdlib>
+#include <fstream>
+#include <stdexcept>
+
+#include "../../reduce.hpp"
+#include "../../rms-norm.hpp"
+#include "../../rope.hpp"
+#include "../../softmax.hpp"
+#include "../la/arm_kml.hpp"
+#include "llama.cpp/ggml-common.h"
+#include "llama.cpp/ggml.h"
+
+void bf16_to_fp16(const ggml_bf16_t* src, ggml_fp16_t* dst, size_t n) {
+  for (size_t i = 0; i < n; ++i) {
+    float x = ggml_bf16_to_fp32(src[i]);
+    dst[i] = ggml_fp32_to_fp16(x);
+  }
+}
+
+void debug_rope() {
+  float16_t* fp16 = new float16_t[1024 * 64];
+
+  for (size_t i = 0; i < 1024 * 64; i++) {
+    fp16[i] = static_cast<double>(std::rand()) / RAND_MAX;
+  }
+  std::ofstream("before_rope", std::ios::binary).write((char*)fp16, 1024 * 64 * sizeof(float16_t));
+
+  DeepseekV3YarnRotaryEmbedding rope(64, 163840, 10000, 40, 4096, 32, 1, 1, 1);
+
+  rope.init(1024);
+
+  Rope<DeepseekV3YarnRotaryEmbedding, float16_t> rope_applier;
+  rope_applier.apply_multiple(rope, fp16, 64, 64, 0, 1024);
+
+  std::ofstream("cos", std::ios::binary).write((char*)rope.cos(0), 1024 * 32 * sizeof(float));
+  std::ofstream("sin", std::ios::binary).write((char*)rope.sin(0), 1024 * 32 * sizeof(float));
+
+  std::ofstream("after_rope", std::ios::binary).write((char*)fp16, 1024 * 64 * sizeof(float16_t));
+}
+
+void debug_softmax() {
+  float16_t* fp16 = new float16_t[64 * 1024];
+
+  for (size_t i = 0; i < 1024 * 64; i++) {
+    fp16[i] = static_cast<double>(std::rand()) / RAND_MAX * 10;
+    if (i % 12 == 0) {
+      fp16[i] -= std::numeric_limits<float16_t>::infinity();
+    }
+  }
+  std::ofstream("before_softmax", std::ios::binary).write((char*)fp16, 1024 * 64 * sizeof(float16_t));
+
+  Softmax<float16_t>::apply_multiple(64, fp16, 1024, 1024);
+  std::ofstream("after_softmax", std::ios::binary).write((char*)fp16, 1024 * 64 * sizeof(float16_t));
+}
+
+void debug_inf() {
+  float16_t x, y;
+  // x = std::numeric_limits<float16_t>::infinity(); // 0.00
+  // y = -std::numeric_limits<float16_t>::infinity(); // -0.00
+  // x = 1e10;
+  x = std::numeric_limits<float>::infinity();   // inf
+  y = -std::numeric_limits<float>::infinity();  // -inf
+  printf("x = %f, y = %f\n", x, y);
+}
+
+void debug_reduce() {
+  std::vector<float16_t*> fp16s(128);
+  for (size_t i = 0; i < 128; i++) {
+    fp16s[i] = new float16_t[1024];
+    for (size_t j = 0; j < 1024; j++) {
+      fp16s[i][j] = i;
+    }
+  }
+
+  reduce_sum(fp16s.data(), 128, 0, 10);
+  for (int i = 0; i < 10; i++) {
+    printf("%f ", fp16s[0][i]);
+  }
+}
+
+int main() {
+  debug_reduce();
+
+  return 0;
+}

kt-kernel/operators/moe_kernel/test/debug.hpp

@@ -0,0 +1,59 @@
+#ifndef KML_DEBUG_HPP
+#define KML_DEBUG_HPP
+
+#include <arm_sve.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <fstream>
+#include <string>
+
+inline std::string get_env_or_default(const char* var_name, const std::string& default_value) {
+  const char* value = std::getenv(var_name);
+  return (value != nullptr) ? std::string(value) : default_value;
+}
+
+inline void dump_bin(std::string file_name, float16_t* data, size_t count) {
+  file_name = get_env_or_default("KML_DEBUG_PATH", "debug") + "/" + file_name + ".f16";
+  std::ofstream f(file_name, std::ios::binary);
+  f.write(reinterpret_cast<const char*>(data), count * sizeof(*data));
+  f.close();
+}
+inline void dump_bin(std::string file_name, float* data, size_t count) {
+  file_name = get_env_or_default("KML_DEBUG_PATH", "debug") + "/" + file_name + ".f32";
+  std::ofstream f(file_name, std::ios::binary);
+  f.write(reinterpret_cast<const char*>(data), count * sizeof(*data));
+  f.close();
+}
+inline void dump_bin(std::string file_name, int64_t* data, size_t count) {
+  file_name = get_env_or_default("KML_DEBUG_PATH", "debug") + "/" + file_name + ".int64";
+  std::ofstream f(file_name, std::ios::binary);
+  f.write(reinterpret_cast<const char*>(data), count * sizeof(*data));
+  f.close();
+}
+
+inline void dump_bin(std::string file_name, int8_t* data, size_t count) {
+  file_name = get_env_or_default("KML_DEBUG_PATH", "debug") + "/" + file_name + ".int8";
+  std::ofstream f(file_name, std::ios::binary);
+  f.write(reinterpret_cast<const char*>(data), count * sizeof(*data));
+  f.close();
+}
+
+inline void dump_bin(std::string file_name, int32_t* data, size_t count) {
+  file_name = get_env_or_default("KML_DEBUG_PATH", "debug") + "/" + file_name + ".int32";
+  std::ofstream f(file_name, std::ios::binary);
+  f.write(reinterpret_cast<const char*>(data), count * sizeof(*data));
+  f.close();
+}
+
+inline void load_bin(std::string file_name, float* data, size_t count) {
+  file_name = get_env_or_default("KML_DEBUG_PATH", "debug") + "/" + file_name + ".f32";
+  std::ifstream f(file_name, std::ios::binary);
+  if (!f.is_open()) {
+    throw std::runtime_error("Failed to open file: " + file_name);
+  }
+  f.read(reinterpret_cast<char*>(data), count * sizeof(*data));
+  f.close();
+}
+
+#endif

kt-kernel/operators/moe_kernel/test/int4_mul-test.cpp

@@ -0,0 +1,110 @@
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+
+#include "../la/arm_kml.hpp"
+#include "debug.hpp"
+#include "kblas.h"
+const int M = 1, K = 7168, N = 8;
+
+int main() {
+  // 随机生成a, b, c矩阵
+  arm_kml::GemmKernelInt4::BufferA buffer_a(M, K);
+  arm_kml::GemmKernelInt4::BufferB buffer_b(N, K, true);
+  arm_kml::GemmKernelInt4::BufferC buffer_c(M, N);
+
+  arm_kml::GemmKernelInt8::BufferA buffer_a_check(M, K);
+  arm_kml::GemmKernelInt8::BufferB buffer_b_check(N, K, true);
+  arm_kml::GemmKernelInt8::BufferC buffer_c_check(M, N);
+
+  float* a = (float*)aligned_alloc(64, sizeof(float) * M * K);
+  float* b = (float*)aligned_alloc(64, sizeof(float) * K * N);
+  float* c = (float*)aligned_alloc(64, sizeof(float) * M * N);
+  float* c_check = (float*)aligned_alloc(64, sizeof(float) * M * N);
+  int8_t* buffer_a_data = (int8_t*)aligned_alloc(64, buffer_a.required_size());
+  int4_2_t* buffer_b_data = (int4_2_t*)aligned_alloc(64, buffer_b.required_size());
+  int32_t* c_data = (int32_t*)aligned_alloc(64, buffer_c.required_size());
+  int8_t* buffer_a_data_check = (int8_t*)aligned_alloc(64, buffer_a_check.required_size());
+  int8_t* buffer_b_data_check = (int8_t*)aligned_alloc(64, buffer_b_check.required_size());
+  int32_t* c_data_check = (int32_t*)aligned_alloc(64, buffer_c_check.required_size());
+  // 初始化元素内容
+  load_bin("input.bin", a, M * K);
+  load_bin("local_q_a_proj_quant.bin", b, N * K);
+
+  // for (int i = 0; i < M * K; i++) {
+  //   // 随机浮点数
+  //   // a[i] = (static_cast<float>(rand()) / (float)RAND_MAX) / 25 - 0.02;
+  //   a[i] = -(static_cast<float>(rand()) / (float)RAND_MAX) / 25;
+  //   // a[i] = i % 10;
+  //   // a[i] = 1;
+  // }
+  // for (int i = 0; i < K * N; i++) {
+  //   // 随机浮点数
+  //   // b[i] = (static_cast<float>(rand()) / (float)RAND_MAX) / 25 - 0.02;
+  //   b[i] = -(static_cast<float>(rand()) / (float)RAND_MAX) / 25;
+  //   // b[i] = i % 10;
+  //   // b[i] = 1;
+  // }
+  // // // // 设置离群值
+  // for (int i = 0; i < N; i++) {
+  //   b[i * K] = 0.06f; // 设置第一列为离群值
+  // }
+  // // 打印一下输入矩阵和权重矩阵
+  // printf("Input matrix a:\n");
+  // for (int i = 0; i < M; i++) {
+  //   for (int j = 0; j < K; j++) {
+  //     printf("%f ", a[i * K + j]);
+  //   }
+  //   printf("\n");
+  // }
+  // printf("Weight matrix b:\n");
+  // for (int i = 0; i < N; i++) {
+  //   for (int j = 0; j < K; j++) {
+  //     printf("%f ", b[i * K + j]);
+  //   }
+  //   printf("\n");
+  // }
+  buffer_a.set_data(buffer_a_data);
+  buffer_b.set_data(buffer_b_data);
+  buffer_c.set_data(c_data);
+  buffer_a_check.set_data(buffer_a_data_check);
+  buffer_b_check.set_data(buffer_b_data_check);
+  buffer_c_check.set_data(c_data_check);
+  //   调用 from mat 进行量化
+  buffer_a.from_mat(M, a, 0, M);
+  for (int i = 0; i <= arm_kml::GemmKernelInt4::recommended_nth(N); i++) {
+    buffer_b.from_mat(b, i, arm_kml::GemmKernelInt4::recommended_nth(N));
+  }
+  buffer_a_check.from_mat(M, a, 0, M);
+  for (int i = 0; i <= arm_kml::GemmKernelInt8::recommended_nth(N); i++) {
+    buffer_b_check.from_mat(b, i, arm_kml::GemmKernelInt8::recommended_nth(N));
+  }
+  // 进行乘法
+  arm_kml::MatRef<int8_t> a_ref(buffer_a.a, M, K, K, CblasRowMajor);
+  arm_kml::MatRef<int4_2_t> b_ref(buffer_b.b, K, N, K, CblasColMajor, CblasNoTrans, buffer_b.if_pack);
+  arm_kml::MatRef<int32_t> c_ref(buffer_c.c, M, N, N, CblasRowMajor);
+  b_ref = b_ref.offset_col(0, N);
+
+  arm_kml::MatRef<int8_t> a_ref_check(buffer_a_check.a, M, K, K, CblasRowMajor);
+  arm_kml::MatRef<int8_t> b_ref_check(buffer_b_check.b, K, N, K, CblasColMajor, CblasNoTrans, buffer_b_check.if_pack);
+  arm_kml::MatRef<int32_t> c_ref_check(buffer_c_check.c, M, N, N, CblasRowMajor);
+
+  arm_kml::decode_mul_mat_clearc(a_ref, b_ref, c_ref);
+  arm_kml::decode_mul_mat_clearc(a_ref_check, b_ref_check, c_ref_check);
+  //   反量化，apply scale
+  arm_kml::GemmKernelInt4::apply_scale(c, N, &buffer_a, &buffer_b, &buffer_c, 0, M, 0, N, true);
+  arm_kml::GemmKernelInt8::apply_scale(c_check, N, &buffer_a_check, &buffer_b_check, &buffer_c_check, 0, M, 0, N, true);
+  // 打印结果,比较 c 和 c_check
+  const float threashold = 0.05;
+  for (int i = 0; i < M * N; i++) {
+    float diff_relative = (c[i] - c_check[i]) / (c_check[i] + 1e-6);
+
+    if (diff_relative > threashold || diff_relative < -threashold) {
+      printf("diff_relative: %f\n", diff_relative);
+      printf("Mismatch at index %d: c = %f, c_check = %f\n", i, c[i], c_check[i]);
+    } else {
+      printf("Match at index %d: c = %f, c_check = %f\n", i, c[i], c_check[i]);
+    }
+  }
+  return 0;
+}

kt-kernel/operators/moe_kernel/test/mat_test.cpp

@@ -0,0 +1,22 @@
+#include "arm_kml.hpp"
+
+int main() {
+  const size_t M = 128, N = 64;
+  float16_t* a = new float16_t[M * N];
+  float16_t* b = new float16_t[M * N];
+  float16_t* c = new float16_t[M * M];
+  float16_t* c_check = new float16_t[M * M];
+  for (size_t i = 0; i < M * N; i++) {
+    a[i] = static_cast<double>(std::rand()) / RAND_MAX / 10.0;
+    b[i] = static_cast<double>(std::rand()) / RAND_MAX / 10.0;
+  }
+
+  arm_kml::MatRef<float16_t> aref(a, M, N, M, CblasColMajor);
+  arm_kml::MatRef<float16_t> bref(b, N, M, M, CblasColMajor);
+  arm_kml::MatRef<float16_t> cref(c, M, M, M, CblasColMajor);
+  {
+    memset(c, 0, M * M * sizeof(float16_t));
+    memset(c_check, 0, M * M * sizeof(float16_t));
+    arm_kml::mul_mat(aref, bref, cref);
+  }
+}

kt-kernel/operators/moe_kernel/test/utils_test.cpp

@@ -0,0 +1,63 @@
+// #pragma once
+#ifdef TEST_UTIL
+#include <arm_neon.h>
+#include <arm_sve.h>
+#include <stdio.h>
+
+static inline void sve_32xbf16_to_32xfp32(const bfloat16_t* src, float* dst0, float* dst1) {
+#ifdef __ARM_FEATURE_SVE
+  // 全真谓词，对应每个 16‑bit 元素
+#else
+// fallback: scalar or NEON
+#endif
+}
+
+static inline void neon_32xbf16_to_32xfp32(const uint16_t* src, float* dst0, float* dst1) {
+  // src 指向 32 个连续的 BF16（uint16_t）
+  // dst0、dst1 各指向 16 个 float 的缓冲
+
+  for (int block = 0; block < 4; ++block) {
+    // 每次处理 8 个 BF16 → 8 个 FP32（拆为两次 4→4 存储）
+    uint16x8_t v_bf16 = vld1q_u16(src + block * 8);  // load 8×BF16 :contentReference[oaicite:6]{index=6}
+
+    // 拆低半、高半各 4 个到 u32
+    uint32x4_t lo_u32 = vmovl_u16(vget_low_u16(v_bf16));   // lower 4 → u32 :contentReference[oaicite:7]{index=7}
+    uint32x4_t hi_u32 = vmovl_u16(vget_high_u16(v_bf16));  // upper 4 → u32 :contentReference[oaicite:8]{index=8}
+
+    // 左移 16 位，相当于将 BF16 的 16 位 mantissa+exp 放到 FP32 高位
+    lo_u32 = vshlq_n_u32(lo_u32, 16);  // shift left 16 :contentReference[oaicite:9]{index=9}
+    hi_u32 = vshlq_n_u32(hi_u32, 16);  // shift left 16 :contentReference[oaicite:10]{index=10}
+
+    // 重新解释为 float32x4_t
+    float32x4_t lo_f32 = vreinterpretq_f32_u32(lo_u32);  // bits → FP32 :contentReference[oaicite:11]{index=11}
+    float32x4_t hi_f32 = vreinterpretq_f32_u32(hi_u32);  // bits → FP32 :contentReference[oaicite:12]{index=12}
+
+    // 存储到 dst0 或 dst1，每次存 8 个
+    if (block < 2) {
+      vst1q_f32(dst0 + block * 4, lo_f32);      // store 4 floats :contentReference[oaicite:13]{index=13}
+      vst1q_f32(dst0 + block * 4 + 4, hi_f32);  // store next 4 floats :contentReference[oaicite:14]{index=14}
+    } else {
+      int b = block - 2;
+      vst1q_f32(dst1 + b * 4, lo_f32);      // store 4 floats :contentReference[oaicite:15]{index=15}
+      vst1q_f32(dst1 + b * 4 + 4, hi_f32);  // store next 4 floats :contentReference[oaicite:16]{index=16}
+    }
+  }
+}
+
+int main() {
+  // 测试代码
+  uint16_t bf16_data[32] = {0};  // 假设这里填充了一些 BF16 数据
+  float f32_data0[16] = {0};
+  float f32_data1[16] = {0};
+
+  neon_32xbf16_to_32xfp32(bf16_data, f32_data0, f32_data1);
+
+  // 打印结果
+  for (int i = 0; i < 16; ++i) {
+    printf("f32_data0[%d]: %f\n", i, f32_data0[i]);
+    printf("f32_data1[%d]: %f\n", i, f32_data1[i]);
+  }
+
+  return 0;
+}
+#endif