diff --git a/onnxruntime/core/mlas/lib/q4_dq.cpp b/onnxruntime/core/mlas/lib/q4_dq.cpp
index cefe98930c832..dc5efeb5814c6 100644
--- a/onnxruntime/core/mlas/lib/q4_dq.cpp
+++ b/onnxruntime/core/mlas/lib/q4_dq.cpp
@@ -152,13 +152,32 @@ MlasJblasNBitsGemmPackB(void* PackedBuf,
                         MLAS_COMPUTE_TYPE CompType,
                         MLAS_THREADPOOL* ThreadPool)
 {
+    GetCPUDevice();
     switch (CompType) {
         case CompInt8:
-            return JblaNBitsGemmPackB(JblasAvxVnniS4Fp32Fp32, PackedBuf, int(BlkSize), QData, Scale,
-                                      Zp, int(N), int(K), isAsym, lastCall, int(ldb), ThreadPool);
+            if (_cd->AVX512_VNNI()) {
+                return JblaNBitsGemmPackB(JblasAvx512VnniS4Fp32Fp32, PackedBuf, int(BlkSize), QData,
+                                          Scale, Zp, int(N), int(K), isAsym, lastCall, int(ldb),
+                                          ThreadPool);
+            }
+            if (_cd->AVX_VNNI()) {
+                return JblaNBitsGemmPackB(JblasAvxVnniS4Fp32Fp32, PackedBuf, int(BlkSize), QData,
+                                          Scale, Zp, int(N), int(K), isAsym, lastCall, int(ldb),
+                                          ThreadPool);
+            }
+            break;
         case CompFp32:
-            return JblaNBitsGemmPackB(JblasAvx512fS4Fp32Fp32, PackedBuf, int(BlkSize), QData, Scale,
-                                      Zp, int(N), int(K), isAsym, lastCall, int(ldb), ThreadPool);
+            if (_cd->AVX512F()) {
+                return JblaNBitsGemmPackB(JblasAvx512fS4Fp32Fp32, PackedBuf, int(BlkSize), QData,
+                                          Scale, Zp, int(N), int(K), isAsym, lastCall, int(ldb),
+                                          ThreadPool);
+            }
+            if (_cd->AVX2()) {
+                return JblaNBitsGemmPackB(JblasAvx2S4Fp32Fp32, PackedBuf, int(BlkSize), QData,
+                                          Scale, Zp, int(N), int(K), isAsym, lastCall, int(ldb),
+                                          ThreadPool);
+            }
+            break;
         case CompBf16:
         case CompFp16:
         default:
@@ -177,18 +196,37 @@ MlasJblasQ4GemmUnPackB(float* FpData,
     auto ptr =
         jblas::storage::gemm::PackedWeightParser::deserialBuffer(const_cast<void*>(PackedBuf));
     ORTThreading orth(ThreadPool);
+    GetCPUDevice();
     if (ptr) {
         if (ptr->mPrologueID == JBLAS_PROLOGUEB_IDS::WeightKBlockS4) {
-            auto coretype = ptr->mCoreType;
-            auto NTile = uint32_t(coretype) & uint32_t(JBLAS_GEMM_CORE::NTILE_MASK);
-            auto CType = uint32_t(coretype) & uint32_t(JBLAS_GEMM_CORE::COMP_MASK);
-            if (NTile == 48 && CType == uint32_t(JBLAS_GEMM_CORE::COMP_FP32)) {
-                JblasAvx512fS4Fp32Fp32.mProB.unpackWeight(int(N), int(K), ptr, FpData, int(ldb),
-                                                          &orth);
+            auto NTile =
+                jblas::gemm::CoreAttr::get_mask_val(ptr->mCoreId, jblas::gemm::CoreAttr::NTILE_MASK,
+                                                    jblas::gemm::CoreAttr::NTILE_SHIFT);
+            auto CType = jblas::gemm::CoreAttr::get_mask_val(
+                ptr->mCoreId, jblas::gemm::CoreAttr::COMP_MASK, jblas::gemm::CoreAttr::COMP_SHIFT);
+            if (CType == uint32_t(jblas::gemm::CompType::COMP_FP32)) {
+                if (NTile == 48 && _cd->AVX512F()) {
+                    JblasAvx512fS4Fp32Fp32.mProB.unpackWeight(int(N), int(K), ptr, FpData, int(ldb),
+                                                              &orth);
+                    return;
+                }
+                if (NTile == 24 && _cd->AVX2()) {
+                    JblasAvx2S4Fp32Fp32.mProB.unpackWeight(int(N), int(K), ptr, FpData, int(ldb),
+                                                           &orth);
+                    return;
+                }
             }
-            if (NTile == 48 && CType == uint32_t(JBLAS_GEMM_CORE::COMP_INT8_US)) {
-                JblasAvx512VnniS4Fp32Fp32.mProB.unpackWeight(int(N), int(K), ptr, FpData, int(ldb),
-                                                             &orth);
+            if (CType == uint32_t(jblas::gemm::CompType::COMP_INT8_US_INT32)) {
+                if (NTile == 48 && _cd->AVX512_VNNI()) {
+                    JblasAvx512VnniS4Fp32Fp32.mProB.unpackWeight(int(N), int(K), ptr, FpData,
+                                                                 int(ldb), &orth);
+                    return;
+                }
+                if (NTile == 24 && _cd->AVX_VNNI()) {
+                    JblasAvxVnniS4Fp32Fp32.mProB.unpackWeight(int(N), int(K), ptr, FpData, int(ldb),
+                                                              &orth);
+                    return;
+                }
             }
         }
         delete ptr;
diff --git a/onnxruntime/core/mlas/lib/q4gemm.cpp b/onnxruntime/core/mlas/lib/q4gemm.cpp
index b003038d3b9fa..891c346c14466 100644
--- a/onnxruntime/core/mlas/lib/q4gemm.cpp
+++ b/onnxruntime/core/mlas/lib/q4gemm.cpp
@@ -155,30 +155,6 @@ jblas::ORTThreading::parallel_for(const jblas::parallel::thread_func& func)
                           [&](ptrdiff_t tid) { func(int(tid)); });
 }
 
-template <class Parallel_T, class Launch_T>
-void
-GemmKBlockRun(Launch_T& launcher,
-              const typename Launch_T::Param& args,
-              parallel::IThreading* threading)
-{
-    device::CpuBase cb;
-    Parallel_T para({
-        threading->num_threads(),
-        cb.mL2Cache,
-        args.M,
-        args.N,
-        args.K,
-        args.KBlock,
-    });
-    threading->parallel_for([&](int tidx) {
-        typename Parallel_T::ThreadProblem thdp{tidx};
-        para.getIndex(thdp);
-        if (thdp.valid) {
-            launcher.run(args, thdp);
-        }
-    });
-}
-
 template <class GemmCore_T>
 void
 JblasQ4GemmCompF32(const int M,
@@ -205,7 +181,7 @@ JblasQ4GemmCompF32(const int M,
         reduceA.template get<float>(), reduceA.lda};
 
     typename Launcher::Param args{M, N, K, B->mBlockSize, {A, K}, {B}, blkargs, {C, N}};
-    GemmKBlockRun<Parallel>(kernel, args, th);
+    jblas::parallel::GemmKBlockRun<Parallel>(kernel, args, th);
 }
 
 template <class GemmCore_T>
@@ -236,10 +212,10 @@ JblasQ4GemmCompInt8(const int M,
         {A, K, &quanA},
         {B},
         {B->template SPtr<int8_t>(), B->mScaT, B->mCStep, quanA.template SPtr<float>(),
-         quanA.mCStep, quanA.template ZPtr<uint8_t>(), B->template RPtr<float>(),
+         quanA.mCStep, quanA.template ZPtr<uint8_t>(), B->template RPtr<float>(), B->mRedT,
          B->template ZPtr<int8_t>(), quanA.template RPtr<float>(), B->mBlockSize},
         {C, N}};
-    GemmKBlockRun<Parallel>(kernel, args, th);
+    jblas::parallel::GemmKBlockRun<Parallel>(kernel, args, th);
 }
 
 void
@@ -259,35 +235,39 @@ JblasQ4GemmBatchDriver(const size_t M,
         if (ptr) {
             if (ptr->mPrologueID == JBLAS_PROLOGUEB_IDS::WeightKBlockS4) {
                 auto kptr = (jblas::storage::gemm::StorageWeightKBlockS4*)ptr;
-                auto coretype = ptr->mCoreType;
-                auto NTile = uint32_t(coretype) & uint32_t(JBLAS_GEMM_CORE::NTILE_MASK);
-                auto CType = uint32_t(coretype) & uint32_t(JBLAS_GEMM_CORE::COMP_MASK);
-                if (NTile == 48 && CType == uint32_t(JBLAS_GEMM_CORE::COMP_FP32)) {
-                    if (_cd->AVX512F()) {
-                        JblasQ4GemmCompF32<gemm::GemmCore_Row_NN_8x48_AVX512F>(
+                auto coretype = ptr->mCoreId;
+                auto NTile = jblas::gemm::CoreAttr::get_mask_val(
+                    ptr->mCoreId, jblas::gemm::CoreAttr::NTILE_MASK,
+                    jblas::gemm::CoreAttr::NTILE_SHIFT);
+                auto CType = jblas::gemm::CoreAttr::get_mask_val(ptr->mCoreId,
+                                                                 jblas::gemm::CoreAttr::COMP_MASK,
+                                                                 jblas::gemm::CoreAttr::COMP_SHIFT);
+                if (CType == uint32_t(gemm::CompType::COMP_FP32)) {
+                    if (NTile == 48 && _cd->AVX512F()) {
+                        JblasQ4GemmCompF32<gemm::SCoreRowNAvx512f<48, 8>>(
                             M, N, K, DataParams[i].A, DataParams[i].lda,
                             (jblas::storage::gemm::StorageWeightKBlockS4*)ptr, DataParams[i].C,
                             DataParams[i].ldc, WorkSpace, &orth);
                         return;
                     }
-                    if (_cd->AVX2()) {
-                        JblasQ4GemmCompF32<gemm::GemmCore_Row_NN_2x48_AVX2>(
+                    if (NTile == 24 && _cd->AVX2()) {
+                        JblasQ4GemmCompF32<gemm::SCoreRowNAvx2<24, 4>>(
                             M, N, K, DataParams[i].A, DataParams[i].lda,
                             (jblas::storage::gemm::StorageWeightKBlockS4*)ptr, DataParams[i].C,
                             DataParams[i].ldc, WorkSpace, &orth);
                         return;
                     }
                 }
-                if (NTile == 48 && CType == uint32_t(JBLAS_GEMM_CORE::COMP_INT8_US)) {
-                    if (_cd->AVX512_VNNI()) {
-                        JblasQ4GemmCompInt8<gemm::GemmCore_Row_NN_8x48_AVX512_VNNI>(
+                if (CType == uint32_t(gemm::CompType::COMP_INT8_US_INT32)) {
+                    if (NTile == 48 && _cd->AVX512_VNNI()) {
+                        JblasQ4GemmCompInt8<gemm::ICoreRowNAvx512vnni<48, 8>>(
                             M, N, K, DataParams[i].A, DataParams[i].lda,
                             (jblas::storage::gemm::StorageWeightKBlockS4*)ptr, DataParams[i].C,
                             DataParams[i].ldc, WorkSpace, &orth);
                         return;
                     }
-                    if (_cd->AVX_VNNI()) {
-                        JblasQ4GemmCompInt8<gemm::GemmCore_Row_NN_2x48_AVX_VNNI>(
+                    if (NTile == 24 && _cd->AVX_VNNI()) {
+                        JblasQ4GemmCompInt8<gemm::ICoreRowNAvxvnni<24, 4>>(
                             M, N, K, DataParams[i].A, DataParams[i].lda,
                             (jblas::storage::gemm::StorageWeightKBlockS4*)ptr, DataParams[i].C,
                             DataParams[i].ldc, WorkSpace, &orth);
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_base.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_base.h
index d2ef3308fdccd..c2f08a0fc9770 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_base.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_base.h
@@ -49,6 +49,21 @@ class JitBase : protected Xbyak::CodeGenerator {
 #endif
   }
 
+  void padto_le(const Xbyak::Reg64& _src, int padding) {
+    // _src=_src/padding*padding
+    if (padding == 1) {
+      return;
+    }
+    for (int i = 1; i < 16; i++) {
+      if ((1 << i) == padding) {
+        shr(_src, i);
+        shl(_src, i);
+        return;
+      }
+    }
+    assert(0);
+  }
+
   void generate_Nbitsmask(const Xbyak::Opmask& _msk, const Xbyak::Reg64& _pos, const Xbyak::Address& _total,
                           const Xbyak::Reg64& _tmp, const Xbyak::Reg64& _tmp1, int N) {
     inLocalLabel();
@@ -86,6 +101,8 @@ class JitBase : protected Xbyak::CodeGenerator {
 class JitAvx : protected JitBase {
  protected:
   static int constexpr VBits = 256;
+  static int constexpr VecBytes = VBits / 8;
+  static int constexpr RegCount = 16;
   typedef Xbyak::Ymm vreg_t;
 };
 
@@ -93,6 +110,7 @@ class JitAvx2 : protected JitAvx {
  protected:
   static int constexpr VBits = 256;
   typedef Xbyak::Ymm vreg_t;
+  void vxor(const vreg_t& x1, const vreg_t& x2, const Xbyak::Operand& op) { vpxor(x1, x2, op); }
 
   void loadbf16_f32(const Xbyak::Ymm& dst, const Xbyak::Address& addr) {
     vpmovzxwd(dst, addr);
@@ -103,8 +121,12 @@ class JitAvx2 : protected JitAvx {
 class JitAvx512f : protected JitAvx2 {
  protected:
   static int constexpr VBits = 512;
+  static int constexpr VecBytes = VBits / 8;
+  static int constexpr RegCount = 32;
   typedef Xbyak::Zmm vreg_t;
 
+  void vxor(const vreg_t& x1, const vreg_t& x2, const Xbyak::Operand& op) { vpxorq(x1, x2, op); }
+
   void interleave_2rows_4regs(Xbyak::Zmm* src_2regs, Xbyak::Zmm* tmp_2reg) {
     vpunpcklwd(tmp_2reg[0], src_2regs[0], src_2regs[1]);
     vpunpckhwd(tmp_2reg[1], src_2regs[0], src_2regs[1]);
@@ -191,18 +213,20 @@ class JitAvx512f : protected JitAvx2 {
   }
 };
 
+class JitAvx512_bf16 : protected JitAvx512f {};
+
 class JitAvx512_fp16 : protected JitAvx512f {};
 
 class JitAvx512vnni : protected JitAvx512f {
  protected:
-  void vpdpbusds_evex(const Xbyak::Xmm& x1, const Xbyak::Xmm& x2, const Xbyak::Operand& op) {
+  void vpdpbusds_(const Xbyak::Xmm& x1, const Xbyak::Xmm& x2, const Xbyak::Operand& op) {
     vpdpbusds(x1, x2, op, Xbyak::EvexEncoding);
   }
 };
 
 class JitAvxvnni : protected JitAvx2 {
  protected:
-  void vpdpbusds_vex(const Xbyak::Xmm& x1, const Xbyak::Xmm& x2, const Xbyak::Operand& op) {
+  void vpdpbusds_(const Xbyak::Xmm& x1, const Xbyak::Xmm& x2, const Xbyak::Operand& op) {
     vpdpbusds(x1, x2, op, Xbyak::VexEncoding);
   }
 };
@@ -215,6 +239,15 @@ class JitAmxtile : protected JitAvx512f {
     uint16_t colb[16];
     uint8_t rows[16];
   };
+  static int constexpr TileCount = 8;
+
+  typedef long long (*configure_t)(void*);
+
+  static void generate_config(Xbyak::CodeGenerator* g) {
+    Xbyak::util::StackFrame st(g, 1, 0, 0);
+    auto& parambase = st.p[0];
+    g->ldtilecfg(g->ptr[parambase]);
+  }
 
   static void configure_tiles(tileconfig_t& tc, int TILE_M, int TILE_N, int TILE_K, int elesize, int ANum, int BNum,
                               int CNum) {
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas.h
index badb571b3926b..2ecb6630d3da0 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas.h
@@ -30,6 +30,7 @@ enum JBLAS_ISA : uint32_t {
   JblasAMX_BF16,
   JblasAMX_INT8,
   JblasAVX512_FP16,
+  JblasAVX512_BF16,
 };
 enum class JBLAS_DTYPE : uint32_t {
   EleBitsMask = 0xff,
@@ -80,57 +81,6 @@ enum JBLAS_ELTWISEOP {
   LINEAR,
 };
 
-enum class JBLAS_GEMM_CORE : uint32_t {
-  // INT32=LSB|**8bits:NTile**||**8bits:PackRow**||**8bits:CompType**||**8bits:Reserve**|
-  Undef = 0,
-  NTILE_MASK = 0xff,
-  NTILE_SHIFT = 0,
-  NTILE_24 = 24,
-  NTILE_48 = 48,
-  NTILE_64 = 64,
-  NTILE_96 = 96,
-  PACKROW_MASK = 0xff00,
-  PACKROW_SHIFT = 8,
-  PACKROW_1 = 1 << PACKROW_SHIFT,
-  PACKROW_2 = 2 << PACKROW_SHIFT,
-  PACKROW_4 = 4 << PACKROW_SHIFT,
-  COMP_MASK = 0xff0000,
-  COMP_SHIFT = 16,
-  COMP_FP32 = 0 << COMP_SHIFT,
-  COMP_BF16 = 1 << COMP_SHIFT,
-  COMP_FP16 = 2 << COMP_SHIFT,
-  COMP_INT_START = 3 << COMP_SHIFT,
-  COMP_INT8_US = COMP_INT_START,
-  COMP_INT8_SS = 4 << COMP_SHIFT,
-  COMP_INT8_SU = 5 << COMP_SHIFT,
-  COMP_INT16_SS = 6 << COMP_SHIFT,
-  COMP_INT8_US_FP32 = 7 << COMP_SHIFT,
-  COMP_INT8_SS_FP32 = 8 << COMP_SHIFT,
-  COMP_INT8_SU_FP32 = 9 << COMP_SHIFT,
-  ISA_MASK = 0xff000000,
-  ISA_SHIFT = 24,
-  ISA_AVX2 = (uint32_t)JBLAS_ISA::JblasAVX2 << ISA_SHIFT,
-  ISA_AVX512F = (uint32_t)JBLAS_ISA::JblasAVX512F << ISA_SHIFT,
-  ISA_AVX_VNNI = (uint32_t)JBLAS_ISA::JblasAVX_VNNI << ISA_SHIFT,
-  ISA_AVX512_VNNI = (uint32_t)JBLAS_ISA::JblasAVX512_VNNI << ISA_SHIFT,
-  ISA_AMX_INT8 = (uint32_t)JBLAS_ISA::JblasAMX_INT8 << ISA_SHIFT,
-  ISA_AMX_BF16 = (uint32_t)JBLAS_ISA::JblasAMX_BF16 << ISA_SHIFT,
-  ISA_AVX512_FP16 = (uint32_t)JBLAS_ISA::JblasAVX512_FP16 << ISA_SHIFT,
-  AVX2_4X24 = NTILE_24 | PACKROW_1 | COMP_FP32 | ISA_AVX2,
-  AVX2_2X48 = NTILE_48 | PACKROW_1 | COMP_FP32 | ISA_AVX2,
-  AVX512F_8x48 = NTILE_48 | PACKROW_1 | COMP_FP32 | ISA_AVX512F,
-  AMX_BF16_16x64 = NTILE_64 | PACKROW_2 | COMP_BF16 | ISA_AMX_BF16,
-  AMX_BF16_16x48 = NTILE_48 | PACKROW_2 | COMP_BF16 | ISA_AMX_BF16,
-  AVX512_FP16_8x64 = NTILE_64 | PACKROW_2 | COMP_FP16 | ISA_AVX512_FP16,
-  AVX512_FP16_8x96 = NTILE_96 | PACKROW_2 | COMP_FP16 | ISA_AVX512_FP16,
-  AVX_VNNI_2x48 = NTILE_48 | PACKROW_4 | COMP_INT8_US | ISA_AVX_VNNI,
-  AVX_VNNI_4x24 = NTILE_24 | PACKROW_4 | COMP_INT8_US | ISA_AVX_VNNI,
-  AVX512_VNNI_8x48 = NTILE_48 | PACKROW_4 | COMP_INT8_US | ISA_AVX512_VNNI,
-  AMX_INT8_16x64_US = NTILE_64 | PACKROW_4 | COMP_INT8_US | ISA_AMX_INT8,
-  AMX_INT8_16x64_SS = NTILE_64 | PACKROW_4 | COMP_INT8_SS | ISA_AMX_INT8,
-  AMX_INT8_16x48_US = NTILE_48 | PACKROW_4 | COMP_INT8_US | ISA_AMX_INT8,
-  AMX_INT8_16x48_SS = NTILE_48 | PACKROW_4 | COMP_INT8_SS | ISA_AMX_INT8,
-};
 enum class JBLAS_PROLOGUEB_IDS : uint32_t {
   Undef = (uint32_t)-1,
   Begin = 0,
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_device.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_device.h
index 0b4a21d7d9e53..bf3f45db18f62 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_device.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_device.h
@@ -209,6 +209,7 @@ class CpuDevice {
     }
   }
   inline int getThreads() { return numthreads; }
+  inline int getCores() { return numcores; }
   inline uint32_t getL2CacheSize() { return L2Cache; }
   inline uint32_t getL1CacheSize() { return L1Cache; }
   inline bool AVX() { return mHasAVX; }
@@ -261,15 +262,15 @@ class CpuDevice {
 
 #define GetCPUDevice() auto _cd = jblas::device::CpuDevice::getInstance();
 
-
 class CpuBase {
  public:
   CpuBase() {
     GetCPUDevice();
     mL2Cache = _cd->getL2CacheSize();
+    mL1Cache = _cd->getL1CacheSize();
     mNumThreads = _cd->getThreads();
   }
-  size_t mL2Cache;
+  size_t mL2Cache, mL1Cache;
   int mNumThreads;
 };
 }  // namespace device
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_epilogue.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_epilogue.h
index d12dbac990578..1cab348459a0c 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_epilogue.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_epilogue.h
@@ -79,6 +79,8 @@ class CustomAccumulatorWriteBackWithEltop {
 template <JBLAS_ISA ISA_T>
 using AccumulatorWriteBackFp32 = AccumulatorWriteBack<ISA_T, float, float>;
 template <JBLAS_ISA ISA_T>
+using AccumulatorWriteBackInt32 = AccumulatorWriteBack<ISA_T, int, int>;
+template <JBLAS_ISA ISA_T>
 using AccumulatorWriteBackBf16 = AccumulatorWriteBack<ISA_T, utils::bf16, utils::bf16>;
 template <JBLAS_ISA ISA_T>
 using AccumulatorWriteBackFp16 = AccumulatorWriteBack<ISA_T, utils::fp16, utils::fp16>;
@@ -133,7 +135,7 @@ class CompFp32BlockEpilogue {
           (float*)_param.scales + K_offset * _param.ldsb + N_offset, srcptr, cachestep, dstptr, cachestep, M, N);
       assert(ret == JblasSuccess);
       if (_param.zps != nullptr) {
-        ret = kernel::wrapper::RemoveZeroPointBias::forward<ISA_T>(
+        ret = kernel::wrapper::RemoveZeroPointBias::forward_wei<ISA_T>(
             dstptr, cachestep, M, N, _param.zps + K_offset * _param.ldsb + N_offset,
             (float*)_param.scales + K_offset * _param.ldsb + N_offset, _param.ldra,
             _param.reduce + M_offset * _param.ldra + K_offset);
@@ -184,7 +186,8 @@ class CompInt8BlockEpilogue {
     int ldsa;
     // optional if A asym
     uint8_t* zpA = nullptr;
-    float* reduceB = nullptr;
+    void* reduceB = nullptr;
+    JBLAS_DTYPE reduceBdtype = JBLAS_DTYPE::F32;
     // optional if B asym
     int8_t* zpB = nullptr;
     float* reduceA = nullptr;
@@ -194,46 +197,56 @@ class CompInt8BlockEpilogue {
                      const int K_offset, const int M, const int N, const Param& _param, void* tmpcache,
                      size_t cachesize) {
     JBLAS_CODE ret = JblasNotSupport;
-    if (_param.scaleBdtype == JBLAS_DTYPE::F32) {
-      ret = kernel::wrapper::DequanS32Fp32::template forward<ISA_T>(
-          srcptr, cachestep, (float*)srcptr, cachestep, M, N, _param.scalesA + M_offset * _param.ldsa + K_offset,
-          _param.ldsa, (float*)_param.scalesB + N_offset + K_offset * _param.ldsb);
-      assert(ret == JblasSuccess);
-    } else if (_param.scaleBdtype == JBLAS_DTYPE::BF16) {
-      ret = kernel::wrapper::DequanS32Fp32::template forward<ISA_T>(
-          srcptr, cachestep, (float*)srcptr, cachestep, M, N, _param.scalesA + M_offset * _param.ldsa + K_offset,
-          _param.ldsa, (utils::bf16*)_param.scalesB + N_offset + K_offset * _param.ldsb);
+    float* scab = nullptr;
+    size_t ScaleBTmpSize = N * sizeof(float);
+    size_t ReduceBTmpSize = N * sizeof(float);
+    assert(cachesize >= (ScaleBTmpSize + ReduceBTmpSize));
+    if (_param.scaleBdtype == JBLAS_DTYPE::BF16) {
+      auto scache = (float*)tmpcache;
+      ret = kernel::wrapper::Memcpy2DBf16CvtFp32::template forward<ISA_T>(
+          (utils::bf16*)_param.scalesB + N_offset + K_offset * _param.ldsb, scache, 1, N, N, N, false);
       assert(ret == JblasSuccess);
+      scab = scache;
+    } else if (_param.scaleBdtype == JBLAS_DTYPE::F32) {
+      scab = (float*)_param.scalesB + N_offset + K_offset * _param.ldsb;
     }
-    ret = kernel::wrapper::AccumulateFp32::template forward<ISA_T>((float*)srcptr, cachestep, dstptr, cachestep, M, N);
-    if (ret != JblasSuccess) {
-      assert(0);
-      return ret;
+    float* redb = nullptr;
+    if (_param.reduceB) {
+      if (_param.reduceBdtype == JBLAS_DTYPE::BF16) {
+        auto rcache = (float*)((char*)tmpcache + ScaleBTmpSize);
+        ret = kernel::wrapper::Memcpy2DBf16CvtFp32::template forward<ISA_T>(
+            (utils::bf16*)_param.reduceB + N_offset + K_offset * _param.ldsb, rcache, 1, N, N, N, false);
+        assert(ret == JblasSuccess);
+        redb = rcache;
+      } else if (_param.reduceBdtype == JBLAS_DTYPE::F32) {
+        redb = (float*)_param.reduceB + N_offset + K_offset * _param.ldsb;
+      }
     }
+    ret = kernel::wrapper::DequanS32Fp32::template forward<ISA_T>(srcptr, cachestep, (float*)srcptr, cachestep, M, N,
+                                                                  _param.scalesA + M_offset * _param.ldsa + K_offset,
+                                                                  _param.ldsa, scab);
+    assert(ret == JblasSuccess);
+    ret = kernel::wrapper::AccumulateFp32::template forward<ISA_T>((float*)srcptr, cachestep, dstptr, cachestep, M, N);
+    assert(ret == JblasSuccess);
 
-    if (_param.zpA == nullptr && _param.zpB == nullptr) {
-      return ret;
-    } else if (_param.zpA != nullptr && _param.zpB == nullptr) {
-      ret = kernel::wrapper::RemoveZeroPointBias::template forward<ISA_T>(
-          dstptr, cachestep, M, N, _param.zpA + M_offset * _param.ldsa + K_offset,
-          _param.scalesA + M_offset * _param.ldsa + K_offset, _param.ldsa,
-          _param.reduceB + N_offset + K_offset * _param.ldsb);
-
-    } else if (_param.zpA == nullptr && _param.zpB != nullptr) {
-      if (_param.scaleBdtype == JBLAS_DTYPE::F32) {
-        ret = kernel::wrapper::RemoveZeroPointBias::template forward<ISA_T>(
-            dstptr, cachestep, M, N, _param.zpB + N_offset + K_offset * _param.ldsb,
-            (float*)_param.scalesB + N_offset + K_offset * _param.ldsb, _param.ldsa,
+    if (_param.zpA == nullptr) {
+      if (_param.zpB == nullptr) {
+        return ret;
+      } else {
+        ret = kernel::wrapper::RemoveZeroPointBias::template forward_wei<ISA_T>(
+            dstptr, cachestep, M, N, _param.zpB + N_offset + K_offset * _param.ldsb, scab, _param.ldsa,
             _param.reduceA + M_offset * _param.ldsa + K_offset);
       }
-
     } else {
-      if (_param.scaleBdtype == JBLAS_DTYPE::F32) {
-        ret = kernel::wrapper::RemoveZeroPointBias::template forward<ISA_T>(
+      if (_param.zpB == nullptr) {
+        ret = kernel::wrapper::RemoveZeroPointBias::template forward_act<ISA_T>(
+            dstptr, cachestep, M, N, _param.zpA + M_offset * _param.ldsa + K_offset,
+            _param.scalesA + M_offset * _param.ldsa + K_offset, _param.ldsa, redb);
+      } else {
+        ret = kernel::wrapper::RemoveZeroPointBias::template forward_both<ISA_T>(
             dstptr, cachestep, M, N, _param.zpA + M_offset * _param.ldsa + K_offset,
-            _param.zpB + N_offset + K_offset * _param.ldsb, _param.scalesA + M_offset * _param.ldsa + K_offset,
-            (float*)_param.scalesB + N_offset + K_offset * _param.ldsb, _param.ldsa, _param.K,
-            _param.reduceA + M_offset * _param.ldsa + K_offset, _param.reduceB + N_offset + K_offset * _param.ldsb);
+            _param.zpB + N_offset + K_offset * _param.ldsb, _param.scalesA + M_offset * _param.ldsa + K_offset, scab,
+            _param.ldsa, _param.K, _param.reduceA + M_offset * _param.ldsa + K_offset, redb);
       }
     }
     return ret;
@@ -271,15 +284,15 @@ class ZpDequantInt32ToFp32 {
     if (_param.zpA == nullptr && _param.zpB == nullptr) {
       return ret;
     } else if (_param.zpA != nullptr && _param.zpB == nullptr) {
-      ret = kernel::wrapper::RemoveZeroPointBias::template forward<ISA_T>(
+      ret = kernel::wrapper::RemoveZeroPointBias::template forward_act<ISA_T>(
           cptr, _param.ldc, M, N, _param.zpA + M_offset * _param.ldsa, _param.scalesA + M_offset * _param.ldsa,
           _param.ldsa, _param.reduceB + N_offset);
     } else if (_param.zpA == nullptr && _param.zpB != nullptr) {
-      ret = kernel::wrapper::RemoveZeroPointBias::template forward<ISA_T>(cptr, _param.ldc, M, N, _param.zpB + N_offset,
-                                                                          _param.scalesB + N_offset, _param.ldsa,
-                                                                          _param.reduceA + M_offset * _param.ldsa);
+      ret = kernel::wrapper::RemoveZeroPointBias::template forward_wei<ISA_T>(
+          cptr, _param.ldc, M, N, _param.zpB + N_offset, _param.scalesB + N_offset, _param.ldsa,
+          _param.reduceA + M_offset * _param.ldsa);
     } else {
-      ret = kernel::wrapper::RemoveZeroPointBias::template forward<ISA_T>(
+      ret = kernel::wrapper::RemoveZeroPointBias::template forward_both<ISA_T>(
           cptr, _param.ldc, M, N, _param.zpA + M_offset * _param.ldsa, _param.zpB + N_offset,
           _param.scalesA + M_offset * _param.ldsa, _param.scalesB + N_offset, _param.ldsa, _param.K,
           _param.reduceA + M_offset * _param.ldsa, _param.reduceB + N_offset);
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_gemm.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_gemm.h
index 5fefbca23b477..e024dcd45643e 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_gemm.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_gemm.h
@@ -19,2852 +19,2152 @@
 
 namespace jblas {
 namespace gemm {
+enum class CompType : uint32_t {
+  COMP_FP32 = 0,
+  COMP_BF16_FP32 = 1,
+  COMP_FP16_FP16 = 2,
+  COMP_INT_START = 3,
+  COMP_INT8_US_INT32 = COMP_INT_START,
+  COMP_INT8_UU_INT32 = 4,
+  COMP_INT8_SS_INT32 = 5,
+  COMP_INT8_SU_INT32 = 6,
+  COMP_INT16_SS_INT32 = 7,
+  COMP_INT8_US_FP32 = 8,
+  COMP_INT8_UU_FP32 = 9,
+  COMP_INT8_SS_FP32 = 10,
+  COMP_INT8_SU_FP32 = 11,
+};
 
-class GemmCore_Row_NN_4x24_AVX2 {
+class CoreAttr {
  public:
-  struct params {
-    float *matA, *matB, *matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-  };
-  typedef long long (*func_t)(params*);
-  typedef float AType;
-  typedef float BType;
-  typedef float CType;
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX2_4X24;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 24, MTILE = 4, KTILE = 4 / sizeof(BType);
-  static int constexpr KUNROLL = 2;
-  static int constexpr PACK_ROW = 1;
-  static int constexpr PREFERED_N = 144;
-  class MicroKernel : protected jblas::xbyak::JitAvx2 {
-   public:
-    MicroKernel() {}
-    static int constexpr VecBytes = 32;
-    static int constexpr VecElements = VecBytes / sizeof(CType);
-    int CRegCount = 12, BRegCount = 3, ARegCount = 1;
-    int CReg = 0, BReg = 12, AReg = 15, TmpReg = BReg;
-    int const NRegs = NTILE / VecElements;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
+  // INT32=LSB|**8bits:NTile**||**8bits:PackRow**||**8bits:CompType**||**8bits:Reserve**|
+  static uint32_t constexpr NTILE_MASK = 0xff, NTILE_SHIFT = 0, PACKROW_MASK = 0xff00, PACKROW_SHIFT = 8,
+                            COMP_MASK = 0xff0000, COMP_SHIFT = 16, ISA_MASK = 0xff000000, ISA_SHIFT = 24;
 
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
-      BRegCount = NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vpxor(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n16", T_NEAR);
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      cmp(reg_tmp, 16);
-      jl(".n8", T_NEAR);
-      generate_kloop(_mtile, 2);
-      write_back(_mtile, 2);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n8");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(_mtile, 1);
-      write_back(_mtile, 1);
-      add(reg_itern, 8);
-      add(reg_matBptr, 8 * sizeof(BType));
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
-    }
-
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
-
-    void generate_fma(int _mtile, int _NRegs, int _ktile, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _ktile; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int i = 0; i < _NRegs; i++) {
-          vmovups(Xbyak::Ymm(BReg + i), ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
-        }
-        for (int mm = 0; mm < _mtile; mm++) {
-          vbroadcastss(Xbyak::Ymm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vfmadd231ps(Xbyak::Ymm(CReg + mm * NRegs + i), Xbyak::Ymm(BReg + i), Xbyak::Ymm(AReg));
-          }
-        }
-      }
-    }
-
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vaddps(Xbyak::Ymm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      L(".LEND");
-      nop();
-      outLocalLabel();
-    }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
+  static inline uint32_t get_mask_val(uint32_t raw, uint32_t mask, uint32_t shift) { return (raw & mask) >> shift; }
+  static constexpr uint32_t make_core_id(uint32_t NTile, uint32_t PackRow, uint32_t CompType, uint32_t ISA) {
+    return (NTile << NTILE_SHIFT) | (PackRow << PACKROW_SHIFT) | (CompType << COMP_SHIFT) | (ISA << ISA_SHIFT);
+  }
 
- public:
-  GemmCore_Row_NN_4x24_AVX2() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
-    }
+  static void parse_id(uint32_t id, uint32_t* vals) {
+    vals[0] = get_mask_val(id, NTILE_MASK, NTILE_SHIFT);
+    vals[1] = get_mask_val(id, PACKROW_MASK, PACKROW_SHIFT);
+    vals[2] = get_mask_val(id, COMP_MASK, COMP_SHIFT);
+    vals[3] = get_mask_val(id, ISA_MASK, ISA_SHIFT);
   }
 
-  void forward(float* matA, float* matB, float* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
-    }
+  static const char* to_str(uint32_t id) {
+    static char tmp[128];
+    uint32_t vals[4];
+    parse_id(id, vals);
+    sprintf(tmp, "N%d_PACK%d_COMP%d_ISA%d", vals[0], vals[1], vals[2], vals[3]);
+    return tmp;
   }
 
- private:
-  std::array<MicroKernel, MTILE> mCodes;
+  static inline size_t get_bsize(uint32_t id) {
+    auto packrow = get_mask_val(id, PACKROW_MASK, PACKROW_SHIFT);
+    return size_t(4 / packrow);
+  }
 };
 
-class GemmCore_Row_NN_2x48_AVX2 {
+namespace code {
+
+template <int _NTILE, int _MTILE = 0>
+class Avx2N8P1 : protected jblas::xbyak::JitAvx2 {
  public:
-  struct params {
-    float *matA, *matB, *matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-  };
-  typedef long long (*func_t)(params*);
+  static int constexpr RegLen = 8, PackRow = 1;
+  static_assert(_NTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? (RegCount - 1) / NRegs : _MTILE;
+  static_assert(NRegs * MRegs <= RegCount - 1);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs, KTILE = 1;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAVX2;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_FP32;
   typedef float AType;
   typedef float BType;
   typedef float CType;
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX2_2X48;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 48, MTILE = 2, KTILE = 4 / sizeof(BType);
-  static int constexpr KUNROLL = 2;
-  static int constexpr PACK_ROW = 1;
-  static int constexpr PREFERED_N = 144;
-  class MicroKernel : protected jblas::xbyak::JitAvx2 {
-   public:
-    MicroKernel() {}
-    static int constexpr VecBytes = 32;
-    static int constexpr VecElements = VecBytes / sizeof(CType);
-    int CRegCount = 12, BRegCount = 1, ARegCount = 2;
-    int CReg = 0, BReg = 12, AReg = 13, TmpReg = BReg;
-    int const NRegs = NTILE / VecElements;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
-      BRegCount = 1;
-      ARegCount = _mtile;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vpxor(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
 
-      outLocalLabel();  // end of local label
-    }
-
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
-
-    void generate_fma(int _mtile, int _NRegs, int _ktile, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _ktile; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int i = 0; i < _mtile; ++i)
-          vbroadcastss(Xbyak::Ymm(AReg + i), ptr[reg_matAptr + kk * AKStepSize + reg_astep * i]);
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(Xbyak::Ymm(BReg), ptr[_reg_matBptr + kk * BKStepSize + j * VecBytes]);
-          for (int i = 0; i < _mtile; ++i)
-            vfmadd231ps(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(BReg), Xbyak::Ymm(AReg + i));
-        }
-      }
-    }
-
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vaddps(Xbyak::Ymm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      L(".LEND");
-      nop();
-      outLocalLabel();
-    }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
+  struct params {
+    AType* matA;
+    int astride;
+    BType* matB;
+    int bstride;
+    CType* matC;
+    int cstride;
+    int k;
+    int n;
+    int init;
   };
+  typedef long long (*func_t)(params*);
 
- public:
-  GemmCore_Row_NN_2x48_AVX2() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
+  int CRegCount = 0, BRegCount = 0, ARegCount = 0, TmpRegCount = 0;
+  int CReg = 0, BReg = 0, AReg = 0, TmpReg = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
+
+ protected:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_ret = rax;
+  Xbyak::Opmask msk_wr = k1;
+
+  void assign_regs() {
+    CRegCount = MRegs * NRegs;
+    ARegCount = 1;
+    BRegCount = RegCount - ARegCount - CRegCount;
+    if (BRegCount < NRegs) {
+      BRegCount = 0;
+      ARegCount = BRegCount + 1;
+    }
+    if (BRegCount > NRegs) {
+      BRegCount = NRegs;
     }
+    CReg = 0;
+    BReg = CReg + CRegCount;
+    AReg = BReg + BRegCount;
+    TmpReg = AReg + ARegCount;
+    assert(TmpReg <= RegCount);
+    TmpRegCount = RegCount - TmpReg;
   }
 
-  void forward(float* matA, float* matB, float* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();  // use local label for multiple instance
+    Xbyak::util::StackFrame st(this, 1, 10, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
   }
 
- private:
-  std::array<MicroKernel, MTILE> mCodes;
-};
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
 
-class GemmCore_Row_NN_4x24_AVX_VNNI {
- public:
-  struct params {
-    uint8_t* matA;
-    int8_t* matB;
-    int32_t* matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-  };
-  typedef long long (*func_t)(params*);
-  typedef uint8_t AType;
-  typedef int8_t BType;
-  typedef int32_t CType;
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX_VNNI_4x24;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 24, MTILE = 4, KTILE = 4 / sizeof(BType);
-  static int constexpr PACK_ROW = KTILE;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 192;
-
-  class MicroKernel : protected jblas::xbyak::JitAvxvnni {
-   public:
-    MicroKernel() {}
-    int CRegCount = 12, BRegCount = 1, ARegCount = 1;
-    int CReg = 0, BReg = 12, AReg = 13, TmpReg = 14;
-    int const NRegs = 3;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 32;
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vxorps(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j));
+  void generate_fma(int _mtile, int _ktile) {
+    for (int kk = 0; kk < _ktile; kk++) {
+      lea(reg_tmp1, ptr[reg_matAptr + kk * AKStepSize]);
+      if (BRegCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          vmovups(vreg_t(BReg + i), ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
         }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
-    }
-
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KTILE * KUNROLL);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, AKStepSize * KUNROLL);
-      add(reg_tmp1, BKStepSize * KUNROLL);
-      add(reg_iterk, KTILE * KUNROLL);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
-
-    void generate_fma(int _mtile, int _NRegs, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _kunroll; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
         for (int mm = 0; mm < _mtile; mm++) {
-          vpbroadcastd(Xbyak::Ymm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vpdpbusds_vex(Xbyak::Ymm(CReg + mm * NRegs + i), Xbyak::Ymm(AReg),
-                          ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
+          vbroadcastss(vreg_t(AReg), ptr[reg_tmp1]);
+          add(reg_tmp1, reg_astride);
+          for (int i = 0; i < NRegs; i++) {
+            vfmadd231ps(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg), vreg_t(BReg + i));
           }
         }
-      }
-    }
-
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vpaddd(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
+      } else if (BRegCount == 0) {
+        for (int mm = 0; mm < _mtile; mm += ARegCount) {
+          int mm_re = utils::remainsize(mm, _mtile, ARegCount);
+          for (int imm = 0; imm < mm_re; imm++) {
+            vbroadcastss(vreg_t(AReg + imm), ptr[reg_tmp1]);
+            add(reg_tmp1, reg_astride);
+            for (int i = 0; i < NRegs; i++) {
+              vfmadd231ps(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg + imm),
+                          ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
+            }
+          }
         }
-        add(reg_matCptr, reg_cstep);
+      } else {
+        assert(0);
       }
-      L(".LEND");
-      nop();
-      outLocalLabel();
-    }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-  };
-
- public:
-  GemmCore_Row_NN_4x24_AVX_VNNI() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
     }
   }
 
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
-    }
-  }
-
- private:
-  std::array<MicroKernel, MTILE> mCodes;
-};
-
-class GemmCore_Row_NN_2x48_AVX_VNNI {
- public:
-  struct params {
-    uint8_t* matA;
-    int8_t* matB;
-    int32_t* matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-  };
-  typedef long long (*func_t)(params*);
-  typedef uint8_t AType;
-  typedef int8_t BType;
-  typedef int32_t CType;
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX_VNNI_2x48;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 48, MTILE = 2, KTILE = 4 / sizeof(BType);
-  static int constexpr PACK_ROW = KTILE;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 192;
-
-  class MicroKernel : protected jblas::xbyak::JitAvxvnni {
-   public:
-    MicroKernel() {}
-    int CRegCount = 12, BRegCount = 1, ARegCount = 1;
-    int CReg = 0, BReg = 12, AReg = 13, TmpReg = 14;
-    int const NRegs = 6;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 32;
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vxorps(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
-    }
-
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KTILE * KUNROLL);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, AKStepSize * KUNROLL);
-      add(reg_tmp1, BKStepSize * KUNROLL);
-      add(reg_iterk, KTILE * KUNROLL);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
-
-    void generate_fma(int _mtile, int _NRegs, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _kunroll; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int mm = 0; mm < _mtile; mm++) {
-          vpbroadcastd(Xbyak::Ymm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vpdpbusds_vex(Xbyak::Ymm(CReg + mm * NRegs + i), Xbyak::Ymm(AReg),
-                          ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
-          }
-        }
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vxor(vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j));
       }
     }
-
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vpaddd(Xbyak::Ymm(CReg + i * NRegs + j), Xbyak::Ymm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Ymm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(vreg_t(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
       }
-      L(".LEND");
-      nop();
-      outLocalLabel();
-    }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-  };
-
- public:
-  GemmCore_Row_NN_2x48_AVX_VNNI() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
+      add(reg_matCptr, reg_cstride);
     }
+    L(".end");
+    outLocalLabel();
   }
 
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(CReg + i * NRegs + j));
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    outLocalLabel();
   }
-
- private:
-  std::array<MicroKernel, MTILE> mCodes;
 };
 
-class GemmCore_Row_NN_8x48_AVX512F {
+template <int _NTILE, int _MTILE = 0>
+class Avx512fN16P1 : protected jblas::xbyak::JitAvx512f {
  public:
-  struct params {
-    float *matA, *matB, *matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-  };
-  typedef long long (*func_t)(params*);
+  static int constexpr RegLen = 16, PackRow = 1;
+  static_assert(_NTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? (RegCount - 1) / NRegs : _MTILE;
+  static_assert(NRegs * MRegs <= RegCount - 1);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs, KTILE = 1;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAVX512F;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_FP32;
   typedef float AType;
   typedef float BType;
   typedef float CType;
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX512F_8x48;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 48, MTILE = 8, KTILE = 4 / sizeof(BType);
-  static int constexpr KUNROLL = 2;
-  static int constexpr PACK_ROW = 1;
-  static int constexpr PREFERED_N = 144;
-  class MicroKernel : protected jblas::xbyak::JitAvx512f {
-   public:
-    MicroKernel() {}
-    int CRegCount = 24, BRegCount = 6, ARegCount = 1;
-    int CReg = 0, BReg = 24, AReg = 27, TmpReg = 28;
-    int const NRegs = 3;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
 
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
+  struct params {
+    AType* matA;
+    int astride;
+    BType* matB;
+    int bstride;
+    CType* matC;
+    int cstride;
+    int k;
+    int n;
+    int init;
+  };
+  typedef long long (*func_t)(params*);
+
+  int CRegCount = 0, BRegCount = 0, ARegCount = 0, TmpRegCount = 0;
+  int CReg = 0, BReg = 0, AReg = 0, TmpReg = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
+
+ protected:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_ret = rax;
+  Xbyak::Opmask msk_wr = k1;
+
+  void assign_regs() {
+    CRegCount = MRegs * NRegs;
+    ARegCount = 1;
+    BRegCount = RegCount - ARegCount - CRegCount;
+    if (BRegCount < NRegs) {
+      BRegCount = 0;
+      ARegCount = BRegCount + 1;
+    }
+    if (BRegCount > NRegs) {
       BRegCount = NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vpxorq(Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n32", T_NEAR);
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      cmp(reg_tmp, 32);
-      jl(".n16", T_NEAR);
-      generate_kloop(_mtile, 2);
-      write_back(_mtile, 2);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(_mtile, 1);
-      write_back(_mtile, 1);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(BType));
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
     }
+    CReg = 0;
+    BReg = CReg + CRegCount;
+    AReg = BReg + BRegCount;
+    TmpReg = AReg + ARegCount;
+    assert(TmpReg <= RegCount);
+    TmpRegCount = RegCount - TmpReg;
+  }
 
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();  // use local label for multiple instance
+    Xbyak::util::StackFrame st(this, 1, 10, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
 
-    void generate_fma(int _mtile, int _NRegs, int _ktile, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _ktile; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int i = 0; i < _NRegs; i++) {
-          vmovups(Xbyak::Zmm(BReg + i), ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
+
+  void generate_fma(int _mtile, int _ktile) {
+    for (int kk = 0; kk < _ktile; kk++) {
+      lea(reg_tmp1, ptr[reg_matAptr + kk * AKStepSize]);
+      if (BRegCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          vmovups(vreg_t(BReg + i), ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
         }
         for (int mm = 0; mm < _mtile; mm++) {
-          vbroadcastss(Xbyak::Zmm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vfmadd231ps(Xbyak::Zmm(CReg + mm * NRegs + i), Xbyak::Zmm(BReg + i), Xbyak::Zmm(AReg));
+          vbroadcastss(vreg_t(AReg), ptr[reg_tmp1]);
+          add(reg_tmp1, reg_astride);
+          for (int i = 0; i < NRegs; i++) {
+            vfmadd231ps(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg), vreg_t(BReg + i));
+          }
+        }
+      } else if (BRegCount == 0) {
+        for (int mm = 0; mm < _mtile; mm += ARegCount) {
+          int mm_re = utils::remainsize(mm, _mtile, ARegCount);
+          for (int imm = 0; imm < mm_re; imm++) {
+            vbroadcastss(vreg_t(AReg + imm), ptr[reg_tmp1]);
+            add(reg_tmp1, reg_astride);
+            for (int i = 0; i < NRegs; i++) {
+              vfmadd231ps(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg + imm),
+                          ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
+            }
           }
         }
+      } else {
+        assert(0);
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vaddps(Xbyak::Zmm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vxor(vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j));
       }
-      L(".LEND");
-      nop();
-      outLocalLabel();
     }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
-
- public:
-  GemmCore_Row_NN_8x48_AVX512F() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(vreg_t(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    L(".end");
+    outLocalLabel();
   }
 
-  void forward(float* matA, float* matB, float* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(CReg + i * NRegs + j));
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    outLocalLabel();
   }
-
- private:
-  std::array<MicroKernel, MTILE> mCodes;
 };
 
-class GemmCore_Row_NN_8x64_AVX512_FP16 {
+template <int _NTILE, int _MTILE = 0>
+class Avx512fp16N32P1 : protected jblas::xbyak::JitAvx512_fp16 {
  public:
+  static int constexpr RegLen = 32, PackRow = 1;
+  static_assert(_NTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? (RegCount - 1) / NRegs : _MTILE;
+  static_assert(NRegs * MRegs <= RegCount - 1);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs, KTILE = 1;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAVX512_FP16;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_FP16_FP16;
   typedef utils::fp16 AType;
   typedef utils::fp16 BType;
   typedef utils::fp16 CType;
+
   struct params {
     AType* matA;
+    int astride;
     BType* matB;
+    int bstride;
     CType* matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
+    int cstride;
+    int k;
+    int n;
+    int init;
   };
   typedef long long (*func_t)(params*);
 
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX512_FP16_8x64;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 64, MTILE = 12, KTILE = 1;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PACK_ROW = 1;
-  static int constexpr PREFERED_N = 128;
-  class MicroKernel : protected jblas::xbyak::JitAvx512_fp16 {
-   public:
-    MicroKernel() {}
-    int CRegCount = 24, BRegCount = 2, ARegCount = 1;
-    int CReg = 0, BReg = 24, AReg = 26, TmpReg = 27;
-    int const NRegs = 2;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
+  int CRegCount = 0, BRegCount = 0, ARegCount = 0, TmpRegCount = 0;
+  int CReg = 0, BReg = 0, AReg = 0, TmpReg = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
+
+ protected:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_ret = rax;
+  Xbyak::Opmask msk_wr = k1;
+
+  void assign_regs() {
+    CRegCount = MRegs * NRegs;
+    ARegCount = 1;
+    BRegCount = RegCount - ARegCount - CRegCount;
+    if (BRegCount < NRegs) {
+      BRegCount = 0;
+      ARegCount = BRegCount + 1;
+    }
+    if (BRegCount > NRegs) {
       BRegCount = NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vpxorq(Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n32", T_NEAR);
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      generate_kloop(_mtile, 1);
-      write_back(_mtile, 1);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
     }
+    CReg = 0;
+    BReg = CReg + CRegCount;
+    AReg = BReg + BRegCount;
+    TmpReg = AReg + ARegCount;
+    assert(TmpReg <= RegCount);
+    TmpRegCount = RegCount - TmpReg;
+  }
 
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();  // use local label for multiple instance
+    Xbyak::util::StackFrame st(this, 1, 10, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
 
-    void generate_fma(int _mtile, int _NRegs, int _ktile, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _ktile; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int i = 0; i < _NRegs; i++) {
-          vmovups(Xbyak::Zmm(BReg + i), ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
+
+  void generate_fma(int _mtile, int _ktile) {
+    for (int kk = 0; kk < _ktile; kk++) {
+      lea(reg_tmp1, ptr[reg_matAptr + kk * AKStepSize]);
+      if (BRegCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          vmovups(vreg_t(BReg + i), ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
         }
         for (int mm = 0; mm < _mtile; mm++) {
-          vpbroadcastw(Xbyak::Zmm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vfmadd231ph(Xbyak::Zmm(CReg + mm * NRegs + i), Xbyak::Zmm(BReg + i), Xbyak::Zmm(AReg));
+          vpbroadcastw(vreg_t(AReg), ptr[reg_tmp1]);
+          add(reg_tmp1, reg_astride);
+          for (int i = 0; i < NRegs; i++) {
+            vfmadd231ph(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg), vreg_t(BReg + i));
+          }
+        }
+      } else if (BRegCount == 0) {
+        for (int mm = 0; mm < _mtile; mm += ARegCount) {
+          int mm_re = utils::remainsize(mm, _mtile, ARegCount);
+          for (int imm = 0; imm < mm_re; imm++) {
+            vpbroadcastw(vreg_t(AReg + imm), ptr[reg_tmp1]);
+            add(reg_tmp1, reg_astride);
+            for (int i = 0; i < NRegs; i++) {
+              vfmadd231ph(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg + imm),
+                          ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
+            }
           }
         }
+      } else {
+        assert(0);
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vxor(vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j));
       }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vaddph(Xbyak::Zmm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      L(".LEND");
-      nop();
-      outLocalLabel();
     }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
-
- public:
-  GemmCore_Row_NN_8x64_AVX512_FP16() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(vreg_t(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    L(".end");
+    outLocalLabel();
   }
 
-  void forward(utils::fp16* matA, utils::fp16* matB, utils::fp16* matC, int _m, int _n, int _k, int _astride,
-               int _bstride, int _cstride, int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(CReg + i * NRegs + j));
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    outLocalLabel();
   }
-
- private:
-  std::array<MicroKernel, MTILE> mCodes;
 };
 
-class GemmCore_Row_NN_8x96_AVX512_FP16 {
+template <int _NTILE, int _MTILE = 0>
+class Avx512bf16N16P2 : protected jblas::xbyak::JitAvx512_bf16 {
  public:
-  typedef utils::fp16 AType;
-  typedef utils::fp16 BType;
-  typedef utils::fp16 CType;
+  static int constexpr RegLen = 16, PackRow = 2;
+  static_assert(_NTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? (RegCount - 1) / NRegs : _MTILE;
+  static_assert(NRegs * MRegs <= RegCount - 1);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs, KTILE = 2;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAVX512_BF16;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_BF16_FP32;
+  typedef utils::bf16 AType;
+  typedef utils::bf16 BType;
+  typedef float CType;
+
   struct params {
     AType* matA;
+    int astride;
     BType* matB;
+    int bstride;
     CType* matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
+    int cstride;
+    int k;
+    int n;
+    int init;
   };
   typedef long long (*func_t)(params*);
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX512_FP16_8x96;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 96, MTILE = 8, KTILE = 1;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PACK_ROW = 1;
-  static int constexpr PREFERED_N = 192;
-  class MicroKernel : protected jblas::xbyak::JitAvx512_fp16 {
-   public:
-    MicroKernel() {}
-    int CRegCount = 24, BRegCount = 3, ARegCount = 1;
-    int CReg = 0, BReg = 24, AReg = 27, TmpReg = 28;
-    int const NRegs = 3;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
 
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
+  int CRegCount = 0, BRegCount = 0, ARegCount = 0, TmpRegCount = 0;
+  int CReg = 0, BReg = 0, AReg = 0, TmpReg = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
+
+ protected:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_ret = rax;
+  Xbyak::Opmask msk_wr = k1;
+
+  void assign_regs() {
+    CRegCount = MRegs * NRegs;
+    ARegCount = 1;
+    BRegCount = RegCount - ARegCount - CRegCount;
+    if (BRegCount < NRegs) {
+      BRegCount = 0;
+      ARegCount = BRegCount + 1;
+    }
+    if (BRegCount > NRegs) {
       BRegCount = NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vpxorq(Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n64", T_NEAR);
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n64");
-      cmp(reg_tmp, 64);
-      jl(".n32", T_NEAR);
-      generate_kloop(_mtile, 2);
-      write_back(_mtile, 2);
-      add(reg_itern, 64);
-      add(reg_matBptr, 64 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      generate_kloop(_mtile, 1);
-      write_back(_mtile, 1);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
     }
+    CReg = 0;
+    BReg = CReg + CRegCount;
+    AReg = BReg + BRegCount;
+    TmpReg = AReg + ARegCount;
+    assert(TmpReg <= RegCount);
+    TmpRegCount = RegCount - TmpReg;
+  }
 
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();  // use local label for multiple instance
+    Xbyak::util::StackFrame st(this, 1, 10, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
+
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
 
-    void generate_fma(int _mtile, int _NRegs, int _ktile, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _ktile; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int i = 0; i < _NRegs; i++) {
-          vmovups(Xbyak::Zmm(BReg + i), ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
+  void generate_fma(int _mtile, int _ktile) {
+    for (int kk = 0; kk < _ktile; kk++) {
+      lea(reg_tmp1, ptr[reg_matAptr + kk * AKStepSize]);
+      if (BRegCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          vmovups(vreg_t(BReg + i), ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
         }
         for (int mm = 0; mm < _mtile; mm++) {
-          vpbroadcastw(Xbyak::Zmm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vfmadd231ph(Xbyak::Zmm(CReg + mm * NRegs + i), Xbyak::Zmm(BReg + i), Xbyak::Zmm(AReg));
+          vbroadcastss(vreg_t(AReg), ptr[reg_tmp1]);
+          add(reg_tmp1, reg_astride);
+          for (int i = 0; i < NRegs; i++) {
+            vdpbf16ps(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg), vreg_t(BReg + i));
           }
         }
+      } else if (BRegCount == 0) {
+        for (int mm = 0; mm < _mtile; mm += ARegCount) {
+          int mm_re = utils::remainsize(mm, _mtile, ARegCount);
+          for (int imm = 0; imm < mm_re; imm++) {
+            vbroadcastss(vreg_t(AReg + imm), ptr[reg_tmp1]);
+            add(reg_tmp1, reg_astride);
+            for (int i = 0; i < NRegs; i++) {
+              vdpbf16ps(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg + imm),
+                        ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
+            }
+          }
+        }
+      } else {
+        assert(0);
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vxor(vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j));
       }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vaddph(Xbyak::Zmm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      L(".LEND");
-      nop();
-      outLocalLabel();
     }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
-
- public:
-  GemmCore_Row_NN_8x96_AVX512_FP16() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(vreg_t(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    L(".end");
+    outLocalLabel();
   }
 
-  void forward(utils::fp16* matA, utils::fp16* matB, utils::fp16* matC, int _m, int _n, int _k, int _astride,
-               int _bstride, int _cstride, int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(CReg + i * NRegs + j));
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    outLocalLabel();
   }
-
- private:
-  std::array<MicroKernel, MTILE> mCodes;
 };
 
-class GemmCore_Row_NN_8x48_AVX512_VNNI {
+template <int _NTILE, int _MTILE = 0>
+class Avx512vnniN16P4 : protected jblas::xbyak::JitAvx512vnni {
  public:
-  struct params {
-    uint8_t* matA;
-    int8_t* matB;
-    int32_t* matC;
-    int k, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-  };
-  typedef long long (*func_t)(params*);
+  static int constexpr RegLen = 16, PackRow = 4;
+  static_assert(_NTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? (RegCount - 1) / NRegs : _MTILE;
+  static_assert(NRegs * MRegs <= RegCount - 1);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs, KTILE = 4;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAVX512_VNNI;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_INT8_US_INT32;
   typedef uint8_t AType;
   typedef int8_t BType;
   typedef int32_t CType;
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AVX512_VNNI_8x48;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 48, MTILE = 8, KTILE = 4 / sizeof(BType);
-  static int constexpr PACK_ROW = KTILE;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 192;
-
-  class MicroKernel : protected jblas::xbyak::JitAvx512vnni {
-   public:
-    MicroKernel() {}
-    int CRegCount = 24, BRegCount = 6, ARegCount = 1;
-    int CReg = 0, BReg = 24, AReg = 27, TmpReg = 28;
-    int const NRegs = 3;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code(int _mtile) {
-      reset();
-      generate_mtile(_mtile);
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
+  struct params {
+    AType* matA;
+    int astride;
+    BType* matB;
+    int bstride;
+    CType* matC;
+    int cstride;
+    int k;
+    int n;
+    int init;
+  };
+  typedef long long (*func_t)(params*);
+
+  int CRegCount = 0, BRegCount = 0, ARegCount = 0, TmpRegCount = 0;
+  int CReg = 0, BReg = 0, AReg = 0, TmpReg = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
 
-   protected:
-    void generate_mtile(int _mtile) {
-      CRegCount = _mtile * NRegs;
+ private:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_ret = rax;
+
+ protected:
+  void assign_regs() {
+    CRegCount = MRegs * NRegs;
+    ARegCount = 1;
+    BRegCount = RegCount - ARegCount - CRegCount;
+    if (BRegCount < NRegs) {
+      BRegCount = 0;
+      ARegCount = BRegCount + 1;
+    }
+    if (BRegCount > NRegs) {
       BRegCount = NRegs;
-      BReg = CReg + CRegCount;
-      AReg = BReg + BRegCount;
-      TmpReg = AReg + ARegCount;
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < NRegs; j++) {
-          vpxorq(Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n32", T_NEAR);
-      generate_kloop(_mtile, NRegs);
-      write_back(_mtile, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      cmp(reg_tmp, 32);
-      jl(".n16", T_NEAR);
-      generate_kloop(_mtile, 2);
-      write_back(_mtile, 2);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(int8_t) * 4);
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(_mtile, 1);
-      write_back(_mtile, 1);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(int8_t) * 4);
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
     }
+    CReg = 0;
+    BReg = CReg + CRegCount;
+    AReg = BReg + BRegCount;
+    TmpReg = AReg + ARegCount;
+    assert(TmpReg <= RegCount);
+    TmpRegCount = RegCount - TmpReg;
+  }
 
-    void generate_kloop(int _mtile, int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KTILE * KUNROLL);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_mtile, _nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, AKStepSize * KUNROLL);
-      add(reg_tmp1, BKStepSize * KUNROLL);
-      add(reg_iterk, KTILE * KUNROLL);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_mtile, _nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();
+    Xbyak::util::StackFrame st(this, 1, 10, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
+
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
 
-    void generate_fma(int _mtile, int _NRegs, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      for (int kk = 0; kk < _kunroll; kk++) {
-        lea(reg_tmp, ptr[reg_matAptr + kk * AKStepSize]);
-        for (int i = 0; i < _NRegs; i++) {
-          vmovups(Xbyak::Zmm(BReg + i), ptr[_reg_matBptr + kk * BKStepSize + i * VecBytes]);
+  void generate_fma(int _mtile, int _kunroll) {
+    for (int kk = 0; kk < _kunroll; kk++) {
+      lea(reg_tmp1, ptr[reg_matAptr + kk * AKStepSize]);
+      if (BRegCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          vmovups(vreg_t(BReg + i), ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
         }
         for (int mm = 0; mm < _mtile; mm++) {
-          vpbroadcastd(Xbyak::Zmm(AReg), ptr[reg_tmp]);
-          add(reg_tmp, reg_astep);
-          for (int i = 0; i < _NRegs; i++) {
-            vpdpbusds_evex(Xbyak::Zmm(CReg + mm * NRegs + i), Xbyak::Zmm(AReg), Xbyak::Zmm(BReg + i));
+          vpbroadcastd(vreg_t(AReg), ptr[reg_tmp1]);
+          add(reg_tmp1, reg_astride);
+          for (int i = 0; i < NRegs; i++) {
+            vpdpbusds_(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg), vreg_t(BReg + i));
           }
         }
+      } else if (BRegCount == 0) {
+        for (int mm = 0; mm < _mtile; mm += ARegCount) {
+          int mm_re = utils::remainsize(mm, _mtile, ARegCount);
+          for (int imm = 0; imm < mm_re; imm++) {
+            vpbroadcastd(vreg_t(AReg + imm), ptr[reg_tmp1]);
+            add(reg_tmp1, reg_astride);
+            for (int i = 0; i < NRegs; i++) {
+              vpdpbusds_(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg + imm),
+                         ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
+            }
+          }
+        }
+      } else {
+        assert(0);
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
-      }
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      for (int i = 0; i < _mtile; i++) {
-        for (int j = 0; j < _NRegs; j++) {
-          vpaddd(Xbyak::Zmm(CReg + i * NRegs + j), Xbyak::Zmm(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
-          vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + i * NRegs + j));
-        }
-        add(reg_matCptr, reg_cstep);
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vxor(vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j));
       }
-      L(".LEND");
-      nop();
-      outLocalLabel();
     }
-
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-  };
-
- public:
-  GemmCore_Row_NN_8x48_AVX512_VNNI() {
-    for (int i = 0; i < MTILE; i++) {
-      mCodes[i].generate_code(i + 1);
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(vreg_t(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    L(".end");
+    outLocalLabel();
   }
 
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    auto param = params{matA, matB, matC, _k, _n, _astride, _bstride, _cstride, kpos};
-    if (_m <= MTILE) {
-      mCodes[_m - 1].mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(CReg + i * NRegs + j));
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    outLocalLabel();
   }
-
- private:
-  std::array<MicroKernel, MTILE> mCodes;
 };
 
-class GemmCore_Row_NN_16x64_AMX_BF16 {
+template <int _NTILE, int _MTILE = 0>
+class AvxvnniN8P4 : protected jblas::xbyak::JitAvxvnni {
  public:
-  typedef utils::bf16 AType;
-  typedef utils::bf16 BType;
-  typedef float CType;
+  static int constexpr RegLen = 8, PackRow = 4;
+  static_assert(_NTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? (RegCount - 1) / NRegs : _MTILE;
+  static_assert(NRegs * MRegs <= RegCount - 1);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs, KTILE = 4;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAVX_VNNI;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_INT8_US_INT32;
+  typedef uint8_t AType;
+  typedef int8_t BType;
+  typedef int32_t CType;
   struct params {
     AType* matA;
+    int astride;
     BType* matB;
+    int bstride;
     CType* matC;
-    int k, msize, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-    void *workspace, *cfg;
+    int cstride;
+    int k;
+    int n;
+    int init;
   };
   typedef long long (*func_t)(params*);
 
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AMX_BF16_16x64;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 64, MTILE = 16, KTILE = 64 / sizeof(BType);
-  static int constexpr PACK_ROW = 2;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 256;
-  class MicroKernel : protected jblas::xbyak::JitAmxbf16 {
-   public:
-    friend GemmCore_Row_NN_16x64_AMX_BF16;
-    MicroKernel() {}
-    static int constexpr CReg = 0, TmpReg = 4;
-    static int constexpr NRegs = 4;
-    static int constexpr CRegCount = NRegs;
-    static int constexpr C_tilenum = 4, A_tilenum = 1, B_tilenum = 3;
-    static int constexpr CTile = 0, ATile = CTile + C_tilenum, BTile = ATile + A_tilenum;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code() {
-      reset();
-      generate_mtile();
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile() {
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-      mov(reg_tmp, ptr[parambase + OFFSET(cfg)]);
-      ldtilecfg(ptr[reg_tmp]);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < C_tilenum; i++) {
-        tilezero(Xbyak::Tmm(CTile + i));
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n48", T_NEAR);
-      generate_kloop(NRegs);
-      write_back(MTILE, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n48");
-      cmp(reg_tmp, 48);
-      jl(".n32", T_NEAR);
-      generate_kloop(3);
-      write_back(MTILE, 3);
-      add(reg_itern, 48);
-      add(reg_matBptr, 48 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      cmp(reg_tmp, 32);
-      jl(".n16", T_NEAR);
-      generate_kloop(2);
-      write_back(MTILE, 2);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(1);
-      write_back(MTILE, 1);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(BType));
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
-    }
+  int CRegCount = 0, BRegCount = 0, ARegCount = 0, TmpRegCount = 0;
+  int CReg = 0, BReg = 0, AReg = 0, TmpReg = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
 
-    void generate_kloop(int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
+ private:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_ret = rax;
+  Xbyak::Opmask msk_wr = k1;
+
+ protected:
+  void assign_regs() {
+    CRegCount = MRegs * NRegs;
+    ARegCount = 1;
+    BRegCount = RegCount - ARegCount - CRegCount;
+    if (BRegCount < NRegs) {
+      BRegCount = 0;
+      ARegCount = BRegCount + 1;
+    }
+    if (BRegCount > NRegs) {
+      BRegCount = NRegs;
     }
+    CReg = 0;
+    BReg = CReg + CRegCount;
+    AReg = BReg + BRegCount;
+    TmpReg = AReg + ARegCount;
+    assert(TmpReg <= RegCount);
+    TmpRegCount = RegCount - TmpReg;
+  }
 
-    void generate_fma(int _NTile, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      mov(reg_tmp, NTILE * 4);
-      if (_NTile <= B_tilenum) {
-        for (int kk = 0; kk < _kunroll; kk++) {
-          for (int i = 0; i < _NTile; i++) {
-            tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
-          }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();
+    Xbyak::util::StackFrame st(this, 1, 10, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
 
-          for (int mm = 0; mm < 1; mm++) {
-            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-            for (int i = 0; i < _NTile; i++) {
-              tdpbf16ps(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm), Xbyak::Tmm(BTile + i));
-            }
-          }
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
+
+  void generate_fma(int _mtile, int _kunroll) {
+    for (int kk = 0; kk < _kunroll; kk++) {
+      lea(reg_tmp1, ptr[reg_matAptr + kk * AKStepSize]);
+      if (BRegCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          vmovups(vreg_t(BReg + i), ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
         }
-      } else {
-        for (int kk = 0; kk < _kunroll; kk++) {
-          for (int i = 0; i < _NTile - 1; i++) {
-            tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
+        for (int mm = 0; mm < _mtile; mm++) {
+          vpbroadcastd(vreg_t(AReg), ptr[reg_tmp1]);
+          add(reg_tmp1, reg_astride);
+          for (int i = 0; i < NRegs; i++) {
+            vpdpbusds_(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg), vreg_t(BReg + i));
           }
-
-          for (int mm = 0; mm < 1; mm++) {
-            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-            for (int i = 0; i < _NTile - 1; i++) {
-              tdpbf16ps(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm), Xbyak::Tmm(BTile + i));
+        }
+      } else if (BRegCount == 0) {
+        for (int mm = 0; mm < _mtile; mm += ARegCount) {
+          int mm_re = utils::remainsize(mm, _mtile, ARegCount);
+          for (int imm = 0; imm < mm_re; imm++) {
+            vpbroadcastd(vreg_t(AReg + imm), ptr[reg_tmp1]);
+            add(reg_tmp1, reg_astride);
+            for (int i = 0; i < NRegs; i++) {
+              vpdpbusds_(vreg_t(CReg + mm * NRegs + i), vreg_t(AReg + imm),
+                         ptr[reg_matBptr + kk * BKStepSize + i * VecBytes]);
             }
-            tileloaddt1(Xbyak::Tmm(BTile + 0), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + (_NTile - 1) * 64]);
-            tdpbf16ps(Xbyak::Tmm(CTile + mm * C_tilenum + _NTile - 1), Xbyak::Tmm(ATile + mm), Xbyak::Tmm(BTile + 0));
           }
         }
+      } else {
+        assert(0);
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      mov(reg_tmp, dword[parambase + OFFSET(workspace)]);
-      mov(reg_tmp1, NTILE * 4);
-      for (int mm = 0; mm < 1; mm++) {
-        for (int i = 0; i < _NRegs; i++) {
-          tilestored(ptr[reg_tmp + reg_tmp1 + i * 64 + mm * 16 * NTILE * 4], Xbyak::Tmm(CTile + mm * C_tilenum + i));
-        }
-      }
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vxor(vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j), vreg_t(CReg + i * NRegs + j));
       }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vaddps(Xbyak::Zmm(CReg + j), ptr[reg_matCptr + j * VecBytes]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
+    }
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(vreg_t(CReg + i * NRegs + j), ptr[reg_matCptr + j * VecBytes]);
       }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      L(".LEND");
-      nop();
-      outLocalLabel();
+      add(reg_matCptr, reg_cstride);
     }
+    L(".end");
+    outLocalLabel();
+  }
 
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
-
- public:
-  GemmCore_Row_NN_16x64_AMX_BF16() { mCodes.generate_code(); }
-
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    assert((NTILE * MTILE * sizeof(CType))<= cachesize);
-    MicroKernel::tileconfig_t mCfg;
-    memset(&mCfg, 0, sizeof(mCfg));
-    auto param = params{matA, matB, matC, _k, _m, _n, _astride, _bstride, _cstride, kpos, tmpcache, &mCfg};
-    if (_m <= MTILE) {
-      jblas::xbyak::JitAmxtile::configure_tiles(mCfg, _m < 16 ? _m : 16, _n < 16 ? _n : 16, _k < KTILE ? _k : KTILE,
-                                                sizeof(BType), MicroKernel::A_tilenum, MicroKernel::B_tilenum,
-                                                MicroKernel::C_tilenum);
-      mCodes.mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    for (int i = 0; i < _mtile; i++) {
+      for (int j = 0; j < NRegs; j++) {
+        vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(CReg + i * NRegs + j));
+      }
+      add(reg_matCptr, reg_cstride);
     }
+    outLocalLabel();
   }
-
- private:
-  MicroKernel mCodes;
 };
 
-class GemmCore_Row_NN_16x48_AMX_BF16 {
+template <int _NTILE, int _MTILE = 0>
+class Amxbf16N16P2 : protected jblas::xbyak::JitAmxbf16 {
  public:
+  static int constexpr RegLen = 16, PackRow = 2;
+  static_assert(_NTILE % RegLen == 0);
+  static_assert(_MTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? 1 : _MTILE / RegLen;
+  static_assert(NRegs * MRegs + 2 <= TileCount);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs * RegLen, KTILE = 32;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAMX_BF16;
+  static uint32_t constexpr COMPUTE = (uint32_t)CompType::COMP_BF16_FP32;
   typedef utils::bf16 AType;
   typedef utils::bf16 BType;
   typedef float CType;
+
   struct params {
     AType* matA;
+    int astride;
     BType* matB;
+    int bstride;
     CType* matC;
-    int k, msize, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-    void *workspace, *cfg;
+    int cstride;
+    int k;
+    int n;
+    int init;
+    void* workspace;
   };
   typedef long long (*func_t)(params*);
 
-  static JBLAS_GEMM_CORE constexpr TYPE = JBLAS_GEMM_CORE::AMX_BF16_16x48;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 48, MTILE = 16, KTILE = 64 / sizeof(BType);
-  static int constexpr PACK_ROW = 2;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 240;
-  class MicroKernel : protected jblas::xbyak::JitAmxbf16 {
-   public:
-    friend GemmCore_Row_NN_16x48_AMX_BF16;
-    MicroKernel() {}
-    static int constexpr CReg = 0, TmpReg = 4;
-    static int constexpr NRegs = 3;
-    static int constexpr CRegCount = NRegs;
-    static int constexpr C_tilenum = 3, A_tilenum = 1, B_tilenum = 3;
-    static int constexpr CTile = 0, ATile = CTile + C_tilenum, BTile = ATile + A_tilenum;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code() {
-      reset();
-      generate_mtile();
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile() {
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-      mov(reg_tmp, ptr[parambase + OFFSET(cfg)]);
-      ldtilecfg(ptr[reg_tmp]);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < C_tilenum; i++) {
-        tilezero(Xbyak::Tmm(CTile + i));
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n32", T_NEAR);
-      generate_kloop(NRegs);
-      write_back(MTILE, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      cmp(reg_tmp, 32);
-      jl(".n16", T_NEAR);
-      generate_kloop(2);
-      write_back(MTILE, 2);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(1);
-      write_back(MTILE, 1);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(BType));
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
+  int TmpRegCount = RegCount;
+  int TmpReg = 0;
+  int CTileCount = 0, ATileCount = 0, BTileCount = 0;
+  int CTile = 0, ATile = 0, BTile = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
+
+ protected:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_tmp3;
+  Xbyak::Reg64 reg_ret = rax;
+
+  void assign_regs() {
+    CTileCount = NRegs * MRegs;
+    auto tile_re = TileCount - CTileCount;
+    if (tile_re - 1 >= NRegs) {
+      BTileCount = NRegs;
+      ATileCount = tile_re - BTileCount;
+    } else if (tile_re - 1 >= MRegs) {
+      ATileCount = MRegs;
+      BTileCount = tile_re - ATileCount;
+    } else {
+      ATileCount = 1;
+      BTileCount = tile_re - ATileCount;
     }
+    CTile = 0;
+    ATile = CTile + CTileCount;
+    BTile = ATile + ATileCount;
+  }
 
-    void generate_kloop(int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();  // use local label for multiple instance
+    Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_tmp3 = st.t[10];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
 
-    void generate_fma(int _NTile, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      mov(reg_tmp, NTILE * 4);
-      for (int kk = 0; kk < _kunroll; kk++) {
-        for (int i = 0; i < _NTile; i++) {
-          tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
-        }
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
+
+  void generate_fma(int _mtile, int kunrll) {
+    auto& reg_Bstride = reg_tmp1;
+    mov(reg_Bstride, NTILE * 4);
+    int mtiles = _mtile / RegLen;
 
-        for (int mm = 0; mm < 1; mm++) {
-          tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-          for (int i = 0; i < _NTile; i++) {
-            tdpbf16ps(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm), Xbyak::Tmm(BTile + i));
+    for (int kk = 0; kk < kunrll; kk++) {
+      auto& reg_Atmp = reg_tmp2;
+      if (mtiles == 1) {
+        reg_Atmp = reg_matAptr;
+      } else {
+        mov(reg_Atmp, reg_matAptr);
+      }
+      if (BTileCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          tileloaddt1(Xbyak::Tmm(BTile + i), ptr[reg_matBptr + reg_Bstride + kk * BKStepSize + i * 64]);
+        }
+        for (int mm = 0; mm < mtiles; mm++) {
+          tileloadd(Xbyak::Tmm(ATile), ptr[reg_Atmp + reg_astride + kk * AKStepSize]);
+          for (int i = 0; i < NRegs; i++) {
+            tdpbf16ps(Xbyak::Tmm(CTile + mm * NRegs + i), Xbyak::Tmm(ATile), Xbyak::Tmm(BTile + i));
+          }
+          if (mm != mtiles - 1) {
+            lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+            lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+          }
+        }
+      } else {
+        if (ATileCount == mtiles) {
+          for (int mm = 0; mm < mtiles; mm++) {
+            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_Atmp + reg_astride + kk * AKStepSize]);
+            if (mm != mtiles - 1) {
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+            }
+          }
+          for (int i = 0; i < NRegs; i++) {
+            tileloaddt1(Xbyak::Tmm(BTile), ptr[reg_matBptr + reg_Bstride + kk * BKStepSize + i * 64]);
+            for (int mm = 0; mm < mtiles; mm++) {
+              tdpbf16ps(Xbyak::Tmm(CTile + mm * NRegs + i), Xbyak::Tmm(ATile + mm), Xbyak::Tmm(BTile));
+            }
+          }
+        } else {
+          for (int mm = 0; mm < mtiles; mm++) {
+            tileloadd(Xbyak::Tmm(ATile), ptr[reg_Atmp + reg_astride + kk * AKStepSize]);
+            for (int i = 0; i < NRegs; i++) {
+              tileloaddt1(Xbyak::Tmm(BTile), ptr[reg_matBptr + reg_Bstride + kk * BKStepSize + i * 64]);
+              tdpbf16ps(Xbyak::Tmm(CTile + mm * NRegs + i), Xbyak::Tmm(ATile), Xbyak::Tmm(BTile));
+            }
+            if (mm != mtiles - 1) {
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+            }
           }
         }
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      mov(reg_tmp, dword[parambase + OFFSET(workspace)]);
-      mov(reg_tmp1, NTILE * 4);
-      for (int mm = 0; mm < 1; mm++) {
-        for (int i = 0; i < _NRegs; i++) {
-          tilestored(ptr[reg_tmp + reg_tmp1 + i * 64 + mm * 16 * NTILE * 4], Xbyak::Tmm(CTile + mm * C_tilenum + i));
-        }
-      }
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < CTileCount; i++) {
+      tilezero(Xbyak::Tmm(CTile + i));
+    }
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    int mtnum = _mtile / 16;
+    for (int mm = 0; mm < mtnum; mm++) {
+      for (int i = 0; i < NRegs; i++) {
+        tileloaddt1(Xbyak::Tmm(CTile + mm * NRegs + i), ptr[reg_matCptr + reg_cstride + i * 64]);
       }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vaddps(Xbyak::Zmm(CReg + j), ptr[reg_matCptr + j * VecBytes]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
+      if (mm != mtnum - 1) {
+        lea(reg_matCptr, ptr[reg_matCptr + 8 * reg_cstride]);
+        lea(reg_matCptr, ptr[reg_matCptr + 8 * reg_cstride]);
       }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      L(".LEND");
-      nop();
-      outLocalLabel();
     }
+    L(".end");
+    outLocalLabel();
+  }
 
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
-
- public:
-  GemmCore_Row_NN_16x48_AMX_BF16() { mCodes.generate_code(); }
-
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    assert((NTILE * MTILE * sizeof(CType)) <= cachesize);
-    MicroKernel::tileconfig_t mCfg;
-    std::memset(&mCfg, 0, sizeof(mCfg));
-    auto param = params{matA, matB, matC, _k, _m, _n, _astride, _bstride, _cstride, kpos, tmpcache, &mCfg};
-    if (_m <= MTILE) {
-      jblas::xbyak::JitAmxtile::configure_tiles(mCfg, _m < 16 ? _m : 16, _n < 16 ? _n : 16, _k < KTILE ? _k : KTILE,
-                                                sizeof(BType), MicroKernel::A_tilenum, MicroKernel::B_tilenum,
-                                                MicroKernel::C_tilenum);
-      mCodes.mKernel(&param);
-    } else {
-      assert(0);
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, dword[parambase + OFFSET(workspace)]);
+    mov(reg_tmp1, NTILE * 4);
+    for (int mm = 0; mm < MRegs; mm++) {
+      for (int i = 0; i < NRegs; i++) {
+        tilestored(ptr[reg_tmp + reg_tmp1 + i * 64 + mm * 16 * NTILE * 4], Xbyak::Tmm(CTile + mm * NRegs + i));
+      }
+    }
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    int zunroll = TmpRegCount / NRegs;
+    for (int i = 0; i < _mtile; i += zunroll) {
+      int m_re = utils::remainsize(i, _mtile, zunroll);
+      for (int im = 0; im < m_re; im++) {
+        for (int j = 0; j < NRegs; j++) {
+          vmovups(vreg_t(TmpReg + im * NRegs + j), ptr[reg_tmp + j * 64 + (i + im) * NTILE * 4]);
+          vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(TmpReg + im * NRegs + j));
+        }
+        add(reg_matCptr, reg_cstride);
+      }
     }
+    outLocalLabel();
   }
-
- private:
-  MicroKernel mCodes;
 };
 
-template <class T_A_, class T_B_, JBLAS_GEMM_CORE _ID>
-class GemmCore_Row_NN_16x64_AMX_I8 {
+template <typename AT, typename BT, int _NTILE, int _MTILE = 0>
+class Amxint8N16P4 : protected jblas::xbyak::JitAmxint8 {
  public:
-  typedef T_A_ AType;
-  typedef T_B_ BType;
+  static int constexpr RegLen = 16, PackRow = 4;
+  static_assert(_NTILE % RegLen == 0);
+  static_assert(_MTILE % RegLen == 0);
+  static int constexpr NRegs = _NTILE / RegLen;
+  static int constexpr MRegs = _MTILE == 0 ? 1 : _MTILE / RegLen;
+  static_assert(NRegs * MRegs + 2 <= TileCount);
+  static int constexpr NTILE = RegLen * NRegs, MTILE = MRegs * RegLen, KTILE = 64;
+  static int constexpr KUNROLL = 2;
+  static uint32_t constexpr ISA = (uint32_t)JBLAS_ISA::JblasAMX_INT8;
+  static uint32_t constexpr COMPUTE =
+      (uint32_t)(std::is_same_v<AT, int8_t>
+                     ? std::is_same_v<BT, int8_t> ? CompType::COMP_INT8_SS_INT32 : CompType::COMP_INT8_SU_INT32
+                 : std::is_same_v<BT, int8_t> ? CompType::COMP_INT8_US_INT32
+                                              : CompType::COMP_INT8_UU_INT32);
+  using AType = AT;
+  using BType = BT;
   typedef int32_t CType;
+
   struct params {
     AType* matA;
+    int astride;
     BType* matB;
+    int bstride;
     CType* matC;
-    int k, msize, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-    void *workspace, *cfg;
+    int cstride;
+    int k;
+    int n;
+    int init;
+    void* workspace;
   };
   typedef long long (*func_t)(params*);
 
-  static JBLAS_GEMM_CORE constexpr TYPE = _ID;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 64, MTILE = 16, KTILE = 64 / sizeof(BType);
-  static int constexpr PACK_ROW = 4;
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 256;
-  class MicroKernel : protected jblas::xbyak::JitAmxint8 {
-   public:
-    friend GemmCore_Row_NN_16x64_AMX_I8<T_A_, T_B_, _ID>;
-    MicroKernel() {}
-    static int constexpr CReg = 0, TmpReg = 4;
-    static int constexpr NRegs = 4;
-    static int constexpr CRegCount = NRegs;
-    static int constexpr C_tilenum = 4, A_tilenum = 1, B_tilenum = 3;
-    static int constexpr CTile = 0, ATile = CTile + C_tilenum, BTile = ATile + A_tilenum;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code() {
-      reset();
-      generate_mtile();
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile() {
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-      mov(reg_tmp, ptr[parambase + OFFSET(cfg)]);
-      ldtilecfg(ptr[reg_tmp]);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < C_tilenum; i++) {
-        tilezero(Xbyak::Tmm(CTile + i));
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n48", T_NEAR);
-      generate_kloop(NRegs);
-      write_back(MTILE, NRegs);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n48");
-      cmp(reg_tmp, 48);
-      jl(".n32", T_NEAR);
-      generate_kloop(3);
-      write_back(MTILE, 3);
-      add(reg_itern, 48);
-      add(reg_matBptr, 48 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      cmp(reg_tmp, 32);
-      jl(".n16", T_NEAR);
-      generate_kloop(2);
-      write_back(MTILE, 2);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(1);
-      write_back(MTILE, 1);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(BType));
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
+  int TmpRegCount = RegCount;
+  int TmpReg = 0;
+  int CTileCount = 0, ATileCount = 0, BTileCount = 0;
+  int CTile = 0, ATile = 0, BTile = 0;
+  static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
+  static int constexpr AKStepSize = KTILE * sizeof(AType);
+
+  void generate_code(int _mtile) {
+    assign_regs();
+    reset();
+    generate_mtile(_mtile);
+    ready();
+    mKernel = getCode<func_t>();
+  }
+  func_t mKernel = nullptr;
+
+ protected:
+  Xbyak::Reg64 parambase;
+  Xbyak::Reg64 reg_matAptr;
+  Xbyak::Reg64 reg_matBptr;
+  Xbyak::Reg64 reg_matCptr;
+  Xbyak::Reg64 reg_ksize;
+  Xbyak::Reg64 reg_nsize;
+  Xbyak::Reg64 reg_cstride;
+  Xbyak::Reg64 reg_astride;
+  Xbyak::Reg64 reg_iterk;
+  Xbyak::Reg64 reg_itern;
+  Xbyak::Reg64 reg_tmp;
+  Xbyak::Reg64 reg_tmp1;
+  Xbyak::Reg64 reg_tmp2;
+  Xbyak::Reg64 reg_tmp3;
+  Xbyak::Reg64 reg_ret = rax;
+
+  void assign_regs() {
+    CTileCount = NRegs * MRegs;
+    auto tile_re = TileCount - CTileCount;
+    if (tile_re - 1 >= NRegs) {
+      BTileCount = NRegs;
+      ATileCount = tile_re - BTileCount;
+    } else if (tile_re - 1 >= MRegs) {
+      ATileCount = MRegs;
+      BTileCount = tile_re - ATileCount;
+    } else {
+      ATileCount = 1;
+      BTileCount = tile_re - ATileCount;
     }
+    CTile = 0;
+    ATile = CTile + CTileCount;
+    BTile = ATile + ATileCount;
+  }
 
-    void generate_kloop(int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
-    }
+  void generate_mtile(int _mtile) {
+    inLocalLabel();  // use local label for multiple instance
+    Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
+    parambase = st.p[0];
+    reg_matAptr = st.t[0];
+    reg_matBptr = st.t[1];
+    reg_matCptr = st.t[0];
+    reg_ksize = st.t[2];
+    reg_astride = st.t[3];
+    reg_cstride = st.t[3];
+    reg_iterk = st.t[4];
+    reg_tmp = st.t[5];
+    reg_tmp1 = st.t[6];
+    reg_tmp2 = st.t[7];
+    reg_tmp3 = st.t[10];
+    reg_nsize = st.t[8];
+    reg_itern = st.t[9];
+    reg_ret = rax;
+
+    vreg_push(rsp);
+
+    load32(reg_ksize, ptr[parambase + OFFSET(k)]);
+    load32(reg_nsize, ptr[parambase + OFFSET(n)]);
+    xor_(reg_itern, reg_itern);
+    L(".nloop");
+    init_regs(_mtile);
+    mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
+    load32(reg_astride, ptr[parambase + OFFSET(astride)]);
+    mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
+    load32(reg_tmp, ptr[parambase + OFFSET(bstride)]);
+    imul(reg_tmp, reg_itern);
+    lea(reg_matBptr, ptr[reg_matBptr + reg_tmp]);
+    xor_(reg_iterk, reg_iterk);
+    generate_kloop(_mtile);
+    write_back(_mtile);
+    add(reg_itern, NTILE);
+    cmp(reg_itern, reg_nsize);
+    jb(".nloop");
+    mov(reg_ret, 0);
+    vreg_pop(rsp);
+
+    outLocalLabel();  // end of local label
+  }
 
-    void generate_fma(int _NTile, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      mov(reg_tmp, NTILE * 4);
-      if (_NTile <= B_tilenum) {
-        for (int kk = 0; kk < _kunroll; kk++) {
-          for (int i = 0; i < _NTile; i++) {
-            tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
-          }
+  void generate_kloop(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, reg_ksize);
+    padto_le(reg_tmp, KUNROLL * KTILE);
+    cmp(reg_tmp, 0);
+    jz(".kloop", T_NEAR);
+    L(".unkloop");
+    generate_fma(_mtile, KUNROLL);
+    add(reg_matAptr, KUNROLL * AKStepSize);
+    add(reg_matBptr, KUNROLL * BKStepSize);
+    add(reg_iterk, KUNROLL * KTILE);
+    cmp(reg_iterk, reg_tmp);  // k iteration variable
+    jb(".unkloop");
+    cmp(reg_tmp, reg_ksize);
+    jge(".kend", T_NEAR);
+    L(".kloop");
+    generate_fma(_mtile, 1);
+    add(reg_matAptr, 1 * AKStepSize);
+    add(reg_matBptr, 1 * BKStepSize);
+    add(reg_iterk, 1 * KTILE);
+    cmp(reg_iterk, reg_ksize);  // k iteration variable
+    jb(".kloop");
+    L(".kend");
+    outLocalLabel();
+  }
 
-          for (int mm = 0; mm < 1; mm++) {
-            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-            for (int i = 0; i < _NTile; i++) {
-              _tdpb<AType, BType>(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm),
-                                  Xbyak::Tmm(BTile + i));
-            }
+  void generate_fma(int _mtile, int kunrll) {
+    auto& reg_Bstride = reg_tmp1;
+    mov(reg_Bstride, NTILE * 4);
+    int mtiles = _mtile / RegLen;
+
+    for (int kk = 0; kk < kunrll; kk++) {
+      auto& reg_Atmp = reg_tmp2;
+      if (mtiles == 1) {
+        reg_Atmp = reg_matAptr;
+      } else {
+        mov(reg_Atmp, reg_matAptr);
+      }
+      if (BTileCount == NRegs) {
+        for (int i = 0; i < NRegs; i++) {
+          tileloaddt1(Xbyak::Tmm(BTile + i), ptr[reg_matBptr + reg_Bstride + kk * BKStepSize + i * 64]);
+        }
+        for (int mm = 0; mm < mtiles; mm++) {
+          tileloadd(Xbyak::Tmm(ATile), ptr[reg_Atmp + reg_astride + kk * AKStepSize]);
+          for (int i = 0; i < NRegs; i++) {
+            _tdpb<AT, BT>(Xbyak::Tmm(CTile + mm * NRegs + i), Xbyak::Tmm(ATile), Xbyak::Tmm(BTile + i));
+          }
+          if (mm != mtiles - 1) {
+            lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+            lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
           }
         }
       } else {
-        for (int kk = 0; kk < _kunroll; kk++) {
-          for (int i = 0; i < _NTile - 1; i++) {
-            tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
+        if (ATileCount == mtiles) {
+          for (int mm = 0; mm < mtiles; mm++) {
+            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_Atmp + reg_astride + kk * AKStepSize]);
+            if (mm != mtiles - 1) {
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+            }
           }
-
-          for (int mm = 0; mm < 1; mm++) {
-            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-            for (int i = 0; i < _NTile - 1; i++) {
-              _tdpb<AType, BType>(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm),
-                                  Xbyak::Tmm(BTile + i));
+          for (int i = 0; i < NRegs; i++) {
+            tileloaddt1(Xbyak::Tmm(BTile), ptr[reg_matBptr + reg_Bstride + kk * BKStepSize + i * 64]);
+            for (int mm = 0; mm < mtiles; mm++) {
+              _tdpb<AT, BT>(Xbyak::Tmm(CTile + mm * NRegs + i), Xbyak::Tmm(ATile + mm), Xbyak::Tmm(BTile));
+            }
+          }
+        } else {
+          for (int mm = 0; mm < mtiles; mm++) {
+            tileloadd(Xbyak::Tmm(ATile), ptr[reg_Atmp + reg_astride + kk * AKStepSize]);
+            for (int i = 0; i < NRegs; i++) {
+              tileloaddt1(Xbyak::Tmm(BTile), ptr[reg_matBptr + reg_Bstride + kk * BKStepSize + i * 64]);
+              _tdpb<AT, BT>(Xbyak::Tmm(CTile + mm * NRegs + i), Xbyak::Tmm(ATile), Xbyak::Tmm(BTile));
+            }
+            if (mm != mtiles - 1) {
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
+              lea(reg_Atmp, ptr[reg_Atmp + 8 * reg_astride]);
             }
-            tileloaddt1(Xbyak::Tmm(BTile + 0), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + (_NTile - 1) * 64]);
-            _tdpb<AType, BType>(Xbyak::Tmm(CTile + mm * C_tilenum + _NTile - 1), Xbyak::Tmm(ATile + mm),
-                                Xbyak::Tmm(BTile + 0));
           }
         }
       }
     }
+  }
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      mov(reg_tmp, dword[parambase + OFFSET(workspace)]);
-      mov(reg_tmp1, NTILE * 4);
-      for (int mm = 0; mm < 1; mm++) {
-        for (int i = 0; i < _NRegs; i++) {
-          tilestored(ptr[reg_tmp + reg_tmp1 + i * 64 + mm * 16 * NTILE * 4], Xbyak::Tmm(CTile + mm * C_tilenum + i));
-        }
+  void init_regs(int _mtile) {
+    inLocalLabel();
+    load32(reg_tmp, ptr[parambase + OFFSET(init)]);
+    cmp(reg_tmp, 0);
+    je(".read", T_NEAR);
+    for (int i = 0; i < CTileCount; i++) {
+      tilezero(Xbyak::Tmm(CTile + i));
+    }
+    jmp(".end", T_NEAR);
+    L(".read");
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    int mtnum = _mtile / 16;
+    for (int mm = 0; mm < mtnum; mm++) {
+      for (int i = 0; i < NRegs; i++) {
+        tileloaddt1(Xbyak::Tmm(CTile + mm * NRegs + i), ptr[reg_matCptr + reg_cstride + i * 64]);
+      }
+      if (mm != mtnum - 1) {
+        lea(reg_matCptr, ptr[reg_matCptr + 8 * reg_cstride]);
+        lea(reg_matCptr, ptr[reg_matCptr + 8 * reg_cstride]);
       }
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
+    }
+    L(".end");
+    outLocalLabel();
+  }
+
+  void write_back(int _mtile) {
+    inLocalLabel();
+    mov(reg_tmp, dword[parambase + OFFSET(workspace)]);
+    mov(reg_tmp1, NTILE * 4);
+    for (int mm = 0; mm < MRegs; mm++) {
+      for (int i = 0; i < NRegs; i++) {
+        tilestored(ptr[reg_tmp + reg_tmp1 + i * 64 + mm * 16 * NTILE * 4], Xbyak::Tmm(CTile + mm * NRegs + i));
       }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vpaddd(Xbyak::Zmm(CReg + j), Xbyak::Zmm(CReg + j), ptr[reg_matCptr + j * VecBytes]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
+    }
+    mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
+    load32(reg_cstride, ptr[parambase + OFFSET(cstride)]);
+    lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
+    int zunroll = TmpRegCount / NRegs;
+    for (int i = 0; i < _mtile; i += zunroll) {
+      int m_re = utils::remainsize(i, _mtile, zunroll);
+      for (int im = 0; im < m_re; im++) {
+        for (int j = 0; j < NRegs; j++) {
+          vmovups(vreg_t(TmpReg + im * NRegs + j), ptr[reg_tmp + j * 64 + (i + im) * NTILE * 4]);
+          vmovups(ptr[reg_matCptr + j * VecBytes], vreg_t(TmpReg + im * NRegs + j));
+        }
+        add(reg_matCptr, reg_cstride);
       }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      L(".LEND");
-      nop();
-      outLocalLabel();
     }
+    outLocalLabel();
+  }
+};
+template <int N, int M>
+using Amxint8N16P4US = Amxint8N16P4<uint8_t, int8_t, N, M>;
 
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
+template <int N, int M>
+using Amxint8N16P4SS = Amxint8N16P4<int8_t, int8_t, N, M>;
 
+class AmxConfigure : protected jblas::xbyak::JitAmxtile {
  public:
-  GemmCore_Row_NN_16x64_AMX_I8() { mCodes.generate_code(); }
+  typedef long long (*func_t)(tileconfig_t*);
+
+  static void configure(int TILE_M, int TILE_N, int TILE_K, int elesize, int ANum, int BNum, int CNum) {
+    static AmxConfigure code;
+    tileconfig_t cfg;
+    std::memset(&cfg, 0, sizeof(cfg));
+    configure_tiles(cfg, TILE_M, TILE_N, TILE_K, elesize, ANum, BNum, CNum);
+    code.mKernel(&cfg);
+  }
 
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
+ protected:
+  AmxConfigure() {
+    generate_config(this);
+    mKernel = getCode<func_t>();
+  }
+
+  func_t mKernel = nullptr;
+};
+
+}  // namespace code
+template <template <int, int> class CodeT, int _NTILE, int _MTILE = 0>
+class CoreCodeBase {
+ public:
+  using Code = CodeT<_NTILE, _MTILE>;
+  using AType = typename Code::AType;
+  using BType = typename Code::BType;
+  using CType = typename Code::CType;
+  static int constexpr NTILE = Code::NTILE;
+  static int constexpr MTILE = Code::MTILE;
+  static int constexpr KTILE = Code::KTILE;
+  static int constexpr PACK_ROW = Code::PackRow;
+  static int constexpr COMP = Code::COMPUTE;
+  static int constexpr PREFERED_N = 144;
+  static JBLAS_ISA constexpr ISA = (JBLAS_ISA)Code::ISA;
+  static uint32_t constexpr ID = CoreAttr::make_core_id(NTILE, PACK_ROW, COMP, ISA);
+  void configure() { (void)(0); }
+
+ protected:
+  CoreCodeBase() {
+    for (int i = 0; i < mCodes.size(); i++) {
+      mCodes[i].generate_code(i + 1);
+    }
+  }
+  std::array<Code, Code::MTILE> mCodes;
+};
+
+template <template <int, int> class CodeT, int _NTILE, int _MTILE = 0>
+class CoreCodeBaseAMX {
+ public:
+  using Code = CodeT<_NTILE, _MTILE>;
+  using AType = typename Code::AType;
+  using BType = typename Code::BType;
+  using CType = typename Code::CType;
+  static int constexpr NTILE = Code::NTILE;
+  static int constexpr MTILE = Code::MTILE;
+  static int constexpr KTILE = Code::KTILE;
+  static int constexpr PACK_ROW = Code::PackRow;
+  static int constexpr COMP = Code::COMPUTE;
+  static int constexpr PREFERED_N = 144;
+  static JBLAS_ISA constexpr ISA = (JBLAS_ISA)Code::ISA;
+  static uint32_t constexpr ID = CoreAttr::make_core_id(_NTILE, PACK_ROW, COMP, ISA);
+  Xbyak::CodeGenerator cfgcode;
+
+ protected:
+  CoreCodeBaseAMX() {
+    for (int i = 0; i < mCodes.size(); i++) {
+      mCodes[i].generate_code((i + 1) * 16);
+    }
+  }
+  std::array<Code, Code::MRegs> mCodes;
+};
+
+template <int _NTILE, int _MTILE = 0>
+class SCoreRowNAvx2 : public CoreCodeBase<code::Avx2N8P1, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBase<code::Avx2N8P1, _NTILE, _MTILE>::Code;
+  void forward(float* matA, float* matB, float* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
                int kpos, void* tmpcache, size_t cachesize) {
-    assert((NTILE * MTILE * sizeof(CType)) <= cachesize);
-    typename MicroKernel::tileconfig_t mCfg;
-    memset(&mCfg, 0, sizeof(mCfg));
-    auto param = params{matA, matB, matC, _k, _m, _n, _astride, _bstride, _cstride, kpos, tmpcache, &mCfg};
-    if (_m <= MTILE) {
-      jblas::xbyak::JitAmxint8::configure_tiles(mCfg, _m < 16 ? _m : 16, _n < 16 ? _n : 16, _k < KTILE ? _k : KTILE,
-                                                sizeof(BType), MicroKernel::A_tilenum, MicroKernel::B_tilenum,
-                                                MicroKernel::C_tilenum);
-      mCodes.mKernel(&param);
+    auto param = typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0};
+    if (_m <= Code::MTILE) {
+      this->mCodes[_m - 1].mKernel(&param);
     } else {
       assert(0);
     }
   }
+};
 
- private:
-  MicroKernel mCodes;
+template <int _NTILE, int _MTILE = 0>
+class SCoreRowNAvx512f : public CoreCodeBase<code::Avx512fN16P1, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBase<code::Avx512fN16P1, _NTILE, _MTILE>::Code;
+  void forward(float* matA, float* matB, float* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
+               int kpos, void* tmpcache, size_t cachesize) {
+    auto param = typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0};
+    if (_m <= Code::MTILE) {
+      this->mCodes[_m - 1].mKernel(&param);
+    } else {
+      assert(0);
+    }
+  }
 };
-using GemmCore_Row_NN_16x64_AMX_U8S8 =
-    GemmCore_Row_NN_16x64_AMX_I8<uint8_t, int8_t, JBLAS_GEMM_CORE::AMX_INT8_16x64_US>;
-using GemmCore_Row_NN_16x64_AMX_S8S8 =
-    GemmCore_Row_NN_16x64_AMX_I8<uint8_t, int8_t, JBLAS_GEMM_CORE::AMX_INT8_16x64_SS>;
 
-template <class T_A_, class T_B_, JBLAS_GEMM_CORE _ID>
-class GemmCore_Row_NN_16x48_AMX_I8 {
+template <int _NTILE, int _MTILE = 0>
+class HCoreRowNAvx512fp16 : public CoreCodeBase<code::Avx512fp16N32P1, _NTILE, _MTILE> {
  public:
-  typedef T_A_ AType;
-  typedef T_B_ BType;
-  typedef int32_t CType;
-  struct params {
-    AType* matA;
-    BType* matB;
-    CType* matC;
-    int k, msize, nsize;
-    int astep, bstep, cstep;
-    int kpos;
-    void *workspace, *cfg;
-  };
-  typedef long long (*func_t)(params*);
+  using Code = typename CoreCodeBase<code::Avx512fp16N32P1, _NTILE, _MTILE>::Code;
 
-  static JBLAS_GEMM_CORE constexpr TYPE = _ID;
-  static JBLAS_ISA constexpr ISA =
-      utils::gemm_core_mask<TYPE, JBLAS_GEMM_CORE::ISA_MASK, JBLAS_GEMM_CORE::ISA_SHIFT, JBLAS_ISA>();
-  static int constexpr NTILE = 48, MTILE = 16, KTILE = 64 / sizeof(BType);
-  static int constexpr PACK_ROW = 4 / sizeof(BType);
-  static int constexpr KUNROLL = 2;
-  static int constexpr PREFERED_N = 240;  // TODO?
-  class MicroKernel : protected jblas::xbyak::JitAmxint8 {
-   public:
-    friend GemmCore_Row_NN_16x48_AMX_I8<T_A_, T_B_, _ID>;
-    MicroKernel() {}
-    static int constexpr CReg = 0;
-    static int constexpr C_tilenum = 3, A_tilenum = 1, B_tilenum = 3;
-    static int constexpr CTile = 0, ATile = CTile + C_tilenum, BTile = ATile + A_tilenum;
-    static int constexpr BKStepSize = KTILE * NTILE * sizeof(BType);
-    static int constexpr AKStepSize = KTILE * sizeof(AType);
-    static int constexpr VecBytes = 64;
-
-    void generate_code() {
-      reset();
-      generate_mtile();
-      ready();
-      mKernel = getCode<func_t>();
-    }
-    func_t mKernel = nullptr;
-
-   protected:
-    void generate_mtile() {
-      inLocalLabel();  // use local label for multiple instance
-      Xbyak::util::StackFrame st(this, 1, 11, 16 * 10);
-      parambase = st.p[0];
-      reg_matAptr = st.t[0];
-      reg_matBptr = st.t[1];
-      reg_matCptr = st.t[0];
-      reg_ksize = st.t[2];
-      reg_nsize = st.t[9];
-      reg_cstep = st.t[3];
-      reg_astep = st.t[5];
-      reg_iterk = st.t[4];
-      reg_itern = st.t[7];
-      reg_tmp = st.t[6];
-      reg_tmp1 = st.t[8];
-      reg_tmp2 = st.t[10];
-      reg_ret = rax;
-
-      vreg_push(rsp);
-      mov(reg_tmp, ptr[parambase + OFFSET(cfg)]);
-      ldtilecfg(ptr[reg_tmp]);
-
-      mov(reg_matBptr, ptr[parambase + OFFSET(matB)]);
-      load32(reg_ksize, ptr[parambase + OFFSET(k)]);
-      load32(reg_nsize, ptr[parambase + OFFSET(nsize)]);
-      load32(reg_astep, ptr[parambase + OFFSET(astep)]);
-
-      xor_(reg_itern, reg_itern);
-      L(".nloop");
-      for (int i = 0; i < C_tilenum; i++) {
-        tilezero(Xbyak::Tmm(CTile + i));
-      }
-      mov(reg_matAptr, ptr[parambase + OFFSET(matA)]);
-      mov(reg_tmp1, reg_matBptr);
-
-      xor_(reg_iterk, reg_iterk);
-
-      mov(reg_tmp, reg_nsize);
-      sub(reg_tmp, reg_itern);
-      cmp(reg_tmp, NTILE);
-      jl(".n32", T_NEAR);
-      generate_kloop(C_tilenum);
-      write_back(MTILE, C_tilenum);
-      load32(reg_tmp, ptr[parambase + OFFSET(bstep)]);
-      imul(reg_tmp, reg_tmp, NTILE);
-      add(reg_matBptr, reg_tmp);
-      add(reg_itern, NTILE);
-      jmp(".nend", T_NEAR);
-
-      L(".n32");
-      cmp(reg_tmp, 32);
-      jl(".n16", T_NEAR);
-      generate_kloop(2);
-      write_back(MTILE, 2);
-      add(reg_itern, 32);
-      add(reg_matBptr, 32 * sizeof(BType));
-      jmp(".nend", T_NEAR);
-
-      L(".n16");
-      xor_(reg_iterk, reg_iterk);
-      generate_kloop(1);
-      write_back(MTILE, 1);
-      add(reg_itern, 16);
-      add(reg_matBptr, 16 * sizeof(BType));
-      L(".nend");
-      cmp(reg_itern, reg_nsize);
-      jb(".nloop");
-
-      mov(reg_ret, 0);
-      vreg_pop(rsp);
-
-      outLocalLabel();  // end of local label
+  void forward(utils::fp16* matA, utils::fp16* matB, utils::fp16* matC, int _m, int _n, int _k, int _astride,
+               int _bstride, int _cstride, int kpos, void* tmpcache, size_t cachesize) {
+    auto param = typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0};
+    if (_m <= Code::MTILE) {
+      this->mCodes[_m - 1].mKernel(&param);
+    } else {
+      assert(0);
     }
+  }
+};
 
-    void generate_kloop(int _nregs) {
-      inLocalLabel();
-      L(".kloop");
-      mov(reg_tmp, reg_ksize);
-      sub(reg_tmp, reg_iterk);
-      cmp(reg_tmp, KUNROLL * KTILE);
-      jl(".k1loop", T_NEAR);
-      generate_fma(_nregs, KUNROLL, reg_tmp1);
-      add(reg_matAptr, KUNROLL * AKStepSize);
-      add(reg_tmp1, KUNROLL * BKStepSize);
-      add(reg_iterk, KUNROLL * KTILE);
-      jmp(".kloopend", T_NEAR);
-
-      L(".k1loop");
-      generate_fma(_nregs, 1, reg_tmp1);
-      add(reg_matAptr, 1 * AKStepSize);
-      add(reg_tmp1, 1 * BKStepSize);
-      add(reg_iterk, 1 * KTILE);
-      L(".kloopend");
-      cmp(reg_iterk, reg_ksize);  // k iteration variable
-      jb(".kloop");
-      outLocalLabel();
+template <int _NTILE, int _MTILE = 0>
+class HCoreRowNAvx512bf16 : public CoreCodeBase<code::Avx512bf16N16P2, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBase<code::Avx512bf16N16P2, _NTILE, _MTILE>::Code;
+  void forward(utils::bf16* matA, utils::bf16* matB, float* matC, int _m, int _n, int _k, int _astride, int _bstride,
+               int _cstride, int kpos, void* tmpcache, size_t cachesize) {
+    auto param = typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0};
+    if (_m <= Code::MTILE) {
+      this->mCodes[_m - 1].mKernel(&param);
+    } else {
+      assert(0);
     }
+  }
+};
 
-    void generate_fma(int _NTile, int _kunroll, const Xbyak::Reg64& _reg_matBptr) {
-      mov(reg_tmp, NTILE * 4);
-      if (_NTile <= B_tilenum) {
-        for (int kk = 0; kk < _kunroll; kk++) {
-          for (int i = 0; i < _NTile; i++) {
-            tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
-          }
-
-          for (int mm = 0; mm < 1; mm++) {
-            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-            for (int i = 0; i < _NTile; i++) {
-              _tdpb<AType, BType>(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm),
-                                  Xbyak::Tmm(BTile + i));
-            }
-          }
-        }
-      } else {
-        for (int kk = 0; kk < _kunroll; kk++) {
-          for (int i = 0; i < _NTile - 1; i++) {
-            tileloaddt1(Xbyak::Tmm(BTile + i), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + i * 64]);
-          }
+template <int _NTILE, int _MTILE = 0>
+class HCoreRowNAmxbf16 : public CoreCodeBaseAMX<code::Amxbf16N16P2, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBaseAMX<code::Amxbf16N16P2, _NTILE, _MTILE>::Code;
+  using AType = typename Code::AType;
+  using BType = typename Code::BType;
+  using CType = typename Code::CType;
+
+  void configure() {
+    code::AmxConfigure::configure(16, 16, Code::KTILE, sizeof(BType), this->mCodes[0].ATileCount,
+                                  this->mCodes[0].BTileCount, this->mCodes[0].CTileCount);
+  }
 
-          for (int mm = 0; mm < 1; mm++) {
-            tileloadd(Xbyak::Tmm(ATile + mm), ptr[reg_matAptr + reg_astep + kk * AKStepSize]);
-            for (int i = 0; i < _NTile - 1; i++) {
-              _tdpb<AType, BType>(Xbyak::Tmm(CTile + mm * C_tilenum + i), Xbyak::Tmm(ATile + mm),
-                                  Xbyak::Tmm(BTile + i));
-            }
-            tileloaddt1(Xbyak::Tmm(BTile + 0), ptr[_reg_matBptr + reg_tmp + kk * BKStepSize + (_NTile - 1) * 64]);
-            _tdpb<AType, BType>(Xbyak::Tmm(CTile + mm * C_tilenum + _NTile - 1), Xbyak::Tmm(ATile + mm),
-                                Xbyak::Tmm(BTile + 0));
-          }
-        }
-      }
+  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
+               int kpos, void* tmpcache, size_t cachesize) {
+    auto param =
+        typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0, tmpcache};
+    if (_m <= Code::MTILE) {
+      int idx = utils::updiv(_m, 16) - 1;
+      this->mCodes[idx].mKernel(&param);
+    } else {
+      assert(0);
     }
+  }
+};
 
-    void write_back(int _mtile, int _NRegs) {
-      inLocalLabel();
-      mov(reg_tmp, dword[parambase + OFFSET(workspace)]);
-      mov(reg_tmp1, NTILE * 4);
-      for (int mm = 0; mm < 1; mm++) {
-        for (int i = 0; i < _NRegs; i++) {
-          tilestored(ptr[reg_tmp + reg_tmp1 + i * 64 + mm * 16 * NTILE * 4], Xbyak::Tmm(CTile + mm * C_tilenum + i));
-        }
-      }
-      load32(reg_matCptr, ptr[parambase + OFFSET(kpos)]);
-      cmp(reg_matCptr, 0);
-      jg(".LACC", T_NEAR);
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
-      }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      jmp(".LEND", T_NEAR);
-      L(".LACC");
-      mov(reg_matCptr, ptr[parambase + OFFSET(matC)]);
-      lea(reg_matCptr, ptr[reg_matCptr + reg_itern * sizeof(CType)]);
-      load32(reg_cstep, ptr[parambase + OFFSET(cstep)]);
-      inLocalLabel();
-      xor_(reg_tmp1, reg_tmp1);
-      L(".mloop");
-      for (int j = 0; j < _NRegs; j++) {
-        vmovups(Xbyak::Zmm(CReg + j), ptr[reg_tmp + j * 64]);
-        vpaddd(Xbyak::Zmm(CReg + j), Xbyak::Zmm(CReg + j), ptr[reg_matCptr + j * VecBytes]);
-        vmovups(ptr[reg_matCptr + j * VecBytes], Xbyak::Zmm(CReg + j));
-      }
-      add(reg_matCptr, reg_cstep);
-      add(reg_tmp, NTILE * 4);
-      add(reg_tmp1, 1);
-      cmp(reg_tmp1.cvt32(), ptr[parambase + OFFSET(msize)]);
-      jb(".mloop");
-      outLocalLabel();
-      L(".LEND");
-      nop();
-      outLocalLabel();
+template <int _NTILE, int _MTILE = 0>
+class ICoreRowNAvx512vnni : public CoreCodeBase<code::Avx512vnniN16P4, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBase<code::Avx512vnniN16P4, _NTILE, _MTILE>::Code;
+  void forward(uint8_t* matA, int8_t* matB, int32_t* matC, int _m, int _n, int _k, int _astride, int _bstride,
+               int _cstride, int kpos, void* tmpcache, size_t cachesize) {
+    auto param = typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0};
+    if (_m <= Code::MTILE) {
+      this->mCodes[_m - 1].mKernel(&param);
+    } else {
+      assert(0);
     }
+  }
+};
 
-   private:
-    Xbyak::Reg64 parambase;
-    Xbyak::Reg64 reg_matAptr;
-    Xbyak::Reg64 reg_matBptr;
-    Xbyak::Reg64 reg_matCptr;
-    Xbyak::Reg64 reg_ksize;
-    Xbyak::Reg64 reg_nsize;
-    Xbyak::Reg64 reg_cstep;
-    Xbyak::Reg64 reg_astep;
-    Xbyak::Reg64 reg_iterk;
-    Xbyak::Reg64 reg_itern;
-    Xbyak::Reg64 reg_tmp;
-    Xbyak::Reg64 reg_tmp1;
-    Xbyak::Reg64 reg_tmp2;
-    Xbyak::Reg64 reg_ret = rax;
-    Xbyak::Opmask msk_wr = k1;
-  };
+template <int _NTILE, int _MTILE = 0>
+class ICoreRowNAvxvnni : public CoreCodeBase<code::AvxvnniN8P4, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBase<code::AvxvnniN8P4, _NTILE, _MTILE>::Code;
+
+  void forward(uint8_t* matA, int8_t* matB, int32_t* matC, int _m, int _n, int _k, int _astride, int _bstride,
+               int _cstride, int kpos, void* tmpcache, size_t cachesize) {
+    auto param = typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0};
+    if (_m <= Code::MTILE) {
+      this->mCodes[_m - 1].mKernel(&param);
+    } else {
+      assert(0);
+    }
+  }
+};
 
+template <int _NTILE, int _MTILE = 0>
+class ICoreRowNAmxint8 : public CoreCodeBaseAMX<code::Amxint8N16P4US, _NTILE, _MTILE> {
  public:
-  GemmCore_Row_NN_16x48_AMX_I8() { mCodes.generate_code(); }
+  using Code = typename CoreCodeBaseAMX<code::Amxint8N16P4US, _NTILE, _MTILE>::Code;
+  using AType = typename Code::AType;
+  using BType = typename Code::BType;
+  using CType = typename Code::CType;
+  void configure() {
+    code::AmxConfigure::configure(16, 16, Code::KTILE, sizeof(BType), this->mCodes[0].ATileCount,
+                                  this->mCodes[0].BTileCount, this->mCodes[0].CTileCount);
+  }
 
-  void forward(AType* matA, BType* matB, CType* matC, int _m, int _n, int _k, int _astride, int _bstride, int _cstride,
-               int kpos, void* tmpcache, size_t cachesize) {
-    assert((NTILE * MTILE * sizeof(CType)) <= cachesize);
-    typename MicroKernel::tileconfig_t mCfg;
-    memset(&mCfg, 0, sizeof(mCfg));
-    auto param = params{matA, matB, matC, _k, _m, _n, _astride, _bstride, _cstride, kpos, tmpcache, &mCfg};
-    if (_m <= MTILE) {
-      jblas::xbyak::JitAmxint8::configure_tiles(mCfg, _m < 16 ? _m : 16, _n < 16 ? _n : 16, _k < KTILE ? _k : KTILE,
-                                                sizeof(BType), MicroKernel::A_tilenum, MicroKernel::B_tilenum,
-                                                MicroKernel::C_tilenum);
-      mCodes.mKernel(&param);
+  void forward(uint8_t* matA, int8_t* matB, int32_t* matC, int _m, int _n, int _k, int _astride, int _bstride,
+               int _cstride, int kpos, void* tmpcache, size_t cachesize) {
+    auto param =
+        typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0, tmpcache};
+    if (_m <= Code::MTILE) {
+      int idx = utils::updiv(_m, 16) - 1;
+      this->mCodes[idx].mKernel(&param);
     } else {
       assert(0);
     }
   }
+};
+
+template <int _NTILE, int _MTILE = 0>
+class ICoreRowNAmxint8SS : public CoreCodeBaseAMX<code::Amxint8N16P4SS, _NTILE, _MTILE> {
+ public:
+  using Code = typename CoreCodeBaseAMX<code::Amxint8N16P4SS, _NTILE, _MTILE>::Code;
+  using AType = typename Code::AType;
+  using BType = typename Code::BType;
+  using CType = typename Code::CType;
+  void configure() {
+    code::AmxConfigure::configure(16, 16, Code::KTILE, sizeof(BType), this->mCodes[0].ATileCount,
+                                  this->mCodes[0].BTileCount, this->mCodes[0].CTileCount);
+  }
 
- private:
-  MicroKernel mCodes;
+  void forward(int8_t* matA, int8_t* matB, int32_t* matC, int _m, int _n, int _k, int _astride, int _bstride,
+               int _cstride, int kpos, void* tmpcache, size_t cachesize) {
+    auto param =
+        typename Code::params{matA, _astride, matB, _bstride, matC, _cstride, _k, _n, kpos == 0 ? 1 : 0, tmpcache};
+    if (_m <= Code::MTILE) {
+      int idx = utils::updiv(_m, 16) - 1;
+      this->mCodes[idx].mKernel(&param);
+    } else {
+      assert(0);
+    }
+  }
 };
-using GemmCore_Row_NN_16x48_AMX_U8S8 =
-    GemmCore_Row_NN_16x48_AMX_I8<uint8_t, int8_t, JBLAS_GEMM_CORE::AMX_INT8_16x48_US>;
-using GemmCore_Row_NN_16x48_AMX_S8S8 = GemmCore_Row_NN_16x48_AMX_I8<int8_t, int8_t, JBLAS_GEMM_CORE::AMX_INT8_16x48_SS>;
 }  // namespace gemm
 }  // namespace jblas
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_parallel.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_parallel.h
index 76be9b3569e32..92d1c1d0c1db6 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_parallel.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_parallel.h
@@ -117,8 +117,9 @@ namespace gemm {
 
 struct ConfigGemmBase {
   int threads;
-  size_t l2cache;
   int size[3];
+  size_t l2cache = 1024ULL * 1024;
+  size_t l1cache = 32ULL * 1024;
 };
 
 struct ThreadProblemBase : ThreadProblem2D {
@@ -134,8 +135,8 @@ class SchedulerBase : public Scheduler2D {
   SchedulerBase() = default;
   SchedulerBase(const ConfigGemmBase& config) { update(config); }
   virtual void getIndex(ThreadProblem& problem) {
-    problem.l2cachesize = mL2Size;
     problem.tmpcachesize = mL2Size - mL2Use;
+    problem.l2cachesize = mL2Size;
     problem.block[0] = mBlock[0];
     problem.block[1] = mBlock[1];
     problem.block[2] = mBlock[2];
@@ -149,6 +150,7 @@ class SchedulerBase : public Scheduler2D {
     }
     mThdCount = config.threads;
     mL2Size = config.l2cache;
+    mL1Size = config.l1cache;
     if (mSize[0] <= 0 || mSize[1] <= 0 || mSize[2] <= 0) {
       return;
     }
@@ -190,17 +192,17 @@ class SchedulerBase : public Scheduler2D {
     update_cache_blocking();
     Scheduler2D::set(mThdSize, mSize, mStep);
     mL2Use = size_t(mBlock[0]) * mBlock[1] * mEleSize[2];
-    mL2Use += size_t(mBlock[0]) * mBlock[2] * mEleSize[0];
-    mL2Use += size_t(mBlock[1]) * mBlock[2] * mEleSize[0];
+    mL2Use += size_t(mBlock[1]) * mBlock[2] * mEleSize[1];
+    mL2Use += size_t(mStep[0]) * mBlock[2] * mEleSize[0];
   }
+  const float DensityThres = 32;
 
   float calculate_score() {
     int tmpnstep = mThdSize[1] < _GemmCore_T::PREFERED_N ? mThdSize[1] : _GemmCore_T::PREFERED_N;
     float threadratio = float(mThdValid) / mThdCount;
     float density = float(tmpnstep) * mThdSize[0] / (tmpnstep + mThdSize[0]);
-    const float Thres = 64;
-    if (mDensity < Thres) {
-      return (threadratio * 1.f + density * 0.0016f) * density / mDensity;
+    if (mDensity < DensityThres) {
+      return threadratio;
     }
     return (threadratio * 1.f + density * 0.0016f);
   }
@@ -218,6 +220,14 @@ class SchedulerBase : public Scheduler2D {
   // B Access = M/mMStep
   // A Access = N/mNStep
   void update_cache_blocking() {
+    if (mDensity <= DensityThres) {
+      return cache_block_memory();
+    } else {
+      return cache_blocking_compute();
+    }
+  }
+
+  void cache_blocking_compute() {
     int constexpr KRef = 256;
     size_t csize_total = mL2Size - _GemmCore_T::PREFERED_N * KRef * mEleSize[1];
     int maxM = static_cast<int>(csize_total / _GemmCore_T::PREFERED_N / mEleSize[2]);
@@ -238,27 +248,30 @@ class SchedulerBase : public Scheduler2D {
     } else {
       mBlock[1] = mThdSize[1];
     }
-    update_kstep();
-  }
-  void update_kstep() {
-    auto rawk = static_cast<int>((mL2Size / mBlock[1] - mBlock[0] * mEleSize[2]) / mEleSize[1]);
+    auto rawk = static_cast<int>((mL2Size - mBlock[0] * mBlock[1] * mEleSize[2]) /
+                                 (mBlock[0] * mEleSize[0] + mBlock[1] * mEleSize[1]));
     rawk = std::min(rawk, mSizePadded[2]);
     mBlock[2] = utils::padto_le(rawk, mStep[2]);
+  }
 
-    int const NStage = 16;
-    int const K_Split = utils::padto(utils::updiv(mSize[2], NStage), mStep[2]);
-    if (mBlock[2] > K_Split) {
-      mBlock[2] = K_Split;
-    }
+  void cache_block_memory() {
+    mBlock[0] = mThdSize[0];
+    mBlock[1] = mStep[1];
+    size_t reservsize = (size_t)mBlock[0] * mBlock[1] * mEleSize[2];
+    size_t maxK = (mL1Size - reservsize) / (mBlock[1] * mEleSize[1] + mBlock[0] * mEleSize[0]);
+    size_t Bsize = maxK * mBlock[1] * mEleSize[1];
+    size_t Bsize_1K = utils::padto_le(Bsize, 1024);
+    mBlock[2] = int(Bsize_1K / mEleSize[1] / mBlock[1]);
+    mBlock[2] = utils::padto_le(mBlock[2], mStep[2]);
   }
 
-  size_t mL2Size = 0, mL2Use = 0;
+  size_t mL2Size = 0, mL1Size = 0, mL2Use = 0;
   float mDensity = 0.f;
 
  private:
   int mSize[3] = {0, 0, 0};
   int mThdSize[3] = {0, 0, 0};
-  static constexpr int mStep[3] = {1, _GemmCore_T::NTILE, _GemmCore_T::KTILE};
+  static constexpr int mStep[3] = {_GemmCore_T::MTILE, _GemmCore_T::NTILE, _GemmCore_T::KTILE};
   static constexpr int mEleSize[3] = {sizeof(typename _GemmCore_T::AType), sizeof(typename _GemmCore_T::BType),
                                       sizeof(typename _GemmCore_T::CType)};
   int mSizePadded[3] = {0, 0, 0};
@@ -293,6 +306,7 @@ class SchedulerKBlock : public Scheduler2D {
     }
     mThdCount = config.threads;
     mL2Size = config.l2cache;
+    mL1Size = config.l1cache;
     mKBlock = config.kblock;
     if (mSize[0] <= 0 || mSize[1] <= 0 || mSize[2] <= 0) {
       return;
@@ -335,17 +349,17 @@ class SchedulerKBlock : public Scheduler2D {
     update_cache_blocking();
     Scheduler2D::set(mThdSize, mSize, mStep);
     mL2Use = size_t(mBlock[0]) * mBlock[1] * mEleSize[2] * 2;
-    mL2Use += size_t(mBlock[0]) * mBlock[2] * mEleSize[0];
-    mL2Use += size_t(mBlock[1]) * mBlock[2] * mEleSize[0];
+    mL2Use += size_t(mBlock[1]) * mBlock[2] * mEleSize[1];
+    mL2Use += size_t(mStep[0]) * mBlock[2] * mEleSize[0];
   }
+  const float DensityThres = 32;
 
   float calculate_score() {
     int tmpnstep = mThdSize[1] < _GemmCore_T::PREFERED_N ? mThdSize[1] : _GemmCore_T::PREFERED_N;
     float threadratio = float(mThdValid) / mThdCount;
     float density = float(tmpnstep) * mThdSize[0] / (tmpnstep + mThdSize[0]);
-    const float Thres = 64;
-    if (mDensity < Thres) {
-      return (threadratio * 1.f + density * 0.0016f) * density / mDensity;
+    if (mDensity < DensityThres) {
+      return threadratio * 1.f;
     }
     return (threadratio * 1.f + density * 0.0016f);
   }
@@ -362,6 +376,14 @@ class SchedulerKBlock : public Scheduler2D {
   // B=MBlock*KBlock
   // A=MTILE*KBlock
   void update_cache_blocking() {
+    if (mDensity <= DensityThres) {
+      return cache_block_memory();
+    } else {
+      return cache_blocking_compute();
+    }
+  }
+
+  void cache_blocking_compute() {
     int constexpr KRef = 256;
     int constexpr NRef = _GemmCore_T::PREFERED_N;
     int constexpr MTile = _GemmCore_T::MTILE;
@@ -393,14 +415,34 @@ class SchedulerKBlock : public Scheduler2D {
     }
   }
 
-  size_t mL2Size = 0, mL2Use = 0;
+  void cache_block_memory() {
+    mBlock[0] = _GemmCore_T::MTILE;
+    size_t startK = std::max(16, _GemmCore_T::KTILE);
+    auto getMaxN = [&](size_t refk) {
+      size_t sizeA = refk * mEleSize[0] * mBlock[0];
+      size_t maxN = (mL1Size - sizeA) / (mBlock[0] * mEleSize[2] * 2 + refk * mEleSize[1]);
+      return maxN;
+    };
+    auto maxN = getMaxN(startK);
+    if (maxN <= mThdSize[1]) {
+      mBlock[1] = int(maxN);
+      mBlock[1] = utils::padto_le(mBlock[1], mStep[1]);
+      mBlock[2] = int(startK);
+    } else {
+      mBlock[2] = int(startK * 2);
+      mBlock[1] = int(getMaxN(mBlock[2]));
+      mBlock[1] = utils::padto_le(mBlock[1], mStep[1]);
+      mBlock[1] = std::min(mBlock[1], mThdSize[1]);
+    }
+  }
+  size_t mL2Size = 0, mL1Size = 0, mL2Use = 0;
   float mDensity = 0.f;
   int mKBlock = 0;
 
  private:
   int mSize[3] = {0, 0, 0};
   int mThdSize[3] = {0, 0, 0};
-  static constexpr int mStep[3] = {1, _GemmCore_T::NTILE, _GemmCore_T::KTILE};
+  static constexpr int mStep[3] = {_GemmCore_T::MTILE, _GemmCore_T::NTILE, _GemmCore_T::KTILE};
   static constexpr int mEleSize[3] = {sizeof(typename _GemmCore_T::AType), sizeof(typename _GemmCore_T::BType),
                                       sizeof(typename _GemmCore_T::CType)};
   int mSizePadded[3] = {0, 0, 0};
@@ -458,5 +500,44 @@ class StdThreading : public IThreading {
  private:
   std::vector<std::thread> thdset;
 };
+
+template <class Parallel_T, class Launch_T>
+void GemmBaseRun(Launch_T& launcher, const typename Launch_T::Param& args, parallel::IThreading* th) {
+  device::CpuBase cb;
+  Parallel_T para({th->num_threads(), args.M, args.N, args.K, cb.mL2Cache, cb.mL1Cache});
+  static bool flag = false;
+  if (flag) {
+    printf("%s\n", __FUNCTION__);
+    para.print();
+    flag = false;
+  }
+  th->parallel_for([&](int tidx) {
+    typename Parallel_T::ThreadProblem thdp{tidx};
+    para.getIndex(thdp);
+    if (thdp.valid) {
+      launcher.run(args, thdp);
+    }
+  });
+}
+
+template <class Parallel_T, class Launch_T>
+void GemmKBlockRun(Launch_T& launcher, const typename Launch_T::Param& args, parallel::IThreading* th) {
+  device::CpuBase cb;
+  Parallel_T para({th->num_threads(), args.M, args.N, args.K, cb.mL2Cache, cb.mL1Cache, args.KBlock});
+  static bool flag = false;
+  if (flag) {
+    printf("%s\n", __FUNCTION__);
+    para.print();
+    flag = false;
+  }
+  th->parallel_for([&](int tidx) {
+    typename Parallel_T::ThreadProblem thdp{tidx};
+    para.getIndex(thdp);
+    if (thdp.valid) {
+      launcher.run(args, thdp);
+    }
+  });
+}
+
 }  // namespace parallel
 }  // namespace jblas
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_a.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_a.h
index 5fb9e52493171..f5aa76c46ff64 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_a.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_a.h
@@ -40,7 +40,7 @@ class ActivationBase {
   JBLAS_CODE getActivation(AType** dstptr, int* dststep, const Param& _param, int m_size, int k_size, int m_offset,
                            int k_offset, void* tmpcache, size_t cachesize) {
     auto aptr = const_cast<AType*>(_param.A);
-    if (k_size % _GemmCore_T::KTILE == 0) {
+    if (k_size % _GemmCore_T::KTILE == 0 && m_size >= _GemmCore_T::MTILE) {
       *dstptr = aptr + m_offset * _param.lda + k_offset;
       *dststep = _param.lda;
       return JblasSuccess;
@@ -156,8 +156,8 @@ class ActivationKBlockQuantize {
     (void)k_size;
     auto quan = _param.quan;
     auto aptr = quan->template APtr<AType>();
-    *dstptr = aptr + m_offset * _param.lda + k_offset;
-    *dststep = _param.lda;
+    *dstptr = aptr + m_offset * quan->lda + k_offset;
+    *dststep = quan->lda;
     return JblasSuccess;
   }
 };
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_b.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_b.h
index 3d30158d93815..0b43e11679062 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_b.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_prologue_b.h
@@ -51,7 +51,7 @@ class WeightPack {
   StorageType createStorage(int n, int k) {
     int KPad = utils::padto(k, _GemmCore_T::KTILE);
     int NPad = utils::padto(n, _GemmCore_T::NTILE);
-    StorageType tmp(_GemmCore_T::TYPE);
+    StorageType tmp(_GemmCore_T::ID);
     tmp.resize(NPad, KPad, n, k, utils::jblas_dtype<WType>);
     return tmp;
   }
@@ -113,7 +113,7 @@ class WeightKBlockS8 {
   StorageWeight createStorage(int n, int k, int blocksize, JBLAS_DTYPE scat, JBLAS_DTYPE redt, bool is_asym) {
     int KPad = utils::padto(k, _GemmCore_T::KTILE);
     int NPad = utils::padto(n, _GemmCore_T::NTILE);
-    StorageWeight tmp(_GemmCore_T::TYPE);
+    StorageWeight tmp(_GemmCore_T::ID);
     tmp.resize(NPad, KPad, blocksize <= 0 ? KPad : blocksize, n, k, scat, redt, is_asym);
     return tmp;
   }
@@ -306,12 +306,17 @@ class WeightKBlockS8 {
       if (stor->mRedT == JBLAS_DTYPE::F32) {
         reduce(stor->mN, stor->mK, stor->mBlockSize, deq, stor->mN, stor->template RPtr<float>(), stor->mCStep,
                threading);
+      } else if (stor->mRedT == JBLAS_DTYPE::BF16) {
+        reduce(stor->mN, stor->mK, stor->mBlockSize, deq, stor->mN, stor->template RPtr<utils::bf16>(), stor->mCStep,
+               threading);
+      } else {
+        assert(0);
       }
       utils::afree(deq);
     }
   }
-
-  void reduce(const int N, const int K, const int KBlock, const float* B, const int ldb, float* rptr, const int ldr,
+  template<typename RED_T>
+  void reduce(const int N, const int K, const int KBlock, const float* B, const int ldb, RED_T* rptr, const int ldr,
               parallel::IThreading* threading) {
     parallel::Scheduler2D _para({threading->num_threads(), K, N, KBlock, 16});
     threading->parallel_for([&](int tidx) {
@@ -320,7 +325,7 @@ class WeightKBlockS8 {
       if (thdp.valid) {
         const auto src = B + thdp.loc[0] * ldb + thdp.loc[1];
         const auto dst = rptr + thdp.loc[1] + thdp.loc[0] / KBlock * ldr;
-        using RowReduceSum = kernel::wrapper::RowReduceSum<float>;
+        using RowReduceSum = kernel::wrapper::RowReduceSum<RED_T>;
         for (int i = 0; i < thdp.size[0]; i += KBlock) {
           int rowremain = utils::remainsize(thdp.loc[0] + i, K, KBlock);
           auto ret = RowReduceSum::template forward<ISA_T>(  //
@@ -479,7 +484,7 @@ class WeightKBlockS4 : public WeightKBlockS8<_GemmCore_T, ISA_T> {
                               JBLAS_DTYPE redT, bool is_asym = false) {
     int KPad = utils::padto(K, _GemmCore_T::KTILE);
     int NPad = utils::padto(N, _GemmCore_T::NTILE);
-    StorageWeight tmp(_GemmCore_T::TYPE);
+    StorageWeight tmp(_GemmCore_T::ID);
     tmp.resize(NPad, KPad, blocksize <= 0 ? KPad : blocksize, N, K, weiT, scaT, redT, is_asym);
     return tmp;
   }
@@ -718,7 +723,7 @@ class WeightKBlockF4 : public WeightKBlockS4<_GemmCore_T, ISA_T> {
   StorageWeight createStorage(const int N, const int K, int blocksize, JBLAS_DTYPE f4T, JBLAS_DTYPE scaT) {
     int KPad = utils::padto(K, _GemmCore_T::KTILE);
     int NPad = utils::padto(N, _GemmCore_T::NTILE);
-    StorageWeight tmp(_GemmCore_T::TYPE);
+    StorageWeight tmp(_GemmCore_T::ID);
     tmp.resize(NPad, KPad, blocksize <= 0 ? KPad : blocksize, N, K, f4T, scaT);
     return tmp;
   }
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_storage.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_storage.h
index 84f0201bce86e..2c05e5a4f59ab 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_storage.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_storage.h
@@ -14,10 +14,13 @@
 #pragma once
 #include "jit_base.h"
 #include "jit_blas.h"
+#include "jit_blas_gemm.h"
 #include "jit_blas_utils.h"
 
 namespace jblas {
 namespace storage {
+
+constexpr size_t Alignment = 64;
 class ISerialObject {
  protected:
   virtual size_t getSerializedSize() = 0;
@@ -71,11 +74,12 @@ class ISerialBuffer : public ISerialObject {
   virtual size_t getSerializedSize() override {
     size_t totalsize = 0;
     totalsize += sizeof(mBufSize);
-    totalsize += mBufSize;
+    totalsize += mBufSize + Alignment;
     return totalsize;
   }
   virtual void serializeToBuffer(int8_t*& wptr) override {
     utils::serialize(wptr, mBufSize);
+    wptr = utils::pointer_align<Alignment>(wptr);
     if (wptr != (int8_t*)mBufPtr) {
       std::memcpy(wptr, mBufPtr, mBufSize);
     }
@@ -87,6 +91,7 @@ class ISerialBuffer : public ISerialObject {
     } else {
       utils::serialize<size_t>(rptr, mBufSize);
     }
+    rptr = utils::pointer_align<Alignment>(rptr);
     mBufPtr = (int8_t*)rptr;
     rptr += mBufSize;
   }
@@ -102,12 +107,12 @@ namespace gemm {
 class WeightBase : public storage::ISerializable {
  public:
   JBLAS_PROLOGUEB_IDS mPrologueID = JBLAS_PROLOGUEB_IDS::Undef;
-  JBLAS_GEMM_CORE mCoreType = JBLAS_GEMM_CORE::Undef;
+  uint32_t mCoreId = 0;
   JBLAS_DTYPE mDType = JBLAS_DTYPE::F32;
   int mNPad = 0, mKPad = 0;
   int mN = 0, mK = 0;
 
-  WeightBase(JBLAS_GEMM_CORE _type) { mCoreType = _type; }
+  WeightBase(uint32_t _id) { mCoreId = _id; }
 
   // bytes offset to mPrologueID
   static constexpr inline size_t offset() { return sizeof(mSize); }
@@ -126,7 +131,7 @@ class WeightBase : public storage::ISerializable {
   virtual void serializeToBuffer(int8_t*& wptr) {
     ISerializable::serializeToBuffer(wptr);
     utils::serialize(wptr, mPrologueID);
-    utils::serialize(wptr, mCoreType);
+    utils::serialize(wptr, mCoreId);
     utils::serialize(wptr, mNPad);
     utils::serialize(wptr, mKPad);
     utils::serialize(wptr, mN);
@@ -138,7 +143,7 @@ class WeightBase : public storage::ISerializable {
     ISerializable::deserializeBuffer(rptr, map_buf);
     if (!map_buf) {
       mPrologueID = utils::deserialize<JBLAS_PROLOGUEB_IDS>(rptr);
-      mCoreType = utils::deserialize<JBLAS_GEMM_CORE>(rptr);
+      mCoreId = utils::deserialize<uint32_t>(rptr);
       mNPad = utils::deserialize<int>(rptr);
       mKPad = utils::deserialize<int>(rptr);
       mN = utils::deserialize<int>(rptr);
@@ -146,7 +151,7 @@ class WeightBase : public storage::ISerializable {
       mDType = utils::deserialize<JBLAS_DTYPE>(rptr);
     } else {
       utils::serialize<JBLAS_PROLOGUEB_IDS>(rptr, mPrologueID);
-      utils::serialize<JBLAS_GEMM_CORE>(rptr, mCoreType);
+      utils::serialize<uint32_t>(rptr, mCoreId);
       utils::serialize<int>(rptr, mNPad);
       utils::serialize<int>(rptr, mKPad);
       utils::serialize<int>(rptr, mN);
@@ -158,7 +163,7 @@ class WeightBase : public storage::ISerializable {
   inline constexpr size_t getMiscSize() {
     size_t totalsize = 0;
     totalsize += sizeof(mPrologueID);
-    totalsize += sizeof(mCoreType);
+    totalsize += sizeof(mCoreId);
     totalsize += sizeof(mNPad);
     totalsize += sizeof(mKPad);
     totalsize += sizeof(mN);
@@ -171,7 +176,7 @@ class WeightBase : public storage::ISerializable {
 class WeightKBlockBase : public WeightBase {
  public:
   int mBlockSize = 1;
-  WeightKBlockBase(JBLAS_GEMM_CORE _type) : WeightBase(_type) {}
+  WeightKBlockBase(uint32_t _id) : WeightBase(_id) {}
   void resize(int NPad, int KPad, int Block, int N, int K, JBLAS_DTYPE dtype) {
     WeightBase::resize(NPad, KPad, N, K, dtype);
     mBlockSize = Block;
@@ -270,9 +275,9 @@ class StorageQuantCorrection : public ISerialObject {
   }
   virtual size_t getSerializedSize() override {
     size_t totalsize = getMiscSize();
-    totalsize += mCSize * mScaEleSize;
-    if (mIsAsym) totalsize += mCSize * mZpEleSize;
-    if (mHasReduce) totalsize += mCSize * mRedEleSize;
+    totalsize += mCSize * mScaEleSize + Alignment;
+    if (mIsAsym) totalsize += mCSize * mZpEleSize + Alignment;
+    if (mHasReduce) totalsize += mCSize * mRedEleSize + Alignment;
     return totalsize;
   }
   virtual void serializeToBuffer(int8_t*& wptr) override {
@@ -283,17 +288,20 @@ class StorageQuantCorrection : public ISerialObject {
     utils::serialize(wptr, mHasReduce);
     utils::serialize(wptr, mCStep);
     utils::serialize(wptr, mCSize);
+    wptr = utils::pointer_align<Alignment>(wptr);
     if (wptr != mSPtr) {
       std::memcpy(wptr, mSPtr, mScaEleSize);
     }
     wptr += mCSize * mScaEleSize;
     if (mIsAsym) {
+      wptr = utils::pointer_align<Alignment>(wptr);
       if (wptr != mZPtr) {
         std::memcpy(wptr, mZPtr, mZpEleSize);
       }
       wptr += mCSize * mZpEleSize;
     }
     if (mHasReduce) {
+      wptr = utils::pointer_align<Alignment>(wptr);
       if (wptr != mRPtr) {
         std::memcpy(wptr, mRPtr, mCSize * mRedEleSize);
       }
@@ -319,13 +327,16 @@ class StorageQuantCorrection : public ISerialObject {
       utils::serialize<int>(rptr, mCStep);
       utils::serialize<size_t>(rptr, mCSize);
     }
+    rptr = utils::pointer_align<Alignment>(rptr);
     mSPtr = rptr;
     rptr += mCSize * mScaEleSize;
     if (mIsAsym) {
+      rptr = utils::pointer_align<Alignment>(rptr);
       mZPtr = rptr;
       rptr += mCSize * mZpEleSize;
     }
     if (mHasReduce) {
+      rptr = utils::pointer_align<Alignment>(rptr);
       mRPtr = rptr;
       rptr += mCSize * mRedEleSize;
     }
@@ -480,11 +491,11 @@ class StorageQuantActivation : public ISerializable, public ISerialBuffer, publi
 
 class StoragePackedWeight : public WeightBase, public ISerialBuffer {
  public:
-  StoragePackedWeight(JBLAS_GEMM_CORE _type) : WeightBase(_type) { mPrologueID = JBLAS_PROLOGUEB_IDS::WeightPack; }
+  StoragePackedWeight(uint32_t _id) : WeightBase(_id) { mPrologueID = JBLAS_PROLOGUEB_IDS::WeightPack; }
 
   size_t resize(int NPad, int KPad, int N, int K, JBLAS_DTYPE dtype) {
     WeightBase::resize(NPad, KPad, N, K, dtype);
-    auto bsize = (size_t)NPad * KPad * utils::jblas_gemm_weight_element_size(mCoreType);
+    auto bsize = (size_t)NPad * KPad * jblas::utils::jblas_dtype_size(dtype);
     ISerialBuffer::resize(bsize);
     mSize = WeightBase::getSerializedSize() + ISerialBuffer::getSerializedSize();
     return mSize;
@@ -523,18 +534,17 @@ class StorageWeightKBlockS8 : public WeightKBlockBase, public Buffer8Bit, public
   using InfoType = WeightKBlockBase;
   using QWeightType = Buffer8Bit;
   using CorrectionType = StorageQuantCorrection;
-  StorageWeightKBlockS8(JBLAS_GEMM_CORE _type) : WeightKBlockBase(_type) {
-    mPrologueID = JBLAS_PROLOGUEB_IDS::WeightKBlockS8;
-  }
+  StorageWeightKBlockS8(uint32_t _type) : WeightKBlockBase(_type) { mPrologueID = JBLAS_PROLOGUEB_IDS::WeightKBlockS8; }
 
   size_t resize(int NPad, int KPad, int Block, int N, int K, JBLAS_DTYPE scalet, JBLAS_DTYPE redt, bool IsAsym) {
     JBLAS_DTYPE zpt = JBLAS_DTYPE::S8;
     InfoType::resize(NPad, KPad, Block, N, K, JBLAS_DTYPE::S8);
     QWeightType::resize((size_t)NPad * KPad);
     int nk_scale = utils::updiv(KPad, Block);
-    auto gemm_comp = utils::bitand_u32(mCoreType, JBLAS_GEMM_CORE::COMP_MASK);
+    auto gemm_comp = jblas::gemm::CoreAttr::get_mask_val(mCoreId, jblas::gemm::CoreAttr::COMP_MASK,
+                                                         jblas::gemm::CoreAttr::COMP_SHIFT);
     CorrectionType::resize(nk_scale, NPad, scalet, zpt, redt, IsAsym,
-                           gemm_comp >= (uint32_t)JBLAS_GEMM_CORE::COMP_INT_START);
+                           gemm_comp >= (uint32_t)jblas::gemm::CompType::COMP_INT_START);
     mSize = InfoType::getSerializedSize() + QWeightType::getSerializedSize() + CorrectionType::getSerializedSize();
     return mSize;
   }
@@ -563,9 +573,7 @@ class StorageWeightKBlockS4 : public WeightKBlockBase, public Buffer4Bit, public
   using InfoType = WeightKBlockBase;
   using QWeightType = Buffer4Bit;
   using CorrectionType = StorageQuantCorrection;
-  StorageWeightKBlockS4(JBLAS_GEMM_CORE _type) : WeightKBlockBase(_type) {
-    mPrologueID = JBLAS_PROLOGUEB_IDS::WeightKBlockS4;
-  }
+  StorageWeightKBlockS4(uint32_t _type) : WeightKBlockBase(_type) { mPrologueID = JBLAS_PROLOGUEB_IDS::WeightKBlockS4; }
 
   size_t resize(int NPad, int KPad, int Block, int N, int K, JBLAS_DTYPE s4t, JBLAS_DTYPE scalet, JBLAS_DTYPE redt,
                 bool IsAsym) {
@@ -573,9 +581,10 @@ class StorageWeightKBlockS4 : public WeightKBlockBase, public Buffer4Bit, public
     InfoType::resize(NPad, KPad, Block, N, K, s4t);
     QWeightType::resize((size_t)NPad * KPad);
     int nk_scale = utils::updiv(KPad, Block);
-    auto gemm_comp = utils::bitand_u32(mCoreType, JBLAS_GEMM_CORE::COMP_MASK);
+    auto gemm_comp = jblas::gemm::CoreAttr::get_mask_val(mCoreId, jblas::gemm::CoreAttr::COMP_MASK,
+                                                         jblas::gemm::CoreAttr::COMP_SHIFT);
     CorrectionType::resize(nk_scale, NPad, scalet, zpt, redt, IsAsym,
-                           gemm_comp >= (uint32_t)JBLAS_GEMM_CORE::COMP_INT_START);
+                           gemm_comp >= (uint32_t)jblas::gemm::CompType::COMP_INT_START);
     mSize = InfoType::getSerializedSize() + QWeightType::getSerializedSize() + CorrectionType::getSerializedSize();
     return mSize;
   }
@@ -601,7 +610,7 @@ class StorageWeightKBlockS4 : public WeightKBlockBase, public Buffer4Bit, public
 
 class StorageWeightKBlockF4 : public StorageWeightKBlockS4 {
  public:
-  StorageWeightKBlockF4(JBLAS_GEMM_CORE _type) : StorageWeightKBlockS4(_type) {
+  StorageWeightKBlockF4(uint32_t _type) : StorageWeightKBlockS4(_type) {
     mPrologueID = JBLAS_PROLOGUEB_IDS::WeightKBlockF4;
   }
 
@@ -630,16 +639,16 @@ class PackedWeightParser {
       auto type = static_cast<JBLAS_PROLOGUEB_IDS>(mProID);
       switch (type) {
         case JBLAS_PROLOGUEB_IDS::WeightPack:
-          ptr = new gemm::StoragePackedWeight(JBLAS_GEMM_CORE::Undef);
+          ptr = new gemm::StoragePackedWeight(0);
           break;
         case JBLAS_PROLOGUEB_IDS::WeightKBlockS8:
-          ptr = new gemm::StorageWeightKBlockS8(JBLAS_GEMM_CORE::Undef);
+          ptr = new gemm::StorageWeightKBlockS8(0);
           break;
         case JBLAS_PROLOGUEB_IDS::WeightKBlockS4:
-          ptr = new gemm::StorageWeightKBlockS4(JBLAS_GEMM_CORE::Undef);
+          ptr = new gemm::StorageWeightKBlockS4(0);
           break;
         case JBLAS_PROLOGUEB_IDS::WeightKBlockF4:
-          ptr = new gemm::StorageWeightKBlockF4(JBLAS_GEMM_CORE::Undef);
+          ptr = new gemm::StorageWeightKBlockF4(0);
           break;
         default:
           break;
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_utils.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_utils.h
index 4184ccee1a37d..26043a013420e 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_utils.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_utils.h
@@ -89,11 +89,6 @@ inline uint32_t bitand_u32(const T& src, const T& src1) {
   return uint32_t(src) & uint32_t(src1);
 }
 
-template <JBLAS_GEMM_CORE SRC, JBLAS_GEMM_CORE Mask, JBLAS_GEMM_CORE Shift, typename DST_T>
-inline constexpr DST_T gemm_core_mask() {
-  return static_cast<DST_T>((uint32_t(SRC) & uint32_t(Mask)) >> uint32_t(Shift));
-}
-
 struct bf16 {
   uint16_t x;
   union bf16f32 {
@@ -104,15 +99,32 @@ struct bf16 {
   bf16() : x(0) {}
 
 #if CompileBF16()
+#pragma GCC push_options
 #pragma GCC target("avx512vl", "avx512bf16")
-  explicit bf16(float vf32) : x(bit_cast<uint16_t>(_mm_cvtness_sbh(vf32))) {}
+  static uint16_t f32_to_bf16(float v) {
+    auto mm = _mm_load_ss(&v);
+    auto mm2 = _mm_cvtneps_pbh(mm);
+    uint16_t dst;
+    _mm_storeu_si16(reinterpret_cast<uint16_t*>(&dst), reinterpret_cast<__m128i>(mm2));
+    return dst;
+  }
+#pragma GCC pop_options
+  explicit bf16(float vf32) : x(bit_cast<uint16_t>(f32_to_bf16(vf32))) {}
 #else
   explicit bf16(float vf32) { fromfloat(vf32); }
 #endif
 
 #if CompileBF16()
+#pragma GCC push_options
 #pragma GCC target("avx512vl", "avx512bf16")
-  float tofloat() const { return static_cast<float>(bit_cast<__bf16>(this->x)); }
+  float tofloat() const {
+    auto mm = _mm_loadu_si16(&(this->x));
+    auto mm2 = _mm_bslli_si128(mm, 2);
+    float dst;
+    _mm_store_ss(&dst, reinterpret_cast<__m128>(mm2));
+    return dst;
+  }
+#pragma GCC pop_options
 #else
   float tofloat() const {
     bf16f32 tmp = {0.f};
@@ -121,6 +133,12 @@ struct bf16 {
   }
 #endif
 
+  float tofloat_nosimd() const {
+    bf16f32 tmp = {0.f};
+    tmp.bf16[1] = x;
+    return tmp.f32;
+  }
+
   operator float() const { return tofloat(); }
 
   static bf16 from_bin(const uint16_t x) {
@@ -131,7 +149,7 @@ struct bf16 {
 
   void fromfloat(float _v) {
 #if CompileBF16()
-    x = bit_cast<uint16_t>(_mm_cvtness_sbh(_v));
+    x = bit_cast<uint16_t>(f32_to_bf16(_v));
 #else
     bf16f32 tmp = {0.f};
     tmp.f32 = _v;
@@ -141,6 +159,16 @@ struct bf16 {
     x = tmp.bf16[1];
 #endif
   }
+
+  void fromfloat_nosimd(float _v) {
+    bf16f32 tmp = {0.f};
+    tmp.f32 = _v;
+    // See document of VCVTNEPS2BF16 in Intel® 64 and IA-32 Architectures
+    // Software Developer’s Manual Volume 2
+    const auto lsb = tmp.bf16[1] & 1;
+    tmp.u += 0x7fff + lsb;
+    x = tmp.bf16[1];
+  }
 };
 
 struct fp16 {
@@ -295,65 +323,12 @@ inline constexpr const char* dtype_str() {
   return dtype2str(DT);
 }
 
-inline const char* gemmcore2str(JBLAS_GEMM_CORE GC) {
-  switch (GC) {
-    case JBLAS_GEMM_CORE::AVX2_4X24:
-      return "AVX2_4X24";
-    case JBLAS_GEMM_CORE::AVX2_2X48:
-      return "AVX2_2X48";
-    case JBLAS_GEMM_CORE::AVX512F_8x48:
-      return "AVX512F_8x48";
-    case JBLAS_GEMM_CORE::AMX_BF16_16x64:
-      return "AMX_BF16_16x64";
-    case JBLAS_GEMM_CORE::AMX_BF16_16x48:
-      return "AMX_BF16_16x48";
-    case JBLAS_GEMM_CORE::AVX512_FP16_8x64:
-      return "AVX512_FP16_8x64";
-    case JBLAS_GEMM_CORE::AVX512_FP16_8x96:
-      return "AVX512_FP16_8x96";
-    case JBLAS_GEMM_CORE::AVX_VNNI_2x48:
-      return "AVX_VNNI_2x48";
-    case JBLAS_GEMM_CORE::AVX_VNNI_4x24:
-      return "AVX_VNNI_4x24";
-    case JBLAS_GEMM_CORE::AVX512_VNNI_8x48:
-      return "AVX512_VNNI_8x48";
-    case JBLAS_GEMM_CORE::AMX_INT8_16x64_US:
-      return "AMX_INT8_16x64_US";
-    case JBLAS_GEMM_CORE::AMX_INT8_16x64_SS:
-      return "AMX_INT8_16x64_SS";
-    case JBLAS_GEMM_CORE::AMX_INT8_16x48_US:
-      return "AMX_INT8_16x48_US";
-    case JBLAS_GEMM_CORE::AMX_INT8_16x48_SS:
-      return "AMX_INT8_16x48_SS";
-    default:
-      return "Err_Type";
-  }
-}
-
-template <JBLAS_GEMM_CORE GC>
-inline constexpr const char* gemmcore_str() {
-  return gemmcore2str(GC);
-}
 
 inline constexpr size_t jblas_dtype_size(const JBLAS_DTYPE t) {
   auto bits = (uint32_t)t & 0xff;
   return bits >> 3;  // bits to bytes
 }
 
-static inline int jblas_gemm_weight_element_size(JBLAS_GEMM_CORE _gemm) {
-  auto PACKROW = (JBLAS_GEMM_CORE)((uint32_t)_gemm & 0xff00);
-  switch (PACKROW) {
-    case JBLAS_GEMM_CORE::PACKROW_1:
-      return 4;
-    case JBLAS_GEMM_CORE::PACKROW_4:
-      return 1;
-    case JBLAS_GEMM_CORE::PACKROW_2:
-      return 2;
-    case JBLAS_GEMM_CORE::Undef:
-    default:
-      return 0;
-  }
-}
 
 #ifndef _WIN32
 static void request_perm_xtile_data() {
@@ -450,6 +425,12 @@ _T deserialize(int8_t*& buf) {
 static inline int padto(int a, int b) { return updiv(a, b) * b; }
 static inline size_t padto(size_t a, int b) { return updiv(a, b) * b; }
 
+template <int _Alignment, typename _T>
+static inline _T* pointer_align(_T* src) {
+  auto uptr = reinterpret_cast<uint64_t>(src);
+  return reinterpret_cast<_T*>((uptr + _Alignment - 1) / _Alignment * _Alignment);
+}
+
 template <typename _T>
 static inline _T* amalloc(size_t _size, size_t _alignment = 64) {
   if (_size == 0) {
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h
index 88609a2d2d56b..fd1255fcf99ed 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/jit_blas_wrapper.h
@@ -54,11 +54,16 @@ class LauncherBase {
   Epilogue mEpilogue;
 
   void run(const Param& _param, const parallel::gemm::ThreadProblemBase& _config) {
+    mGemmCore.configure();
     auto StackTmp = alloca(_config.l2cachesize);
     auto tmpB = (BType*)(StackTmp);
+    tmpB = utils::pointer_align<64>(tmpB);
     auto tmpA = (AType*)(tmpB + (size_t)_config.block[1] * _config.block[2]);
+    tmpA = utils::pointer_align<64>(tmpA);
     auto tmpC = (CType*)(tmpA + (size_t)GemmCore::MTILE * _config.block[2]);
+    tmpC = utils::pointer_align<64>(tmpC);
     auto tmpCache = (void*)(tmpC + (size_t)_config.block[0] * _config.block[1]);
+    tmpCache = utils::pointer_align<64>(tmpCache);
     for (int itern = 0; itern < _config.size[1]; itern += _config.block[1]) {
       int n_remain = utils::remainsize(itern, _config.size[1], _config.block[1]);
       for (int iterm = 0; iterm < _config.size[0]; iterm += _config.block[0]) {
@@ -144,12 +149,18 @@ class LauncherKBlock {
   Epilogue mEpilogue;
 
   void run(const Param& _param, const parallel::gemm::ThreadProblemBase& _config) {
+    mGemmCore.configure();
     auto StackTmp = alloca(_config.l2cachesize);
     auto tmpB = (BType*)(StackTmp);
+    tmpB = utils::pointer_align<64>(tmpB);
     auto tmpA = (AType*)(tmpB + (size_t)_config.block[1] * _config.block[2]);
+    tmpA = utils::pointer_align<64>(tmpA);
     auto tmpC = (AccType*)(tmpA + (size_t)GemmCore::MTILE * _config.block[2]);
+    tmpC = utils::pointer_align<64>(tmpC);
     auto tmpBlk = (CType*)(tmpC + (size_t)_config.block[0] * _config.block[1]);
+    tmpBlk = utils::pointer_align<64>(tmpBlk);
     auto tmpCache = (void*)(tmpBlk + (size_t)_config.block[0] * _config.block[1]);
+    tmpCache = utils::pointer_align<64>(tmpCache);
     for (int itern = 0; itern < _config.size[1]; itern += _config.block[1]) {
       int n_remain = utils::remainsize(itern, _config.size[1], _config.block[1]);
       for (int iterm = 0; iterm < _config.size[0]; iterm += _config.block[0]) {
@@ -168,42 +179,47 @@ class LauncherKBlock {
   void run_block(const Param& _param, const parallel::gemm::ThreadProblemBase& _config, int blk_m, int blk_n,
                  int blk_msize, int blk_nsize, AType* tmpA, BType* tmpB, CType* tmpBlk, AccType* tmpC, void* tmpcache) {
     int n_padded = utils::padto(blk_nsize, GemmCore::NTILE);
-    assert(_config.block[2] == _param.KBlock);
     for (int iterk = 0; iterk < _param.K; iterk += _config.block[2]) {
       int k_remain = utils::remainsize(iterk, _param.K, _config.block[2]);
       int k_padded = utils::padto(k_remain, GemmCore::KTILE);
-      int k_paddedle = utils::padto_le(k_remain, GemmCore::KTILE);
       auto bptr_cache = tmpB;
       int bcache_step = 0;
       mProB.getKBlockWeight(&bptr_cache, &bcache_step, k_padded, n_padded, iterk, _config.loc[1] + blk_n, _param.paramB,
                             tmpcache, _config.tmpcachesize);
       int bcache_stride = bcache_step * sizeof(BType);
-      for (int i = 0; i < blk_msize; i += GemmCore::MTILE) {
-        int m_remain = utils::remainsize(i, blk_msize, GemmCore::MTILE);
-        auto cptr_cache = tmpBlk + i * _config.block[1];
-        int ccache_stride = _config.block[1] * sizeof(CType);
-        if (k_paddedle) {
-          AType* aptr_cache = tmpA;
-          int acache_step = 0;
-          mProA.getActivation(&aptr_cache, &acache_step, _param.paramA, m_remain, k_paddedle,
-                              (blk_m + i + _config.loc[0]), iterk, tmpcache, _config.tmpcachesize);
-          mGemmCore.forward(aptr_cache, bptr_cache, cptr_cache, m_remain, n_padded, k_paddedle,
-                            acache_step * sizeof(AType), bcache_stride, ccache_stride, 0, tmpcache,
-                            _config.tmpcachesize);
-        }
-        int k_tail = k_remain - k_paddedle;
-        if (k_tail) {
-          AType* aptr_cache = tmpA;
-          int acache_step = 0;
-          mProA.getActivation(&aptr_cache, &acache_step, _param.paramA, m_remain, k_tail, (blk_m + i + _config.loc[0]),
-                              iterk + k_paddedle, tmpcache, _config.tmpcachesize);
-          mGemmCore.forward(aptr_cache, bptr_cache + k_paddedle * GemmCore::NTILE, cptr_cache, m_remain, n_padded,
-                            k_tail, acache_step * sizeof(AType), bcache_stride, ccache_stride, 0 + k_paddedle, tmpcache,
-                            _config.tmpcachesize);
+
+      for (int ikk = 0; ikk < k_remain; ikk += _param.KBlock) {
+        int k_remain1 = utils::remainsize(iterk + ikk, _param.K, _param.KBlock);
+        int k_paddedle1 = utils::padto_le(k_remain1, GemmCore::KTILE);
+        for (int i = 0; i < blk_msize; i += GemmCore::MTILE) {
+          int m_remain = utils::remainsize(i, blk_msize, GemmCore::MTILE);
+          auto cptr_cache = tmpBlk + i * _config.block[1];
+          int ccache_stride = _config.block[1] * sizeof(CType);
+          if (k_paddedle1) {
+            AType* aptr_cache = tmpA;
+            int acache_step = 0;
+            mProA.getActivation(&aptr_cache, &acache_step, _param.paramA, m_remain, k_paddedle1,
+                                (blk_m + i + _config.loc[0]), iterk + ikk, tmpcache, _config.tmpcachesize);
+            mGemmCore.forward(aptr_cache, bptr_cache + ikk * GemmCore::NTILE, cptr_cache, m_remain, n_padded,
+                              k_paddedle1, acache_step * sizeof(AType), bcache_stride, ccache_stride, 0, tmpcache,
+                              _config.tmpcachesize);
+          }
+          int k_tail = k_remain1 - k_paddedle1;
+          if (k_tail) {
+            AType* aptr_cache = tmpA;
+            int acache_step = 0;
+            mProA.getActivation(&aptr_cache, &acache_step, _param.paramA, m_remain, k_tail,
+                                (blk_m + i + _config.loc[0]), iterk + ikk + k_paddedle1, tmpcache,
+                                _config.tmpcachesize);
+            mGemmCore.forward(aptr_cache, bptr_cache + (ikk + k_paddedle1) * GemmCore::NTILE, cptr_cache, m_remain,
+                              n_padded, k_tail, acache_step * sizeof(AType), bcache_stride, ccache_stride,
+                              0 + k_paddedle1, tmpcache, _config.tmpcachesize);
+          }
         }
+        mBlockEpi.forward(tmpBlk, tmpC, _config.block[1], (_config.loc[0] + blk_m), _config.loc[1] + blk_n,
+                          (iterk + ikk) / _param.KBlock, blk_msize, blk_nsize, _param.paramBlk, tmpcache,
+                          _config.tmpcachesize);
       }
-      mBlockEpi.forward(tmpBlk, tmpC, _config.block[1], (_config.loc[0] + blk_m), _config.loc[1] + blk_n,
-                        iterk / _param.KBlock, blk_msize, blk_nsize, _param.paramBlk, tmpcache, _config.tmpcachesize);
     }
     auto cachewithblk = _config.tmpcachesize + (size_t)_config.block[0] * _config.block[1] * sizeof(CType);
     mEpilogue.forward(tmpC, _config.block[1], (_config.loc[0] + blk_m), _config.loc[1] + blk_n, blk_msize, blk_nsize,
@@ -215,6 +231,7 @@ class LauncherKBlock {
                        void* tmpcache) {
     int n_padded = utils::padto(blk_nsize, GemmCore::NTILE);
     for (int iterk = 0; iterk < _param.K; iterk += _param.KBlock) {
+      memset(tmpBlk, 0, sizeof(CType) * blk_msize * _config.block[1]);
       for (int iblkk = 0; iblkk < _param.KBlock; iblkk += _config.block[2]) {
         int k_remain = utils::remainsize(iterk + iblkk, _param.K, _config.block[2]);
         int k_padded = utils::padto(k_remain, GemmCore::KTILE);
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512_bf16.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512_bf16.h
index fdd259f303ca4..e319aaba7ce96 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512_bf16.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512_bf16.h
@@ -19,7 +19,10 @@
 namespace jblas {
 namespace kernel {
 namespace avx512_bf16 {
-
+#if CompileBF16()
+#pragma GCC push_options
+#pragma GCC target("avx512bf16", "avx512vl", "avx512bw")
+#endif
 static inline JBLAS_CODE bf16_cvt_fp32_2D_write_back(const utils::bf16* src_ptr, float* dst_ptr, int row, int col,
                                                      int src_step, int dst_step, bool zeropadding) {
 #if CompileBF16()
@@ -33,10 +36,13 @@ static inline JBLAS_CODE bf16_cvt_fp32_2D_write_back(const utils::bf16* src_ptr,
     auto dst = dst_ptr + i * dst_step;
     int j = 0;
     for (; j < col_body; j += simd_proc_elt)
-      _mm512_storeu_ps(dst + j, _mm512_cvtpbh_ps((__m256bh)_mm256_loadu_ps(reinterpret_cast<float*>(src + j))));
+      _mm512_storeu_ps(
+          dst + j,  //
+          reinterpret_cast<__m512>(_mm512_bslli_epi128(_mm512_cvtepu16_epi32(_mm256_loadu_epi16(src + j)), 2)));
     if (col_tail > 0)
-      _mm512_mask_storeu_ps(dst + j, tail_mask,
-                            _mm512_cvtpbh_ps((__m256bh)_mm256_loadu_ps(reinterpret_cast<float*>(src + j))));
+      _mm512_mask_storeu_ps(
+          dst + j, tail_mask,
+          reinterpret_cast<__m512>(_mm512_bslli_epi128(_mm512_cvtepu16_epi32(_mm256_loadu_epi16(src + j)), 2)));
     if (zeropadding && npadding) std::memset(dst + col, 0, npadding);
   }
   return JblasSuccess;
@@ -78,7 +84,9 @@ static inline JBLAS_CODE fp32_cvt_bf16_2D_write_back(const void* raw_srcptr, voi
 #endif
   return avx512f::fp32_cvt_bf16_2D_write_back(raw_srcptr, raw_dstptr, row, col, srcstride, dststride, zeropadding);
 }
-
+#if CompileBF16()
+#pragma GCC pop_options
+#endif
 }  // namespace avx512_bf16
 }  // namespace kernel
 }  // namespace jblas
\ No newline at end of file
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512f.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512f.h
index 36d117364f39d..dfe701e541b3d 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512f.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_avx512f.h
@@ -1275,23 +1275,42 @@ static inline JBLAS_CODE dequant_s32_fp32(const int32_t* srcptr, const int srcst
                                           const int row, const int col, const float* scaleA, const int ldsa,
                                           const SCAB_T* scaleB) {
   int col16 = utils::padto_le(col, 16);
+  int col64 = utils::padto_le(col, 64);
   for (int irow = 0; irow < row; irow++) {
     auto scale = scaleA[irow * ldsa];
     auto valpha = _mm512_set1_ps(scale);
     int icol = 0;
-    for (; icol < col16; icol += 16) {
-      __m512 vwscale;
-      if constexpr (std::is_same_v<SCAB_T, float>) {
-        vwscale = _mm512_loadu_ps(scaleB + icol);
-      } else if constexpr (std::is_same_v<SCAB_T, utils::bf16>) {
-        auto tmp = _mm256_loadu_si256((const __m256i*)(scaleB + icol));
-        vwscale = zmm_cvt_bf16_fp32(tmp);
+    for (; icol < col64; icol += 64) {
+      for (int ic = 0; ic < 4; ic++) {
+        __m512 vwscale;
+        if constexpr (std::is_same_v<SCAB_T, float>) {
+          vwscale = _mm512_loadu_ps(scaleB + icol + ic * 16);
+        } else if constexpr (std::is_same_v<SCAB_T, utils::bf16>) {
+          auto tmp = _mm256_loadu_si256((const __m256i*)(scaleB + icol + ic * 16));
+          vwscale = zmm_cvt_bf16_fp32(tmp);
+        }
+        auto vscale = _mm512_mul_ps(valpha, vwscale);
+        auto vsrcd = _mm512_loadu_si512(srcptr + irow * srcstep + icol + ic * 16);
+        auto vsrc = _mm512_cvtepi32_ps(vsrcd);
+        vsrc = _mm512_mul_ps(vsrc, vscale);
+        _mm512_storeu_ps(dstptr + irow * dststep + icol + ic * 16, vsrc);
+      }
+    }
+    if (icol + 16 <= col16) {
+      for (; icol < col16; icol += 16) {
+        __m512 vwscale;
+        if constexpr (std::is_same_v<SCAB_T, float>) {
+          vwscale = _mm512_loadu_ps(scaleB + icol);
+        } else if constexpr (std::is_same_v<SCAB_T, utils::bf16>) {
+          auto tmp = _mm256_loadu_si256((const __m256i*)(scaleB + icol));
+          vwscale = zmm_cvt_bf16_fp32(tmp);
+        }
+        auto vscale = _mm512_mul_ps(valpha, vwscale);
+        auto vsrcd = _mm512_loadu_si512(srcptr + irow * srcstep + icol);
+        auto vsrc = _mm512_cvtepi32_ps(vsrcd);
+        vsrc = _mm512_mul_ps(vsrc, vscale);
+        _mm512_storeu_ps(dstptr + irow * dststep + icol, vsrc);
       }
-      auto vscale = _mm512_mul_ps(valpha, vwscale);
-      auto vsrcd = _mm512_loadu_si512(srcptr + irow * srcstep + icol);
-      auto vsrc = _mm512_cvtepi32_ps(vsrcd);
-      vsrc = _mm512_mul_ps(vsrc, vscale);
-      _mm512_storeu_ps(dstptr + irow * dststep + icol, vsrc);
     }
     for (; icol < col; icol += 1) {
       dstptr[irow * dststep + icol] = scale * scaleB[icol] * srcptr[irow * srcstep + icol];
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_jit.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_jit.h
index 22c785353826d..e1e61372c38a9 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_jit.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_jit.h
@@ -245,7 +245,7 @@ class JitMemcpy2DAvx2 : protected jblas::xbyak::JitAvx2 {
 
   template <typename _SRC_T, typename _DST_T, typename... Eltops>
   static JBLAS_CODE forward(const _SRC_T* srcptr, _DST_T* dstptr, int row, int col, int srcstep, int dststep,
-                            void* elt_const_v = nullptr, Eltops... ops) {
+                            void* elt_const_v = nullptr, const Eltops&... ops) {
     if (col * sizeof(_SRC_T) % 4 != 0) {
       return JblasNotSupport;
     }
@@ -279,7 +279,7 @@ class JitMemcpy2DAvx2 : protected jblas::xbyak::JitAvx2 {
 
   template <typename _SRC_T, typename _DST_T, JBLAS_ELTWISEOP Op>
   static JBLAS_CODE forward1(const _SRC_T* srcptr, _DST_T* dstptr, int row, int col, int srcstep, int dststep,
-                            void* elt_const_v = nullptr) {
+                             void* elt_const_v = nullptr) {
     if (col * sizeof(_SRC_T) % 4 != 0) {
       return JblasNotSupport;
     }
@@ -475,7 +475,7 @@ class JitMemcpy2DAvx512f : protected jblas::xbyak::JitAvx512f {
 
   template <typename _SRC_T, typename _DST_T, typename... Eltops>
   static JBLAS_CODE forward(const _SRC_T* srcptr, _DST_T* dstptr, int row, int col, int srcstep, int dststep,
-                            void* elt_const_v = nullptr, Eltops... ops) {
+                            void* elt_const_v = nullptr, const Eltops&... ops) {
     static std::vector<kernel::jit_injector::eltwise_injector> p = {static_cast<JBLAS_ELTWISEOP>(ops)...};
     if constexpr (sizeof...(ops) != 0)
       static_assert(std::is_same<_SRC_T, float>::value && std::is_same<_DST_T, float>::value);
@@ -1205,6 +1205,168 @@ class PaddingTransInterleaveCvt : protected xbyak::JitAvx512f {
   }
 };
 
+// Complex number matrix(interleaved) - vector(as diagonal matrix) multiplication; Typically used for
+// shift-RoPE
+//
+// vector: fp16 values; view every adjacent 2 values on colunm as a complex num
+// src: bf16 ⌈row/row_pack⌉ x n_tile x row_pack; view every adjacent 2 values on colunm as a complex num
+// dst: same as src
+class CScaleInterleavedBF16FP16 : protected xbyak::JitAvx512_fp16 {
+ public:
+  struct params {
+    void* srcptr;
+    const void* scaleptr;
+    int row;
+  };
+  typedef void (*func_t)(params* p);
+  void operator()(params* p) const { mKernel(p); }
+
+ private:
+  explicit CScaleInterleavedBF16FP16(int n_tile, int n_off, int row_pack = 2, int unroll = 2)
+      : xbyak::JitAvx512_fp16() {
+    inLocalLabel();  // use local label for multiple instance
+    assert(("n_tile must be a multiple of 16", n_tile % 16 == 0));
+    assert(row_pack > 0 && row_pack < 3);  // TODO(yi): int8 interleave not implemented
+    int SF_TmpSize = 64;
+    std::shared_ptr<void> epilogue{// generate code at the very end
+                                   nullptr, [&](void*) {
+                                     outLocalLabel();  // end of local label
+                                     this->ready();
+                                     this->mKernel = this->getCode<func_t>();
+                                   }};
+    Xbyak::util::StackFrame st(this, 1, 4, 16 * 10 + SF_TmpSize);
+    const Xbyak::Reg64& parambase = st.p[0];
+    const Xbyak::Reg64& reg_src = st.t[0];
+    const Xbyak::Reg64& reg_scale = st.t[1];
+    const Xbyak::Reg64& reg_rowsize = st.t[2];
+    const Xbyak::Reg64& reg_iterrow = st.t[3];
+    const Xbyak::Zmm& vreg_scale = zmm31;
+    const auto& mask = k1;
+    const auto masked_off = n_off % 16;
+    if (masked_off != 0) {
+      mov(reg_src, ((1ULL << (16 - masked_off)) - 1) << masked_off);
+      kmovw(mask, reg_src.cvt32());
+    }
+
+    vreg_push(rsp);
+    mov(reg_rowsize.cvt32(), ptr[parambase + OFFSET(row)]);
+    mov(reg_src, qword[parambase + OFFSET(srcptr)]);
+    mov(reg_scale, qword[parambase + OFFSET(scaleptr)]);
+
+    std::vector<Xbyak::Zmm> vreg_src(4 * n_tile / 16);
+    const int ZIDX_TranSrc = 0;
+    for (int i = 0; i < 4 * n_tile / 16; i++) vreg_src[i] = Xbyak::Zmm(ZIDX_TranSrc + i);
+
+    xor_(reg_iterrow, reg_iterrow);
+    Xbyak::Label rowloop;
+    L(rowloop);
+    {
+      assert(("only implement for pack2 bf16", row_pack == 2));
+      for (int i = 0; i < unroll * row_pack; i += row_pack) {
+        vpbroadcastd(vreg_scale, dword[reg_scale + reg_iterrow * sizeof(utils::fp16) + i * sizeof(utils::fp16)]);
+
+        if (masked_off != 0) {
+          int j = utils::padto_le(n_off, 16);
+
+          const auto& vreg0 = vreg_src[j / 16 * 4 + 0];
+          const auto& vreg1 = vreg_src[j / 16 * 4 + 1];
+          const auto& vreg2 = vreg_src[j / 16 * 4 + 2];
+          const auto& vreg3 = vreg_src[j / 16 * 4 + 3];
+          vpmovzxwd(vreg0, yword[reg_src + (i * n_tile + j * row_pack) * sizeof(utils::bf16) + 0]);
+          vpmovzxwd(vreg1, yword[reg_src + (i * n_tile + j * row_pack) * sizeof(utils::bf16) + 32]);
+          vpslldq(vreg0, vreg0, 2);
+          vpslldq(vreg1, vreg1, 2);
+          vcvtps2phx(Ymm(vreg0.getIdx()), vreg0);
+          vcvtps2phx(Ymm(vreg1.getIdx()), vreg1);
+          // #UD If (dest_reg == src1_reg) or (dest_reg == src2_reg)
+          vfmulcph(Ymm(vreg2.getIdx()), Ymm(vreg0.getIdx()), Ymm(vreg_scale.getIdx()));
+          vfmulcph(Ymm(vreg3.getIdx()), Ymm(vreg1.getIdx()), Ymm(vreg_scale.getIdx()));
+          vcvtph2psx(vreg0, Ymm(vreg2.getIdx()));
+          vcvtph2psx(vreg1, Ymm(vreg3.getIdx()));
+          vcvtne2ps2bf16(vreg0, vreg1, vreg0);
+          vmovups(zword[reg_src + (i * n_tile + j * row_pack) * sizeof(utils::bf16)] | mask, vreg0);
+        }
+
+        for (int j = utils::padto(n_off, 16); j < n_tile; j += 16) {
+          const auto& vreg0 = vreg_src[j / 16 * 4 + 0];
+          const auto& vreg1 = vreg_src[j / 16 * 4 + 1];
+          const auto& vreg2 = vreg_src[j / 16 * 4 + 2];
+          const auto& vreg3 = vreg_src[j / 16 * 4 + 3];
+          vpmovzxwd(vreg0, yword[reg_src + (i * n_tile + j * row_pack) * sizeof(utils::bf16) + 0]);
+          vpmovzxwd(vreg1, yword[reg_src + (i * n_tile + j * row_pack) * sizeof(utils::bf16) + 32]);
+          vpslldq(vreg0, vreg0, 2);
+          vpslldq(vreg1, vreg1, 2);
+          vcvtps2phx(Ymm(vreg0.getIdx()), vreg0);
+          vcvtps2phx(Ymm(vreg1.getIdx()), vreg1);
+          // #UD If (dest_reg == src1_reg) or (dest_reg == src2_reg)
+          vfmulcph(Ymm(vreg2.getIdx()), Ymm(vreg0.getIdx()), Ymm(vreg_scale.getIdx()));
+          vfmulcph(Ymm(vreg3.getIdx()), Ymm(vreg1.getIdx()), Ymm(vreg_scale.getIdx()));
+          vcvtph2psx(vreg0, Ymm(vreg2.getIdx()));
+          vcvtph2psx(vreg1, Ymm(vreg3.getIdx()));
+          vcvtne2ps2bf16(vreg0, vreg1, vreg0);
+          vmovups(zword[reg_src + (i * n_tile + j * row_pack) * sizeof(utils::bf16)], vreg0);
+        }
+      }
+    }
+    lea(reg_iterrow, ptr[reg_iterrow + unroll * row_pack]);
+    lea(reg_src, ptr[reg_src + unroll * row_pack * n_tile * sizeof(utils::bf16)]);
+    cmp(reg_iterrow, reg_rowsize);
+    jb(rowloop);
+
+    vreg_pop(rsp);
+  }
+
+  func_t mKernel = nullptr;
+
+ public:
+  template <int NTile, int RowPack = 2>
+  static void forward(utils::bf16* src, const utils::fp16* scale, int row, int col, int src_step, int n_offset) {
+    static_assert(RowPack == 2, "Only implement rowpack2 bf16");
+    static_assert(NTile % 16 == 0, "NTile must be a multiple of 16");
+    constexpr auto unroll = 2;
+    assert(("row should be paded", row % (RowPack * unroll) == 0));
+    assert(("cow should be paded", col % NTile == 0));
+    assert(("can not skip more than col", n_offset < col));
+    int j = utils::padto_le(n_offset, NTile);
+    if (n_offset % NTile != 0) {
+      static const CScaleInterleavedBF16FP16 kern_off(NTile, n_offset % NTile, RowPack, unroll);
+      params param = {src + j * src_step, scale, row};
+      kern_off(&param);
+      j += NTile;
+    }
+
+    for (; j < col; j += NTile) {
+      static const CScaleInterleavedBF16FP16 kern(NTile, 0, RowPack, unroll);
+      params param = {src + j * src_step, scale, row};
+      kern(&param);
+    }
+  }
+
+  template <int NTile, int RowPack = 2>
+  static void reference(utils::bf16* src, const utils::fp16* scale, int row, int col, int src_step, int n_offset) {
+    static_assert(RowPack == 2, "Only implement rowpack2 bf16");
+    static_assert(NTile % 16 == 0, "NTile must be a multiple of 16");
+    assert(("row should be paded", row % RowPack == 0));
+    assert(("cow should be paded", col % NTile == 0));
+    assert(("can not skip more than col", n_offset < col));
+    for (int j = 0; j < col; j += NTile) {
+      for (int i = 0; i < row; i += RowPack) {
+        for (int jj = 0; jj < NTile; ++jj) {
+          if (j + jj < n_offset) continue;
+          auto& rel = (src + j * src_step)[i * NTile + jj * RowPack + 0];
+          auto& img = (src + j * src_step)[i * NTile + jj * RowPack + 1];
+          const auto rel_f32 = static_cast<float>(rel);
+          const auto img_f32 = static_cast<float>(img);
+          const auto rel_scale = static_cast<float>(scale[i + 0]);
+          const auto img_scale = static_cast<float>(scale[i + 1]);
+          rel = static_cast<utils::bf16>(rel_f32 * rel_scale - img_f32 * img_scale);
+          img = static_cast<utils::bf16>(rel_f32 * img_scale + img_f32 * rel_scale);
+        }
+      }
+    }
+  }
+};
+
 }  // namespace jit
 }  // namespace kernel
 }  // namespace jblas
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_ref.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_ref.h
index 375e53b221293..cbce84985b94e 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_ref.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_ref.h
@@ -965,12 +965,13 @@ static inline JBLAS_CODE quant_s8_row_reduce_sum(const int8_t* srcptr, int ldsrc
 }
 
 template <typename _RT>
-static inline JBLAS_CODE row_reduce_sum(const _RT* srcptr, int ldsrc, int row, int col, _RT* reduce) {
-  std::memset(reduce, 0, sizeof(reduce[0]) * col);
-  for (int i = 0; i < row; i++) {
-    for (int j = 0; j < col; j++) {
-      reduce[j] += srcptr[i * ldsrc + j];
+static inline JBLAS_CODE row_reduce_sum(const float* srcptr, int ldsrc, int row, int col, _RT* reduce) {
+  for (int j = 0; j < col; j++) {
+    float tmp = 0.f;
+    for (int i = 0; i < row; i++) {
+      tmp += srcptr[i * ldsrc + j];
     }
+    reduce[j] = _RT(tmp);
   }
   return JblasSuccess;
 }
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_wrapper.h b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_wrapper.h
index 96e75b219a2d5..8275a6312a248 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_wrapper.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/jblas/kernel_wrapper.h
@@ -78,7 +78,6 @@ class Memcpy2D {
   static JBLAS_CODE forward(const _SRC_T* srcptr, _DST_T* dstptr, int row, int col, int srcstep, int dststep,
                             void* const_elt_v = nullptr, Eltops... ops) {
     auto ret = JblasNotSupport;
-#if CompileAVX512F()
     if constexpr (utils::isa_base<ISA_T>::avx512f) {
       ret = kernel::jit::JitMemcpy2DAvx512f::forward<_SRC_T, _DST_T>(srcptr, dstptr, row, col, srcstep, dststep,
                                                                      const_elt_v, ops...);
@@ -86,7 +85,6 @@ class Memcpy2D {
         return ret;
       }
     }
-#endif
 #if CompileAVX2()
     if constexpr (utils::isa_base<ISA_T>::avx2) {
       ret = kernel::jit::JitMemcpy2DAvx2::forward<_SRC_T, _DST_T>(srcptr, dstptr, row, col, srcstep, dststep,
@@ -621,7 +619,7 @@ template <typename _RT>
 class RowReduceSum {
  public:
   template <JBLAS_ISA ISA_T>
-  static inline JBLAS_CODE forward(const _RT* srcptr, int ldsrc, int row, int col, _RT* reduce) {
+  static inline JBLAS_CODE forward(const float* srcptr, int ldsrc, int row, int col, _RT* reduce) {
     return ref::row_reduce_sum<_RT>(srcptr, ldsrc, row, col, reduce);
   }
 };
@@ -644,7 +642,7 @@ class ColBlockReduceSum {
 class RemoveZeroPointBias {
  public:
   template <JBLAS_ISA ISA_T>
-  static inline JBLAS_CODE forward(float* accptr, int ldacc, int row, int col, int8_t* zps, float* scales, int lds,
+  static inline JBLAS_CODE forward_wei(float* accptr, int ldacc, int row, int col, int8_t* zps, float* scales, int lds,
                                    const float* reduce) {
 #if CompileAVX512F()
     if constexpr (utils::isa_base<ISA_T>::avx512f) {
@@ -659,7 +657,7 @@ class RemoveZeroPointBias {
     return ref::remove_wei_zeropoint_bias(accptr, ldacc, row, col, zps, scales, lds, reduce);
   }
   template <JBLAS_ISA ISA_T>
-  static inline JBLAS_CODE forward(float* accptr, int ldacc, int row, int col, uint8_t* zps, float* scales, int lds,
+  static inline JBLAS_CODE forward_act(float* accptr, int ldacc, int row, int col, uint8_t* zps, float* scales, int lds,
                                    const float* reduce) {
 #if CompileAVX512F()
     if constexpr (utils::isa_base<ISA_T>::avx512f) {
@@ -674,7 +672,7 @@ class RemoveZeroPointBias {
     return ref::remove_act_zeropoint_bias(accptr, ldacc, row, col, zps, scales, lds, reduce);
   }
   template <JBLAS_ISA ISA_T>
-  static inline JBLAS_CODE forward(float* accptr, int ldacc, int row, int col, uint8_t* zpa, int8_t* zpb, float* scalea,
+  static inline JBLAS_CODE forward_both(float* accptr, int ldacc, int row, int col, uint8_t* zpa, int8_t* zpb, float* scalea,
                                    float* scaleb, int lds, int k, const float* reducea, const float* reduceb) {
 #if CompileAVX512F()
     if constexpr (utils::isa_base<ISA_T>::avx512f) {
diff --git a/onnxruntime/core/mlas/lib/x86_64/jblas/mlas_jblas_defs.h b/onnxruntime/core/mlas/lib/x86_64/jblas/mlas_jblas_defs.h
index ee5169e5551d7..b0d7d1f61e618 100644
--- a/onnxruntime/core/mlas/lib/x86_64/jblas/mlas_jblas_defs.h
+++ b/onnxruntime/core/mlas/lib/x86_64/jblas/mlas_jblas_defs.h
@@ -29,10 +29,10 @@ using JBLAS_INT8_S4_F32F32 = jblas::wrapper::gemm::LauncherKBlock<
     jblas::prologue_b::gemm::WeightKBlockS4, jblas::epilogue::gemm::CompInt8BlockEpilogue,
     jblas::epilogue::gemm::AccumulatorWriteBackFp32>;
 
-using Jblas_Fp32_AVX512F_F32F32 = JBLAS_FP32_S4_F32F32<jblas::gemm::GemmCore_Row_NN_8x48_AVX512F>;
-using Jblas_Fp32_AVX2_F32F32 = JBLAS_FP32_S4_F32F32<jblas::gemm::GemmCore_Row_NN_2x48_AVX2>;
-using Jblas_Int8_AVX512F_F32F32 = JBLAS_INT8_S4_F32F32<jblas::gemm::GemmCore_Row_NN_8x48_AVX512_VNNI>;
-using Jblas_Int8_AVX2_F32F32 = JBLAS_INT8_S4_F32F32<jblas::gemm::GemmCore_Row_NN_2x48_AVX_VNNI>;
+using Jblas_Fp32_AVX512F_F32F32 = JBLAS_FP32_S4_F32F32<jblas::gemm::SCoreRowNAvx512f<48, 8>>;
+using Jblas_Fp32_AVX2_F32F32 = JBLAS_FP32_S4_F32F32<jblas::gemm::SCoreRowNAvx2<24, 4>>;
+using Jblas_Int8_AVX512F_F32F32 = JBLAS_INT8_S4_F32F32<jblas::gemm::ICoreRowNAvx512vnni<48, 8>>;
+using Jblas_Int8_AVX2_F32F32 = JBLAS_INT8_S4_F32F32<jblas::gemm::ICoreRowNAvxvnni<24, 4>>;
 static Jblas_Fp32_AVX512F_F32F32 JblasAvx512fS4Fp32Fp32;
 static Jblas_Fp32_AVX2_F32F32 JblasAvx2S4Fp32Fp32;
 static Jblas_Int8_AVX512F_F32F32 JblasAvx512VnniS4Fp32Fp32;
diff --git a/onnxruntime/test/contrib_ops/matmul_nbits.cc b/onnxruntime/test/contrib_ops/matmul_nbits.cc
index 5ea798fc9d2c2..9157a45d2bfaa 100644
--- a/onnxruntime/test/contrib_ops/matmul_nbits.cc
+++ b/onnxruntime/test/contrib_ops/matmul_nbits.cc
@@ -141,21 +141,21 @@ void MlasJblasQ4Test(int64_t M, int64_t N, int64_t K, int block_size, bool is_as
       .RunWithConfig();
 }
 
-TEST(MatMulNBits, MlasJblasQ4Fp32G128Asym) {
-  MlasJblasQ4Test(2, 4096, 4096, 128, true, CompFp32);
-}
-
-TEST(MatMulNBits, MlasJblasQ4Fp32G32Asym) {
-  MlasJblasQ4Test(2, 4096, 4096, 32, true, CompFp32);
-}
-
-TEST(MatMulNBits, MlasJblasQ4Fp32G128Sym) {
-  MlasJblasQ4Test(2, 4096, 4096, 128, false, CompFp32);
-}
-
-TEST(MatMulNBits, MlasJblasQ4Fp32G32Sym) {
-  MlasJblasQ4Test(2, 4096, 4096, 32, false, CompFp32);
-}
+//TEST(MatMulNBits, MlasJblasQ4Fp32G128Asym) {
+//  MlasJblasQ4Test(2, 4096, 4096, 128, true, CompFp32);
+//}
+//
+//TEST(MatMulNBits, MlasJblasQ4Fp32G32Asym) {
+//  MlasJblasQ4Test(2, 4096, 4096, 32, true, CompFp32);
+//}
+//
+//TEST(MatMulNBits, MlasJblasQ4Fp32G128Sym) {
+//  MlasJblasQ4Test(2, 4096, 4096, 128, false, CompFp32);
+//}
+//
+//TEST(MatMulNBits, MlasJblasQ4Fp32G32Sym) {
+//  MlasJblasQ4Test(2, 4096, 4096, 32, false, CompFp32);
+//}
 
 TEST(MatMulNBits, MlasJblasQ4Int8G128Sym) {
   MlasJblasQ4Test(2, 4096, 4096, 128, false, CompInt8);