Merge remote-tracking branch 'dev/support-amx-2'

.clang-format

+---
+BasedOnStyle: LLVM
+ColumnLimit: 120  # 设置最大行宽为 100
+IndentWidth: 2
+---

csrc/custom_marlin/gptq_marlin/gptq_marlin.cu

@@ -1420,6 +1420,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
             int* locks      // extra global storage for barrier synchronization
         ) {
         int prob_m = *prob_m_ptr;
+        prob_m = min(prob_m, 1024);
         const int thread_m_blocks = min(div_ceil(prob_m, 16), template_thread_m_blocks);
         if(prob_m > 16 * thread_m_blocks)
             prob_m = (16 * thread_m_blocks) * div_ceil(prob_m, (16 * thread_m_blocks));

csrc/ktransformers_ext/CMakeLists.txt

@@ -53,6 +53,21 @@ else ()
     set(CMAKE_GENERATOR_PLATFORM_LWR "")
 endif ()
 
+if(NOT DEFINED _GLIBCXX_USE_CXX11_ABI)
+    find_package(Python3 REQUIRED COMPONENTS Interpreter)
+
+    execute_process(
+        COMMAND ${Python3_EXECUTABLE} -c
+        "import torch; print('1' if torch.compiled_with_cxx11_abi() else '0')"
+        OUTPUT_VARIABLE ABI_FLAG
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+    )
+
+    set(_GLIBCXX_USE_CXX11_ABI ${ABI_FLAG} CACHE STRING "C++11 ABI setting from PyTorch" FORCE)
+endif()
+
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
+
 if (NOT MSVC)
     if (LLAMA_STATIC)
         add_link_options(-static)
@@ -115,6 +130,37 @@ elseif (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LW
         (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND
          CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|i686|AMD64)$"))
     message(STATUS "x86 detected")
+    set(HOST_IS_X86 TRUE)
+    set(HAS_AVX512 TRUE)
+    set(HAS_AMX TRUE)
+    add_compile_definitions(__x86_64__)
+    # check AVX512
+    execute_process(
+        COMMAND lscpu
+        OUTPUT_VARIABLE LSCPU_OUTPUT
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+    )
+    # message(STATUS "LSCPU_OUTPUT: ${LSCPU_OUTPUT}")
+
+    string(FIND "${LSCPU_OUTPUT}" "avx512" COMPILER_SUPPORTS_AVX512F)
+    
+    if (COMPILER_SUPPORTS_AVX512F GREATER -1)
+        message(STATUS "Compiler and CPU support AVX512F (tested by compiling a program)")
+        add_compile_definitions(__HAS_AVX512F__)
+    else()
+        message(STATUS "Compiler and/or CPU do NOT support AVX512F")
+        set(HAS_AVX512 False)
+    endif()
+
+    # check AMX
+    string(FIND "${LSCPU_OUTPUT}" "amx" COMPILER_SUPPORTS_AMX)
+    
+    if(COMPILER_SUPPORTS_AMX GREATER -1)
+        message(STATUS "Compiler supports AMX")
+        add_compile_definitions(HAS_AMX)
+    else()
+        message(STATUS "Compiler does NOT support AMX")
+    endif()
     if (MSVC)
         # instruction set detection for MSVC only
         if (LLAMA_NATIVE)
@@ -294,8 +340,12 @@ aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/llamafile SOURCE_DIR3
 aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile SOURCE_DIR4)
 aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/kvcache SOURCE_DIR5)
 
+if (HOST_IS_X86 AND HAS_AVX512 AND HAS_AMX)
+    aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/amx SOURCE_DIR6)
+endif()
+
 
-set(ALL_SOURCES ${SOURCE_DIR1} ${SOURCE_DIR2} ${SOURCE_DIR3} ${SOURCE_DIR4} ${SOURCE_DIR5})
+set(ALL_SOURCES ${SOURCE_DIR1} ${SOURCE_DIR2} ${SOURCE_DIR3} ${SOURCE_DIR4} ${SOURCE_DIR5} ${SOURCE_DIR6})
 
 file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h")
 

csrc/ktransformers_ext/bench/bench_moe_amx.py

+#!/usr/bin/env python
+# coding=utf-8
+'''
+Description  :  
+Author       : chenht2022
+Date         : 2025-04-25 18:28:12
+Version      : 1.0.0
+LastEditors  : chenht2022 
+LastEditTime : 2025-04-25 18:28:12
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved. 
+'''
+import os, sys
+import time
+sys.path.append(os.path.dirname(__file__) + '/../build')
+import cpuinfer_ext
+import torch
+
+expert_num = 8
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+n_routed_experts = 8
+layer_num = 10
+qlen = 1024
+CPUInfer = cpuinfer_ext.CPUInfer(65)
+warm_up_iter = 100
+test_iter = 100
+
+def bench_moe(quant_mode: str):
+    with torch.inference_mode(mode=True):
+        if quant_mode == "bf16":
+            bytes_per_elem = 2.000000
+        elif quant_mode == "int8":
+            bytes_per_elem = 1.000000
+        else:
+            assert(False)
+
+
+        moes = []
+        gate_projs = []
+        up_projs = []
+        down_projs = []
+        for _ in range(layer_num):
+            gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
+            up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
+            down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
+            config = cpuinfer_ext.moe.AMX_MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr())
+            if quant_mode == "bf16":
+                moe = cpuinfer_ext.moe.AMXBF16_MOE(config)
+                CPUInfer.submit(moe.load_weights())
+                CPUInfer.sync()
+            elif quant_mode == "int8":
+                moe = cpuinfer_ext.moe.AMXInt8_MOE(config)
+                CPUInfer.submit(moe.load_weights())
+                CPUInfer.sync()
+            gate_projs.append(gate_proj)
+            up_projs.append(up_proj)
+            down_projs.append(down_proj)
+            moes.append(moe)
+        expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous()
+        weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
+        input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
+        output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
+        qlen_tensor = torch.tensor([qlen], dtype=torch.int32)
+
+        # warm up
+        for i in range(warm_up_iter):
+            CPUInfer.submit(
+                moes[i % layer_num].forward( 
+                    qlen, 
+                    n_routed_experts, 
+                    expert_ids[i % layer_num].data_ptr(), 
+                    weights[i % layer_num].data_ptr(),
+                    input[i % layer_num].data_ptr(), 
+                    output[i % layer_num].data_ptr(),
+                    qlen_tensor.data_ptr()
+                )
+            )
+            CPUInfer.sync()
+
+        # test
+        start = time.perf_counter()
+        for i in range(test_iter):
+            CPUInfer.submit(
+                moes[i % layer_num].forward( 
+                    qlen, 
+                    n_routed_experts, 
+                    expert_ids[i % layer_num].data_ptr(), 
+                    weights[i % layer_num].data_ptr(),
+                    input[i % layer_num].data_ptr(), 
+                    output[i % layer_num].data_ptr(),
+                    qlen_tensor.data_ptr()
+                )
+            )
+            CPUInfer.sync()
+        end = time.perf_counter()
+        total_time = end - start
+        print('Quant mode: ', quant_mode)
+        print('Time(s): ', total_time)
+        print('Iteration: ', test_iter) 
+        print('Time(us) per iteration: ', total_time / test_iter * 1000000)
+        print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
+        print('Flops: ', hidden_size * intermediate_size * qlen * 3 * n_routed_experts * 2 * test_iter / total_time / 1000 / 1000 / 1000, 'GFLOPS')
+        print('')
+
+bench_moe("bf16")
+bench_moe("int8")

csrc/ktransformers_ext/operators/llamafile/shared_mem_buffer.cpp → csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.cpp

@@ -30,7 +30,8 @@ void SharedMemBuffer::alloc(void* object, std::vector<std::pair<void**, uint64_t
         if (buffer_) {
             free(buffer_);
         }
-        buffer_ = malloc(size);
+        buffer_ = std::aligned_alloc(64, size);
+
         size_ = size;
         for (auto& obj_requests : hist_requests_) {
             for (auto& requests : obj_requests.second) {
@@ -52,4 +53,4 @@ void SharedMemBuffer::arrange(std::vector<std::pair<void**, uint64_t>> requests)
         *(request.first) = (uint8_t*)buffer_ + offset;
         offset += request.second;
     }
-}
+}
\ No newline at end of file

csrc/ktransformers_ext/cpu_backend/shared_mem_buffer.h

+/**
+ * @Description  :
+ * @Author       : chenht2022
+ * @Date         : 2024-08-05 04:49:08
+ * @Version      : 1.0.0
+ * @LastEditors  : chenht2022 
+ * @LastEditTime : 2024-08-05 06:36:41
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+
+ #ifndef CPUINFER_SHAREDMEMBUFFER_H
+ #define CPUINFER_SHAREDMEMBUFFER_H
+ 
+ #include <cstdint>
+ #include <cstdlib>
+ #include <map>
+ #include <vector>
+ 
+ class SharedMemBuffer {
+    public:
+     SharedMemBuffer();
+     ~SharedMemBuffer();
+ 
+     void alloc(void* object, std::vector<std::pair<void**, uint64_t>> requests);
+     void dealloc(void* object);
+ 
+    private:
+     void* buffer_;
+     uint64_t size_;
+     std::map<void*, std::vector<std::vector<std::pair<void**, uint64_t>>>> hist_requests_;
+ 
+     void arrange(std::vector<std::pair<void**, uint64_t>> requests);
+ };
+ 
+ static SharedMemBuffer shared_mem_buffer;
+ 
+ #endif
\ No newline at end of file

csrc/ktransformers_ext/ext_bindings.cpp

@@ -17,6 +17,11 @@
 #include "operators/llamafile/linear.h"
 #include "operators/llamafile/mlp.h"
 #include "operators/llamafile/moe.h"
+
+#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__)
+#include "operators/amx/moe.hpp"
+#endif
+
 #include "pybind11/functional.h"
 #include "pybind11/operators.h"
 #include "pybind11/pybind11.h"
@@ -563,6 +568,78 @@ class MOEBindings {
     };
 };
 
+
+#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__)
+template<class T>
+class AMX_MOEBindings {
+  public:
+    class WarmUpBindings {
+      public:
+        struct Args {
+            CPUInfer *cpuinfer;
+            AMX_MOE<T> *moe;
+        };
+        static void inner(void *args) {
+            Args *args_ = (Args *)args;
+            args_->cpuinfer->enqueue(&AMX_MOE<T>::warm_up, args_->moe);
+        }
+        static std::pair<intptr_t, intptr_t> cpuinfer_interface(AMX_MOE<T> &moe) {
+            Args *args = new Args{nullptr, &moe};
+            return std::make_pair((intptr_t)&inner, (intptr_t)args);
+        }
+    };
+    class LoadWeightsBindings {
+      public:
+        struct Args {
+            CPUInfer *cpuinfer;
+            AMX_MOE<T> *moe;
+        };
+        static void inner(void *args) {
+            Args *args_ = (Args *)args;
+            args_->cpuinfer->enqueue(&AMX_MOE<T>::load_weights, args_->moe);
+        }
+        static std::pair<intptr_t, intptr_t> cpuinfer_interface(AMX_MOE<T> &moe) {
+            Args *args = new Args{nullptr, &moe};
+            return std::make_pair((intptr_t)&inner, (intptr_t)args);
+        }
+    };
+    class ForwardBindings {
+      public:
+        struct Args {
+            CPUInfer *cpuinfer;
+            AMX_MOE<T> *moe;
+            int qlen;
+            int k;
+            const uint64_t *expert_ids;
+            const float *weights;
+            const void *input;
+            void *output;
+            int *batch_size_tensor;
+        };
+        static void inner(void *args) {
+            Args *args_ = (Args *)args;
+            args_->cpuinfer->enqueue(
+                &AMX_MOE<T>::forward, args_->moe, args_->qlen, args_->k,
+                args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor);
+        }
+        static std::pair<intptr_t, intptr_t>
+        cpuinfer_interface(AMX_MOE<T> &moe, int qlen, int k, intptr_t expert_ids,
+                        intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) {
+            Args *args = new Args{nullptr,
+                                &moe,
+                                qlen,
+                                k,
+                                (const uint64_t *)expert_ids,
+                                (const float *)weights,
+                                (const void *)input,
+                                (void *)output,
+                                (int *)batch_size_tensor};
+            return std::make_pair((intptr_t)&inner, (intptr_t)args);
+        }
+    };
+};
+#endif
+
 PYBIND11_MODULE(cpuinfer_ext, m) {
     py::class_<CPUInfer>(m, "CPUInfer")
         .def(py::init<int>())
@@ -621,6 +698,32 @@ PYBIND11_MODULE(cpuinfer_ext, m) {
         .def("warm_up", &MOEBindings::WarmUpBindinds::cpuinfer_interface)
         .def("forward", &MOEBindings::ForwardBindings::cpuinfer_interface);
 
+
+    #if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__)
+    py::class_<AMX_MOEConfig>(moe_module, "AMX_MOEConfig")
+        .def(py::init([](int expert_num, int routed_expert_num, int hidden_size,
+                         int intermediate_size,
+                         int max_len, intptr_t gate_proj,
+                         intptr_t up_proj, intptr_t down_proj) {
+            return AMX_MOEConfig(expert_num, routed_expert_num, hidden_size,
+                                 intermediate_size, 
+                                 max_len, (void *)gate_proj,
+                                 (void *)up_proj, (void *)down_proj);
+        }));
+
+    py::class_<AMX_MOE<amx::GemmKernel224BF>>(moe_module, "AMXBF16_MOE")
+        .def(py::init<AMX_MOEConfig>())
+        .def("warm_up", &AMX_MOEBindings<amx::GemmKernel224BF>::WarmUpBindings::cpuinfer_interface)
+        .def("load_weights", &AMX_MOEBindings<amx::GemmKernel224BF>::LoadWeightsBindings::cpuinfer_interface)
+        .def("forward", &AMX_MOEBindings<amx::GemmKernel224BF>::ForwardBindings::cpuinfer_interface);
+    py::class_<AMX_MOE<amx::GemmKernel224Int8>>(moe_module, "AMXInt8_MOE")
+        .def(py::init<AMX_MOEConfig>())
+        .def("warm_up", &AMX_MOEBindings<amx::GemmKernel224Int8>::WarmUpBindings::cpuinfer_interface)
+        .def("load_weights", &AMX_MOEBindings<amx::GemmKernel224Int8>::LoadWeightsBindings::cpuinfer_interface)
+        .def("forward", &AMX_MOEBindings<amx::GemmKernel224Int8>::ForwardBindings::cpuinfer_interface);
+
+    #endif
+
     auto kvcache_module = m.def_submodule("kvcache");
 
     py::enum_<AnchorType>(kvcache_module, "AnchorType")

csrc/ktransformers_ext/operators/amx/la/amx.hpp

+/**
+ * @Description  :
+ * @Author       : chenht2022
+ * @Date         : 2025-04-25 18:28:12
+ * @Version      : 1.0.0
+ * @LastEditors  : chenht2022
+ * @LastEditTime : 2025-04-25 18:28:12
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+
+#pragma once
+#include <array>
+#include <cassert>
+#include <cstdint>
+#include <cstdio>
+#include <immintrin.h>
+#include <iostream>
+#include <random>
+#include <stdexcept>
+#include <stdlib.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+#include "utils.hpp"
+#include <memory>
+
+#if (defined(_WIN32) || defined(_WIN64))
+#define RESTRICT __restrict
+#else
+#define RESTRICT __restrict__
+#endif
+
+#if (defined(_WIN32) || defined(_WIN64))
+#define ALWAYS_INLINE __forceinline
+#elif __has_attribute(always_inline) || defined(__GNUC__)
+#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
+#else
+#define ALWAYS_INLINE inline
+#endif
+
+namespace amx {
+
+#define ARCH_GET_XCOMP_PERM 0x1022
+#define ARCH_REQ_XCOMP_PERM 0x1023
+#define XFEATURE_XTILECFG 17
+#define XFEATURE_XTILEDATA 18
+
+const int TMMCount = 8;
+const int MaxTileHeight = 16;
+const int MaxTileWidth = 64;
+
+const int AMX_BLK_SIZE = 32;
+
+#define TMM0 0
+#define TMM1 1
+#define TMM2 2
+#define TMM3 3
+#define TMM4 4
+#define TMM5 5
+#define TMM6 6
+#define TMM7 7
+
+inline bool enable_amx() {
+  static thread_local bool initialized = false;
+  if (initialized) {
+    return true;
+  }
+  initialized = true;
+
+  if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) {
+    printf("\n Fail to do XFEATURE_XTILEDATA \n\n");
+    return false;
+  } else {
+    // printf("\n TILE DATA USE SET - OK \n\n");
+    return true;
+  }
+  return true;
+}
+
+struct alignas(64) TileConfig {
+  uint8_t palette;
+  uint8_t start_row;
+  std::array<uint8_t, 14> __0 = {};
+  std::array<uint16_t, 8> colsb;
+  std::array<uint8_t, 16> __1 = {};
+  std::array<uint8_t, 8> rows;
+  std::array<uint8_t, 8> __2 = {};
+
+  TileConfig() {
+    palette = 1;
+    start_row = 0;
+    for (int i = 0; i < 8; i++) {
+      set_row_col(i, 0, 0);
+    }
+  }
+
+  void set_row_col(int i, uint8_t row, uint16_t col) {
+    colsb[i] = col;
+    rows[i] = row;
+  }
+
+  void set_config() { _tile_loadconfig(this); }
+
+  static void load_data(int to, void *from, size_t stride) {
+    switch (to) {
+    case 0:
+      _tile_loadd(0, from, stride);
+      break;
+    case 1:
+      _tile_loadd(1, from, stride);
+      break;
+    case 2:
+      _tile_loadd(2, from, stride);
+      break;
+    case 3:
+      _tile_loadd(3, from, stride);
+      break;
+    case 4:
+      _tile_loadd(4, from, stride);
+      break;
+    case 5:
+      _tile_loadd(5, from, stride);
+      break;
+    case 6:
+      _tile_loadd(6, from, stride);
+      break;
+    case 7:
+      _tile_loadd(7, from, stride);
+      break;
+    default:
+      throw std::runtime_error("no such tile");
+    }
+  }
+
+  static void store_data(int from, void *to, size_t stride) {
+    switch (from) {
+    case 0:
+      _tile_stored(0, to, stride);
+      break;
+    case 1:
+      _tile_stored(1, to, stride);
+      break;
+    case 2:
+      _tile_stored(2, to, stride);
+      break;
+    case 3:
+      _tile_stored(3, to, stride);
+      break;
+    case 4:
+      _tile_stored(4, to, stride);
+      break;
+    case 5:
+      _tile_stored(5, to, stride);
+      break;
+    case 6:
+      _tile_stored(6, to, stride);
+      break;
+    case 7:
+      _tile_stored(7, to, stride);
+      break;
+    default:
+      throw std::runtime_error("no such tile");
+    }
+  }
+};
+
+static_assert(sizeof(TileConfig) == 64);
+
+inline void debug_tile(int t) {
+  printf("Tile %d\n", t);
+  uint8_t data[16][64] = {};
+  TileConfig::store_data(t, data, 64);
+  for (int i = 0; i < 16; i++) {
+    for (int j = 0; j < 64; j++) {
+      printf("%3d ", data[i][j]);
+    }
+    printf("\n");
+  }
+  printf("\n");
+}
+
+inline void debug_tiles(int to = 8) {
+  for (int i = 0; i < to; i++) {
+    debug_tile(i);
+  }
+}
+
+inline void debug_m512(__m512 x) {
+  float data[16];
+  _mm512_storeu_ps(data, x);
+  for (int i = 0; i < 16; i++) {
+    printf("%f ", data[i]);
+  }
+  printf("\n");
+}
+
+// transpose utils
+inline void transpose_16x16_32bit(__m512i *v) {
+  __m512i v1[16];
+  v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
+  v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
+  v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
+  v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
+  v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
+  v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
+  v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
+  v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
+  v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
+  v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
+  v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
+  v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
+  v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
+  v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
+  v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
+  v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
+
+  v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
+  v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
+  v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
+  v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
+  v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
+  v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
+  v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
+  v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
+  v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
+  v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
+  v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
+  v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
+  v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
+  v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
+  v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
+  v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
+
+  v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
+  v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
+  v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
+  v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
+  v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
+  v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
+  v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
+  v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
+  v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
+  v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
+  v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
+  v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
+  v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
+  v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
+  v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
+  v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
+
+  v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
+  v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
+  v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
+  v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
+  v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
+  v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
+  v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
+  v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
+  v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
+  v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
+  v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
+  v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
+  v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
+  v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
+  v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
+  v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
+}
+
+/*
+  Transpose 16x16 32-bit elements
+  Note that v must be 64 byte aligned
+*/
+inline void transpose_16x16_32bit(__m512i *v, size_t stride) {
+  assert(reinterpret_cast<intptr_t>(v) % 64 == 0 && "v must be 64 aligned");
+
+  auto stride_v = [=](int i) { return offset_pointer(v, i * stride); };
+  __m512i v1[16];
+
+  v1[0] = _mm512_unpacklo_epi32(*stride_v(0), *stride_v(1));
+  v1[1] = _mm512_unpackhi_epi32(*stride_v(0), *stride_v(1));
+  v1[2] = _mm512_unpacklo_epi32(*stride_v(2), *stride_v(3));
+  v1[3] = _mm512_unpackhi_epi32(*stride_v(2), *stride_v(3));
+  v1[4] = _mm512_unpacklo_epi32(*stride_v(4), *stride_v(5));
+  v1[5] = _mm512_unpackhi_epi32(*stride_v(4), *stride_v(5));
+  v1[6] = _mm512_unpacklo_epi32(*stride_v(6), *stride_v(7));
+  v1[7] = _mm512_unpackhi_epi32(*stride_v(6), *stride_v(7));
+  v1[8] = _mm512_unpacklo_epi32(*stride_v(8), *stride_v(9));
+  v1[9] = _mm512_unpackhi_epi32(*stride_v(8), *stride_v(9));
+  v1[10] = _mm512_unpacklo_epi32(*stride_v(10), *stride_v(11));
+  v1[11] = _mm512_unpackhi_epi32(*stride_v(10), *stride_v(11));
+  v1[12] = _mm512_unpacklo_epi32(*stride_v(12), *stride_v(13));
+  v1[13] = _mm512_unpackhi_epi32(*stride_v(12), *stride_v(13));
+  v1[14] = _mm512_unpacklo_epi32(*stride_v(14), *stride_v(15));
+  v1[15] = _mm512_unpackhi_epi32(*stride_v(14), *stride_v(15));
+
+  *stride_v(0) = _mm512_unpacklo_epi64(v1[0], v1[2]);
+  *stride_v(1) = _mm512_unpackhi_epi64(v1[0], v1[2]);
+  *stride_v(2) = _mm512_unpacklo_epi64(v1[1], v1[3]);
+  *stride_v(3) = _mm512_unpackhi_epi64(v1[1], v1[3]);
+  *stride_v(4) = _mm512_unpacklo_epi64(v1[4], v1[6]);
+  *stride_v(5) = _mm512_unpackhi_epi64(v1[4], v1[6]);
+  *stride_v(6) = _mm512_unpacklo_epi64(v1[5], v1[7]);
+  *stride_v(7) = _mm512_unpackhi_epi64(v1[5], v1[7]);
+  *stride_v(8) = _mm512_unpacklo_epi64(v1[8], v1[10]);
+  *stride_v(9) = _mm512_unpackhi_epi64(v1[8], v1[10]);
+  *stride_v(10) = _mm512_unpacklo_epi64(v1[9], v1[11]);
+  *stride_v(11) = _mm512_unpackhi_epi64(v1[9], v1[11]);
+  *stride_v(12) = _mm512_unpacklo_epi64(v1[12], v1[14]);
+  *stride_v(13) = _mm512_unpackhi_epi64(v1[12], v1[14]);
+  *stride_v(14) = _mm512_unpacklo_epi64(v1[13], v1[15]);
+  *stride_v(15) = _mm512_unpackhi_epi64(v1[13], v1[15]);
+
+  v1[0] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0x88);
+  v1[1] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0x88);
+  v1[2] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0x88);
+  v1[3] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0x88);
+  v1[4] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0xdd);
+  v1[5] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0xdd);
+  v1[6] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0xdd);
+  v1[7] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0xdd);
+  v1[8] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0x88);
+  v1[9] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0x88);
+  v1[10] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0x88);
+  v1[11] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0x88);
+  v1[12] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0xdd);
+  v1[13] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0xdd);
+  v1[14] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0xdd);
+  v1[15] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0xdd);
+
+  *stride_v(0) = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
+  *stride_v(1) = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
+  *stride_v(2) = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
+  *stride_v(3) = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
+  *stride_v(4) = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
+  *stride_v(5) = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
+  *stride_v(6) = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
+  *stride_v(7) = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
+  *stride_v(8) = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
+  *stride_v(9) = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
+  *stride_v(10) = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
+  *stride_v(11) = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
+  *stride_v(12) = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
+  *stride_v(13) = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
+  *stride_v(14) = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
+  *stride_v(15) = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
+}
+
+struct GemmKernel224BF {
+  using dt = ggml_bf16_t;
+  using output_t = float;
+  static const int TILE_M = 16;
+  static const int TILE_K = 32;
+  static const int TILE_N = 16;
+  static const int VNNI_BLK = 2;
+
+  static const int M_STEP = TILE_M * 2;
+  static const int N_STEP = TILE_N * 2;
+  static const int K_STEP = TILE_K;
+
+  static inline const int N_BLOCK = 256;
+  static inline const int K_BLOCK = 1792;
+
+  static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; }
+
+  static std::pair<int, int> split_range_n(int n, int ith, int nth) {
+    int n_start = N_BLOCK * ith;
+    int n_end = std::min(n, N_BLOCK * (ith + 1));
+    return {n_start, n_end};
+  }
+
+  static void config() {
+    enable_amx();
+    TileConfig tile_config;
+
+    // size is 16 x 32
+    for (int i = 0; i < 2; i++)
+      tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt));
+
+    // size is 16 x 32
+    for (int i = 2; i < 4; i++)
+      tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt));
+
+    // size is 16 x 16
+    for (int i = 4; i < 8; i++)
+      tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t));
+
+    tile_config.set_config();
+  }
+
+  static void load_a(dt *a, size_t lda) {
+    _tile_loadd(0, a, lda);
+    _tile_loadd(1, offset_pointer(a, lda * TILE_M), lda);
+  }
+
+  static void load_b(dt *b, size_t ldb) {
+    _tile_loadd(2, b, ldb);
+    _tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb);
+  }
+
+  static void clean_c() {
+    _tile_zero(4);
+    _tile_zero(5);
+    _tile_zero(6);
+    _tile_zero(7);
+  }
+
+  static void load_c(output_t *c, size_t ldc) {
+    _tile_loadd(4, c, ldc);
+    _tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
+    _tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc);
+    _tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
+  }
+
+  static void store_c(output_t *c, size_t ldc) {
+    _tile_stored(4, c, ldc);
+    _tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
+    _tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc);
+    _tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
+  }
+
+  static void run_tile() {
+    _tile_dpbf16ps(4, 0, 2);
+    _tile_dpbf16ps(5, 0, 3);
+    _tile_dpbf16ps(6, 1, 2);
+    _tile_dpbf16ps(7, 1, 3);
+  }
+
+  struct BufferA {
+    ggml_bf16_t *a;
+    int max_m, k;
+
+    static size_t required_size(int max_m, int k) { return max_m * k * sizeof(ggml_bf16_t); }
+
+    BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) {
+      assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+      assert(max_m % M_STEP == 0);
+      assert(k % K_STEP == 0);
+      a = reinterpret_cast<ggml_bf16_t *>(ptr);
+    }
+
+    void from_mat(int m, ggml_bf16_t *src, int ith, int nth) {
+      assert(m <= max_m);
+      assert(ith == 0 && nth == 1);
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
+        for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
+          int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+          for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
+            for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
+              __m512i *s = (__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin);
+              __m512i *d = (__m512i *)(a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP +
+                                       i * K_STEP);
+              avx512_copy_32xbf16(s, d);
+            }
+          }
+        }
+      }
+    }
+
+    ggml_bf16_t *get_submat(int m, int k, int m_begin, int k_begin) {
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
+      k_begin -= k_block_begin;
+      int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+      return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP;
+    }
+  };
+
+  struct BufferB {
+    ggml_bf16_t *b;
+    int n, k;
+
+    static size_t required_size(int n, int k) { return n * k * sizeof(ggml_bf16_t); }
+
+    BufferB(int n, int k, void *ptr) : n(n), k(k) {
+      assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+      assert(n % N_STEP == 0);
+      assert(k % K_STEP == 0);
+      b = reinterpret_cast<ggml_bf16_t *>(ptr);
+    }
+
+    void from_mat(ggml_bf16_t *src, int ith, int nth) {
+      auto [n_start, n_end] = split_range_n(n, ith, nth);
+      int n_block_begin = n_start;
+      int n_block_size = n_end - n_block_begin;
+      for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
+        for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
+          int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+          for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
+            for (int i = 0; i < N_STEP; i++) {
+              __m512i *s = (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin);
+              __m512i *d = (__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size +
+                                       k_begin * N_STEP + i * K_STEP);
+              avx512_copy_32xbf16(s, d);
+            }
+            transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
+                                              n_begin * k_block_size + k_begin * N_STEP));
+            transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
+                                              n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP));
+          }
+        }
+      }
+    }
+
+    ggml_bf16_t *get_submat(int n, int k, int n_begin, int k_begin) {
+      int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
+      n_begin -= n_block_begin;
+      int n_block_size = std::min(N_BLOCK, n - n_block_begin);
+      int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
+      k_begin -= k_block_begin;
+      int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+      return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP;
+    }
+  };
+
+  struct BufferC {
+    float *c;
+    int max_m, n;
+
+    static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); }
+
+    BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) {
+      assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+      assert(max_m % M_STEP == 0);
+      assert(n % N_STEP == 0);
+      c = reinterpret_cast<float *>(ptr);
+    }
+
+    void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) {
+      assert(m <= max_m);
+      auto [n_start, n_end] = split_range_n(n, ith, nth);
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      int n_block_begin = n_start;
+      int n_block_size = n_end - n_block_begin;
+      for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
+        for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
+          for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
+            __m512 *x0 =
+                (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP);
+            __m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP +
+                                    i * N_STEP + 16);
+            avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin));
+          }
+        }
+      }
+    }
+
+    float *get_submat(int m, int n, int m_begin, int n_begin) {
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
+      int n_block_size = std::min(N_BLOCK, n - n_block_begin);
+      n_begin -= n_block_begin;
+      return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP;
+    }
+  };
+};
+
+struct GemmKernel224Int8 {
+  using dt = int8_t;
+  using output_t = int32_t;
+  static const int TILE_M = 16;
+  static const int TILE_K = 64;
+  static const int TILE_N = 16;
+  static const int VNNI_BLK = 4;
+
+  static const int M_STEP = TILE_M * 2;
+  static const int N_STEP = TILE_N * 2;
+  static const int K_STEP = TILE_K;
+
+  static inline const int N_BLOCK = 256;
+  static inline const int K_BLOCK = 3584;
+
+  static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; }
+
+  static std::pair<int, int> split_range_n(int n, int ith, int nth) {
+    int n_start = N_BLOCK * ith;
+    int n_end = std::min(n, N_BLOCK * (ith + 1));
+    return {n_start, n_end};
+  }
+
+  static void config() {
+    enable_amx();
+    TileConfig tile_config;
+
+    // size is 16 x 64
+    for (int i = 0; i < 2; i++)
+      tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt));
+
+    // size is 16 x 64
+    for (int i = 2; i < 4; i++)
+      tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt));
+
+    // size is 16 x 16
+    for (int i = 4; i < 8; i++)
+      tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t));
+
+    tile_config.set_config();
+  }
+
+  static void load_a(dt *a, size_t lda) {
+    _tile_loadd(0, a, lda);
+    _tile_loadd(1, offset_pointer(a, lda * TILE_M), lda);
+  }
+
+  static void load_b(dt *b, size_t ldb) {
+    _tile_loadd(2, b, ldb);
+    _tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb);
+  }
+
+  static void clean_c() {
+    _tile_zero(4);
+    _tile_zero(5);
+    _tile_zero(6);
+    _tile_zero(7);
+  }
+
+  static void load_c(output_t *c, size_t ldc) {
+    _tile_loadd(4, c, ldc);
+    _tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
+    _tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc);
+    _tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
+  }
+
+  static void store_c(output_t *c, size_t ldc) {
+    _tile_stored(4, c, ldc);
+    _tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
+    _tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc);
+    _tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
+  }
+
+  static void run_tile() {
+    _tile_dpbssd(4, 0, 2);
+    _tile_dpbssd(5, 0, 3);
+    _tile_dpbssd(6, 1, 2);
+    _tile_dpbssd(7, 1, 3);
+  }
+
+  struct BufferA {
+    int8_t *a;
+    float *d;
+    int max_m, k;
+
+    static size_t required_size(int max_m, int k) { return max_m * k * sizeof(int8_t) + max_m * sizeof(float); }
+
+    BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) {
+      assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+      assert(max_m % M_STEP == 0);
+      assert(k % K_STEP == 0);
+      a = reinterpret_cast<int8_t *>(ptr);
+      d = reinterpret_cast<float *>(a + max_m * k);
+    }
+
+    void from_mat(int m, ggml_bf16_t *src, int ith, int nth) {
+      assert(m <= max_m);
+      assert(ith == 0 && nth == 1);
+      for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
+        for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
+          float amax = 0.0f;
+          for (int j = 0; j < k; j += 32) {
+            __m512 f0, f1;
+            avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + j), &f0, &f1);
+            amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0)));
+            amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1)));
+          }
+          d[m_begin + i] = amax / ((1 << 7) - 1);
+        }
+      }
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
+        for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
+          int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+          for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
+            for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
+              __m512 id = _mm512_set1_ps(d[m_begin + i] ? 1.0f / d[m_begin + i] : 0.0f);
+              int8_t *dst = a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP + i * K_STEP;
+              __m512 f0, f1, f2, f3;
+              avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin), &f0, &f1);
+              avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3);
+              __m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id));
+              __m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id));
+              __m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id));
+              __m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id));
+              __m128i s0 = _mm512_cvtsepi32_epi8(i0);
+              __m128i s1 = _mm512_cvtsepi32_epi8(i1);
+              __m128i s2 = _mm512_cvtsepi32_epi8(i2);
+              __m128i s3 = _mm512_cvtsepi32_epi8(i3);
+              _mm_storeu_si128((__m128i *)dst, s0);
+              _mm_storeu_si128((__m128i *)(dst + 16), s1);
+              _mm_storeu_si128((__m128i *)(dst + 32), s2);
+              _mm_storeu_si128((__m128i *)(dst + 48), s3);
+            }
+          }
+        }
+      }
+    }
+
+    int8_t *get_submat(int m, int k, int m_begin, int k_begin) {
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
+      k_begin -= k_block_begin;
+      int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+      return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP;
+    }
+
+    float *get_scale(int m, int m_begin) { return d + m_begin; }
+  };
+
+  struct BufferB {
+    int8_t *b;
+    float *d;
+    int n, k;
+
+    static size_t required_size(int n, int k) { return n * k * sizeof(int8_t) + n * sizeof(float); }
+
+    BufferB(int n, int k, void *ptr) : n(n), k(k) {
+      assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+      assert(n % N_STEP == 0);
+      assert(k % K_STEP == 0);
+      b = reinterpret_cast<int8_t *>(ptr);
+      d = reinterpret_cast<float *>(b + n * k);
+    }
+
+    void from_mat(ggml_bf16_t *src, int ith, int nth) {
+      auto [n_start, n_end] = split_range_n(n, ith, nth);
+      int n_block_begin = n_start;
+      int n_block_size = n_end - n_block_begin;
+      for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
+        for (int i = 0; i < N_STEP; i++) {
+          float amax = 0.0f;
+          for (int j = 0; j < k; j += 32) {
+            __m512 f0, f1;
+            avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + j), &f0, &f1);
+            amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0)));
+            amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1)));
+          }
+          d[n_block_begin + n_begin + i] = amax / ((1 << 7) - 1);
+        }
+      }
+      for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
+        for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
+          int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+          for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
+            for (int i = 0; i < N_STEP; i++) {
+              __m512 id = _mm512_set1_ps(d[n_block_begin + n_begin + i] ? 1.0f / d[n_block_begin + n_begin + i] : 0.0f);
+              int8_t *dst = b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size +
+                            k_begin * N_STEP + i * K_STEP;
+              __m512 f0, f1, f2, f3;
+              avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin),
+                                        &f0, &f1);
+              avx512_32xbf16_to_32xfp32(
+                  (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3);
+              __m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id));
+              __m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id));
+              __m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id));
+              __m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id));
+              __m128i s0 = _mm512_cvtsepi32_epi8(i0);
+              __m128i s1 = _mm512_cvtsepi32_epi8(i1);
+              __m128i s2 = _mm512_cvtsepi32_epi8(i2);
+              __m128i s3 = _mm512_cvtsepi32_epi8(i3);
+              _mm_storeu_si128((__m128i *)dst, s0);
+              _mm_storeu_si128((__m128i *)(dst + 16), s1);
+              _mm_storeu_si128((__m128i *)(dst + 32), s2);
+              _mm_storeu_si128((__m128i *)(dst + 48), s3);
+            }
+            transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
+                                              n_begin * k_block_size + k_begin * N_STEP));
+            transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
+                                              n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP));
+          }
+        }
+      }
+    }
+
+    int8_t *get_submat(int n, int k, int n_begin, int k_begin) {
+      int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
+      n_begin -= n_block_begin;
+      int n_block_size = std::min(N_BLOCK, n - n_block_begin);
+      int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
+      k_begin -= k_block_begin;
+      int k_block_size = std::min(K_BLOCK, k - k_block_begin);
+      return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP;
+    }
+
+    float *get_scale(int n, int n_begin) { return d + n_begin; }
+  };
+
+  struct BufferC {
+    float *c;
+    int max_m, n;
+
+    static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); }
+
+    BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) {
+      assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+      assert(max_m % M_STEP == 0);
+      assert(n % N_STEP == 0);
+      c = reinterpret_cast<float *>(ptr);
+    }
+
+    void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) {
+      assert(m <= max_m);
+      auto [n_start, n_end] = split_range_n(n, ith, nth);
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      int n_block_begin = n_start;
+      int n_block_size = n_end - n_block_begin;
+      for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
+        for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
+          for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
+            __m512 *x0 =
+                (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP);
+            __m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP +
+                                    i * N_STEP + 16);
+            avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin));
+          }
+        }
+      }
+    }
+
+    float *get_submat(int m, int n, int m_begin, int n_begin) {
+      int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
+      int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
+      int n_block_size = std::min(N_BLOCK, n - n_block_begin);
+      n_begin -= n_block_begin;
+      return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP;
+    }
+  };
+};
+
+inline void mat_mul(int m, int n, int k, std::shared_ptr<GemmKernel224BF::BufferA> ba,
+                    std::shared_ptr<GemmKernel224BF::BufferB> bb, std::shared_ptr<GemmKernel224BF::BufferC> bc, int ith,
+                    int nth, bool use_amx) {
+  using K = GemmKernel224BF;
+  assert(n % K::N_STEP == 0);
+  assert(k % K::K_STEP == 0);
+
+  auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+
+  for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) {
+    for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) {
+      for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) {
+
+        float *c = bc->get_submat(m, n, m_begin, n_begin);
+        if (!use_amx) {
+          __m512 *c512 = (__m512 *)c;
+          if (k_block_begin == 0) {
+            for (int m_i = 0; m_i < m; m_i++) {
+              c512[m_i * 2] = _mm512_setzero_ps();
+              c512[m_i * 2 + 1] = _mm512_setzero_ps();
+            }
+          }
+
+          for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
+            int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin);
+            __m512bh *b512 = (__m512bh *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin);
+            for (int m_i = 0; m_i < m; m_i++) {
+              for (int k_i = 0; k_i < 16; k_i++) {
+                __m512bh ma = (__m512bh)_mm512_set1_epi32(a32[m_i * 16 + k_i]);
+                for (int n_i = 0; n_i < 2; n_i++) {
+                  c512[m_i * 2 + n_i] = _mm512_dpbf16_ps(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]);
+                }
+              }
+            }
+          }
+
+        } else {
+          if (k_block_begin == 0) {
+            K::clean_c();
+          } else {
+            K::load_c(c, K::N_STEP * sizeof(float));
+          }
+          for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
+            K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t));
+            K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t));
+            K::run_tile();
+          }
+          K::store_c(c, K::N_STEP * sizeof(float));
+        }
+      }
+    }
+  }
+}
+
+inline __m512i _mm512_dpbssd_epi32(__m512i src, __m512i a, __m512i b) {
+  __m256i a_lo = _mm512_extracti64x4_epi64(a, 0);
+  __m256i a_hi = _mm512_extracti64x4_epi64(a, 1);
+  __m256i b_lo = _mm512_extracti64x4_epi64(b, 0);
+  __m256i b_hi = _mm512_extracti64x4_epi64(b, 1);
+
+  b_lo = _mm256_sign_epi8(b_lo, a_lo);
+  b_hi = _mm256_sign_epi8(b_hi, a_hi);
+
+  b = _mm512_inserti64x4(b, b_lo, 0);
+  b = _mm512_inserti64x4(b, b_hi, 1);
+
+  a = _mm512_abs_epi8(a);
+
+  return _mm512_dpbusd_epi32(src, a, b);
+}
+
+inline void mat_mul(int m, int n, int k, std::shared_ptr<GemmKernel224Int8::BufferA> ba,
+                    std::shared_ptr<GemmKernel224Int8::BufferB> bb, std::shared_ptr<GemmKernel224Int8::BufferC> bc,
+                    int ith, int nth, bool use_amx) {
+  using K = GemmKernel224Int8;
+  assert(n % K::N_STEP == 0);
+  assert(k % K::K_STEP == 0);
+
+  auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+
+  for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) {
+    for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) {
+      for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) {
+        float *c = bc->get_submat(m, n, m_begin, n_begin);
+
+        if (!use_amx) {
+          __m512i *c512 = (__m512i *)c;
+          if (k_block_begin == 0) {
+            for (int m_i = 0; m_i < m; m_i++) {
+              c512[m_i * 2] = _mm512_setzero_si512();
+              c512[m_i * 2 + 1] = _mm512_setzero_si512();
+            }
+          }
+
+          for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
+            static_assert(K::K_STEP * sizeof(int8_t) == sizeof(__m512i));
+            static_assert(K::N_STEP / K::TILE_N == 2, "Must be lke this");
+
+            int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin);
+            __m512i *b512 = (__m512i *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin);
+            for (int m_i = 0; m_i < m; m_i++) {
+              for (int k_i = 0; k_i < 16; k_i++) {
+                __m512i ma = _mm512_set1_epi32(a32[m_i * 16 + k_i]);
+                for (int n_i = 0; n_i < 2; n_i++) {
+                  c512[m_i * 2 + n_i] = _mm512_dpbssd_epi32(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]);
+                }
+              }
+            }
+          }
+        } else {
+          if (k_block_begin == 0) {
+            K::clean_c();
+          } else {
+            K::load_c((int32_t *)c, K::N_STEP * sizeof(int32_t));
+          }
+          for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
+            K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t));
+            K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t));
+            K::run_tile();
+          }
+          K::store_c((int32_t *)c, K::N_STEP * sizeof(int32_t));
+        }
+
+        if (k_block_begin + K::K_BLOCK >= k) {
+          int to = m - m_begin;
+          if (m - m_begin > K::M_STEP) {
+            to = K::M_STEP;
+          }
+          for (int i = 0; i < to; i++) {
+            __m512 as = _mm512_set1_ps(*ba->get_scale(m, m_begin + i));
+            __m512 bs = _mm512_load_ps(bb->get_scale(n, n_begin));
+            __m512i now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP));
+            __m512 result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now));
+            _mm512_store_ps((__m512 *)(c + i * K::N_STEP), result);
+            bs = _mm512_load_ps(bb->get_scale(n, n_begin) + K::TILE_N);
+            now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP + K::TILE_N));
+            result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now));
+            _mm512_store_ps((__m512 *)(c + i * K::N_STEP + K::TILE_N), result);
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace amx
\ No newline at end of file

csrc/ktransformers_ext/operators/amx/la/utils.hpp

+/**
+ * @Description  :
+ * @Author       : chenht2022
+ * @Date         : 2025-04-25 18:28:12
+ * @Version      : 1.0.0
+ * @LastEditors  : chenht2022
+ * @LastEditTime : 2025-04-25 18:28:12
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+
+#pragma once
+#include <cstdint>
+
+
+template <typename T>
+T* offset_pointer(T* ptr, std::size_t byte_offset) {
+  return reinterpret_cast<T*>(reinterpret_cast<char*>(ptr) + byte_offset);
+}
+
+template <typename T>
+const T* offset_pointer(const T* ptr, std::size_t byte_offset) {
+  return reinterpret_cast<const T*>(reinterpret_cast<const char*>(ptr) + byte_offset);
+}
+
+template <typename T>
+T* offset_pointer_row_major(T* t, int row, int col, std::size_t ld) {
+  return offset_pointer(t, row * ld) + col;
+}
+
+template <typename T>
+T* offset_pointer_col_major(T* t, int row, int col, std::size_t ld) {
+  return offset_pointer(t, col * ld) + row;
+}
+
+static inline void avx512_copy_32xbf16(__m512i* src, __m512i* dst) {
+  _mm512_storeu_si512(dst, _mm512_loadu_si512(src));
+}
+
+static inline void avx512_32xfp32_to_32xbf16(__m512* src0, __m512* src1, __m512i* dst) {
+  _mm512_storeu_si512(dst, __m512i(_mm512_cvtne2ps_pbh(*src1, *src0)));
+}
+
+static inline void avx512_32xbf16_to_32xfp32(__m512i* src, __m512* dst0, __m512* dst1) {
+  _mm512_storeu_ps(dst0, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src))), 16)));
+  _mm512_storeu_ps(dst1, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src) + 1)), 16)));
+}
\ No newline at end of file

csrc/ktransformers_ext/operators/amx/moe.hpp

+/**
+ * @Description  :
+ * @Author       : chenht2022
+ * @Date         : 2025-04-25 18:28:12
+ * @Version      : 1.0.0
+ * @LastEditors  : chenht2022
+ * @LastEditTime : 2025-04-25 18:28:12
+ * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+ **/
+#ifndef CPUINFER_OPERATOR_AMX_MOE_H
+#define CPUINFER_OPERATOR_AMX_MOE_H
+
+#include <cmath>
+#include <cstdio>
+#include <functional>
+#include <mutex>
+#include <vector>
+
+#include "../../cpu_backend/backend.h"
+#include "../../cpu_backend/shared_mem_buffer.h"
+#include "llama.cpp/ggml-impl.h"
+#include "llama.cpp/ggml-quants.h"
+#include "llama.cpp/ggml.h"
+#include "llamafile/sgemm.h"
+
+#include "la/amx.hpp"
+
+#ifdef USE_NUMA
+#include <numa.h>
+#include <numaif.h>
+void *numa_alloc_aligned(size_t size, int node, size_t alignment) {
+  void *ptr = numa_alloc_onnode(size, node);
+  assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+  return ptr;
+}
+#endif
+
+static inline __m512 exp_avx512(__m512 x) {
+  const __m512 log2e = _mm512_set1_ps(1.44269504089f);
+  const __m512 c1 = _mm512_set1_ps(0.69314718056f);
+
+  __m512 y = _mm512_mul_ps(x, log2e);
+  __m512i int_part = _mm512_cvtps_epi32(y);
+  __m512 frac_part = _mm512_sub_ps(y, _mm512_cvtepi32_ps(int_part));
+
+  const __m512 poly_1 = _mm512_set1_ps(0.9999999995f);
+  const __m512 poly_2 = _mm512_set1_ps(0.6931471805f);
+  const __m512 poly_3 = _mm512_set1_ps(0.2402265069f);
+  const __m512 poly_4 = _mm512_set1_ps(0.0555041087f);
+  const __m512 poly_5 = _mm512_set1_ps(0.0096181291f);
+  const __m512 poly_6 = _mm512_set1_ps(0.0013333558f);
+
+  __m512 frac_exp = _mm512_fmadd_ps(
+      frac_part, poly_6,
+      _mm512_fmadd_ps(frac_part, poly_5,
+                      _mm512_fmadd_ps(frac_part, poly_4,
+                                      _mm512_fmadd_ps(frac_part, poly_3, _mm512_fmadd_ps(frac_part, poly_2, poly_1)))));
+
+  __m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part));
+  return _mm512_mul_ps(two_pow_i, frac_exp);
+}
+
+static inline __m512 act_fn(__m512 gate_val, __m512 up_val) {
+  __m512 neg_gate_val = _mm512_sub_ps(_mm512_setzero_ps(), gate_val);
+  __m512 exp_neg_gate = exp_avx512(neg_gate_val);
+  __m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg_gate);
+  __m512 act_val = _mm512_div_ps(gate_val, denom);
+
+  return _mm512_mul_ps(act_val, up_val);
+}
+
+struct AMX_MOEConfig {
+  int expert_num;
+  int routed_expert_num;
+  int hidden_size;
+  int intermediate_size;
+  int max_len;
+  void *gate_proj;
+  void *up_proj;
+  void *down_proj;
+
+  AMX_MOEConfig() {}
+
+  AMX_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int max_len,
+                void *gate_proj, void *up_proj, void *down_proj)
+      : expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size),
+        intermediate_size(intermediate_size), max_len(max_len), gate_proj(gate_proj), up_proj(up_proj),
+        down_proj(down_proj) {}
+};
+
+template <class T> class AMX_MOE {
+private:
+  AMX_MOEConfig config_;
+  void *gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)]
+  void *up_proj_;   // [expert_num * intermediate_size * hidden_size ( /32 if quantized)]
+  void *down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)]
+
+  ggml_bf16_t *m_local_input_;       // [routed_expert_num * max_len * hidden_size]
+  ggml_bf16_t *m_local_gate_output_; // [routed_expert_num * max_len * intermediate_size]
+  ggml_bf16_t *m_local_up_output_;   // [routed_expert_num * max_len * intermediate_size]
+  ggml_bf16_t *m_local_down_output_; // [routed_expert_num * max_len * hidden_size]
+
+  std::vector<std::vector<int>> m_local_pos_;          // [max_len, routed_expert_num]
+  std::vector<int> m_local_num_;                       // [expert_num]
+  std::vector<int> m_expert_id_map_;                   // [expert_num]
+  std::vector<ggml_bf16_t *> m_local_input_ptr_;       // [expert_num]
+  std::vector<ggml_bf16_t *> m_local_gate_output_ptr_; // [expert_num]
+  std::vector<ggml_bf16_t *> m_local_up_output_ptr_;   // [expert_num]
+  std::vector<ggml_bf16_t *> m_local_down_output_ptr_; // [expert_num]
+
+  std::vector<std::shared_ptr<typename T::BufferA>> gate_up_ba_;
+  std::vector<std::shared_ptr<typename T::BufferC>> gate_bc_;
+  std::vector<std::shared_ptr<typename T::BufferC>> up_bc_;
+  std::vector<std::shared_ptr<typename T::BufferA>> down_ba_;
+  std::vector<std::shared_ptr<typename T::BufferC>> down_bc_;
+
+#ifdef USE_NUMA
+  std::vector<std::vector<std::shared_ptr<typename T::BufferB>>> gate_bb_numa_;
+  std::vector<std::vector<std::shared_ptr<typename T::BufferB>>> up_bb_numa_;
+  std::vector<std::vector<std::shared_ptr<typename T::BufferB>>> down_bb_numa_;
+#else
+  std::vector<std::shared_ptr<typename T::BufferB>> gate_bb_;
+  std::vector<std::shared_ptr<typename T::BufferB>> up_bb_;
+  std::vector<std::shared_ptr<typename T::BufferB>> down_bb_;
+#endif
+
+public:
+  AMX_MOE(AMX_MOEConfig config) {
+    config_ = config;
+    gate_proj_ = config_.gate_proj;
+    up_proj_ = config_.up_proj;
+    down_proj_ = config_.down_proj;
+
+    std::vector<std::pair<void **, uint64_t>> m_mem_requests;
+    m_mem_requests.push_back({(void **)&m_local_input_,
+                              sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size});
+    m_mem_requests.push_back({(void **)&m_local_gate_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num *
+                                                                  config_.max_len * config_.intermediate_size});
+    m_mem_requests.push_back({(void **)&m_local_up_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num *
+                                                                config_.max_len * config_.intermediate_size});
+    m_mem_requests.push_back({(void **)&m_local_down_output_,
+                              sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size});
+    std::vector<void *> gate_up_ba_ptr(config_.expert_num);
+    std::vector<void *> gate_bc_ptr(config_.expert_num);
+    std::vector<void *> up_bc_ptr(config_.expert_num);
+    std::vector<void *> down_ba_ptr(config_.expert_num);
+    std::vector<void *> down_bc_ptr(config_.expert_num);
+    for (int i = 0; i < config_.expert_num; i++) {
+      m_mem_requests.push_back(
+          {(void **)&gate_up_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)});
+      m_mem_requests.push_back(
+          {(void **)&gate_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)});
+      m_mem_requests.push_back(
+          {(void **)&up_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)});
+      m_mem_requests.push_back(
+          {(void **)&down_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)});
+      m_mem_requests.push_back(
+          {(void **)&down_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)});
+    }
+    shared_mem_buffer.alloc(this, m_mem_requests);
+
+    m_local_pos_.resize(config_.max_len);
+    for (int i = 0; i < config_.max_len; i++) {
+      m_local_pos_[i].resize(config_.routed_expert_num);
+    }
+    m_expert_id_map_.resize(config_.expert_num);
+    m_local_num_.resize(config_.expert_num);
+    m_local_input_ptr_.resize(config_.expert_num);
+    m_local_gate_output_ptr_.resize(config_.expert_num);
+    m_local_up_output_ptr_.resize(config_.expert_num);
+    m_local_down_output_ptr_.resize(config_.expert_num);
+
+    for (uint64_t i = 0; i < config_.expert_num; i++) {
+      gate_up_ba_.push_back(
+          std::make_shared<typename T::BufferA>(config_.max_len, config_.hidden_size, gate_up_ba_ptr[i]));
+      gate_bc_.push_back(
+          std::make_shared<typename T::BufferC>(config_.max_len, config_.intermediate_size, gate_bc_ptr[i]));
+      up_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.intermediate_size, up_bc_ptr[i]));
+      down_ba_.push_back(
+          std::make_shared<typename T::BufferA>(config_.max_len, config_.intermediate_size, down_ba_ptr[i]));
+      down_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.hidden_size, down_bc_ptr[i]));
+
+#ifdef USE_NUMA
+      int numa_nodes = numa_num_configured_nodes();
+      gate_bb_numa_.resize(numa_nodes);
+      up_bb_numa_.resize(numa_nodes);
+      down_bb_numa_.resize(numa_nodes);
+      for (int j = 0; j < numa_nodes; j++) {
+        void *gate_bb_ptr =
+            numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64);
+        gate_bb_numa_[j].push_back(
+            std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, gate_bb_ptr));
+        void *up_bb_ptr =
+            numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64);
+        up_bb_numa_[j].push_back(
+            std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, up_bb_ptr));
+        void *down_bb_ptr =
+            numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64);
+        down_bb_numa_[j].push_back(
+            std::make_shared<typename T::BufferB>(config_.hidden_size, config_.intermediate_size, down_bb_ptr));
+      }
+#else
+      void *gate_bb_ptr =
+          std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size));
+      gate_bb_.push_back(
+          std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, gate_bb_ptr));
+
+      void *up_bb_ptr =
+          std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size));
+      up_bb_.push_back(
+          std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, up_bb_ptr));
+
+      void *down_bb_ptr =
+          std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size));
+      down_bb_.push_back(
+          std::make_shared<typename T::BufferB>(config_.hidden_size, config_.intermediate_size, down_bb_ptr));
+#endif
+    }
+  }
+
+  ~AMX_MOE() { shared_mem_buffer.dealloc(this); }
+
+  void load_weights(Backend *backend) {
+    int nth = T::recommended_nth(config_.intermediate_size);
+    backend->do_work_stealing_job(
+        nth * config_.expert_num, nullptr,
+        [&](int task_id) {
+          uint64_t expert_idx = task_id / nth;
+          int ith = task_id % nth;
+#ifdef USE_NUMA
+          int numa_nodes = numa_num_configured_nodes();
+          for (int j = 0; j < numa_nodes; j++) {
+            gate_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj +
+                                                       expert_idx * config_.intermediate_size * config_.hidden_size,
+                                                   ith, nth);
+            up_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.up_proj +
+                                                     expert_idx * config_.intermediate_size * config_.hidden_size,
+                                                 ith, nth);
+          }
+#else
+          gate_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj +
+                                             expert_idx * config_.intermediate_size * config_.hidden_size,
+                                         ith, nth);
+          up_bb_[expert_idx]->from_mat(
+              (ggml_bf16_t *)config_.up_proj + expert_idx * config_.intermediate_size * config_.hidden_size, ith, nth);
+#endif
+        },
+        nullptr);
+    nth = T::recommended_nth(config_.hidden_size);
+    backend->do_work_stealing_job(
+        nth * config_.expert_num, nullptr,
+        [&](int task_id) {
+          uint64_t expert_idx = task_id / nth;
+          int ith = task_id % nth;
+#ifdef USE_NUMA
+          int numa_nodes = numa_num_configured_nodes();
+          for (int j = 0; j < numa_nodes; j++) {
+            down_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj +
+                                                       expert_idx * config_.hidden_size * config_.intermediate_size,
+                                                   ith, nth);
+          }
+#else
+          down_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj +
+                                             expert_idx * config_.hidden_size * config_.intermediate_size,
+                                         ith, nth);
+#endif
+        },
+        nullptr);
+  }
+
+  void warm_up(Backend *backend) {}
+
+  void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output,
+               int *batch_size_tensor, Backend *backend) {
+    bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num);
+    qlen = batch_size_tensor[0];
+    int activated_expert = 0;
+    for (int i = 0; i < config_.expert_num; i++) {
+      m_local_num_[i] = 0;
+    }
+    for (int i = 0; i < qlen; i++) {
+      for (int j = 0; j < k; j++) {
+        m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++;
+      }
+    }
+    for (int i = 0; i < config_.expert_num; i++) {
+      if (m_local_num_[i] > 0) {
+        m_expert_id_map_[activated_expert] = i;
+        activated_expert++;
+      }
+    }
+    uint64_t offset = 0;
+    for (int i = 0; i < config_.expert_num; i++) {
+      m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size;
+      m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size;
+      m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size;
+      m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size;
+      offset += m_local_num_[i];
+    }
+    backend->do_work_stealing_job(
+        qlen, nullptr,
+        [&](int i) {
+          for (int j = 0; j < k; j++) {
+            memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size,
+                   (ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size);
+          }
+        },
+        nullptr);
+    backend->do_work_stealing_job(
+        activated_expert, nullptr,
+        [&](int task_id) {
+          int expert_idx = m_expert_id_map_[task_id];
+
+          gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1);
+        },
+        nullptr);
+    int nth = T::recommended_nth(config_.intermediate_size);
+    backend->do_work_stealing_job(
+        nth * activated_expert, [&](int _) { T::config(); },
+        [&](int task_id) {
+          int expert_idx = m_expert_id_map_[task_id / nth];
+          int ith = task_id % nth;
+#ifdef USE_NUMA
+          amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
+                       gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx],
+                       ith, nth, use_amx);
+          amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
+                       gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith,
+                       nth, use_amx);
+#else
+          amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
+                       gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx);
+          amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
+                       gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx);
+#endif
+          gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth);
+          up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth);
+          auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth);
+          for (int i = 0; i < m_local_num_[expert_idx]; i++) {
+            ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size];
+            ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size];
+            for (int j = n_start; j < n_end; j += 32) {
+              __m512 gate_val0, gate_val1, up_val0, up_val1;
+              avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1);
+              avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1);
+              __m512 result0 = act_fn(gate_val0, up_val0);
+              __m512 result1 = act_fn(gate_val1, up_val1);
+              avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j));
+            }
+          }
+        },
+        nullptr);
+    backend->do_work_stealing_job(
+        activated_expert, nullptr,
+        [&](int task_id) {
+          int expert_idx = m_expert_id_map_[task_id];
+          down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1);
+        },
+        nullptr);
+    nth = T::recommended_nth(config_.hidden_size);
+    backend->do_work_stealing_job(
+        nth * activated_expert, [&](int _) { T::config(); },
+        [&](int task_id) {
+          int expert_idx = m_expert_id_map_[task_id / nth];
+          int ith = task_id % nth;
+#ifdef USE_NUMA
+          amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx],
+                       down_bb_numa_[Backend::numa_node][expert_idx], down_bc_[expert_idx], ith, nth, use_amx);
+#else
+          amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx],
+                       down_bb_[expert_idx], down_bc_[expert_idx], ith, nth, use_amx);
+#endif
+          down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth);
+        },
+        nullptr);
+    backend->do_work_stealing_job(
+        qlen, nullptr,
+        [&](int i) {
+          for (int e = 0; e < config_.hidden_size; e += 32) {
+            __m512 x0 = _mm512_setzero_ps();
+            __m512 x1 = _mm512_setzero_ps();
+            for (int j = 0; j < k; j++) {
+              __m512 weight = _mm512_set1_ps(weights[i * k + j]);
+              __m512 down_output0, down_output1;
+              avx512_32xbf16_to_32xfp32((__m512i *)(m_local_down_output_ptr_[expert_ids[i * k + j]] +
+                                                    m_local_pos_[i][j] * config_.hidden_size + e),
+                                        &down_output0, &down_output1);
+              x0 = _mm512_fmadd_ps(down_output0, weight, x0);
+              x1 = _mm512_fmadd_ps(down_output1, weight, x1);
+            }
+            avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + i * config_.hidden_size + e));
+          }
+        },
+        nullptr);
+  }
+};
+
+#endif
\ No newline at end of file

csrc/ktransformers_ext/operators/llamafile/linear.h

@@ -17,12 +17,12 @@
 #include <vector>
 
 #include "../../cpu_backend/backend.h"
+#include "../../cpu_backend/shared_mem_buffer.h"
 #include "conversion.h"
 #include "llama.cpp/ggml-impl.h"
 #include "llama.cpp/ggml-quants.h"
 #include "llama.cpp/ggml.h"
 #include "llamafile/sgemm.h"
-#include "shared_mem_buffer.h"
 
 struct LinearConfig {
     int input_size;

csrc/ktransformers_ext/operators/llamafile/mlp.h

@@ -17,12 +17,12 @@
 #include <vector>
 
 #include "../../cpu_backend/backend.h"
+#include "../../cpu_backend/shared_mem_buffer.h"
 #include "conversion.h"
 #include "llama.cpp/ggml-impl.h"
 #include "llama.cpp/ggml-quants.h"
 #include "llama.cpp/ggml.h"
 #include "llamafile/sgemm.h"
-#include "shared_mem_buffer.h"
 
 struct MLPConfig {
     int hidden_size;

csrc/ktransformers_ext/operators/llamafile/moe.h

@@ -17,12 +17,12 @@
 #include <vector>
 
 #include "../../cpu_backend/backend.h"
+#include "../../cpu_backend/shared_mem_buffer.h"
 #include "conversion.h"
 #include "llama.cpp/ggml-impl.h"
 #include "llama.cpp/ggml-quants.h"
 #include "llama.cpp/ggml.h"
 #include "llamafile/sgemm.h"
-#include "shared_mem_buffer.h"
 
 struct MOEConfig {
     int expert_num;

csrc/ktransformers_ext/operators/llamafile/shared_mem_buffer.h

-/**
- * @Description  :
- * @Author       : chenht2022
- * @Date         : 2024-08-05 04:49:08
- * @Version      : 1.0.0
- * @LastEditors  : chenht2022 
- * @LastEditTime : 2024-08-05 06:36:41
- * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
- **/
-
-#ifndef CPUINFER_SHAREDMEMBUFFER_H
-#define CPUINFER_SHAREDMEMBUFFER_H
-
-#include <cstdint>
-#include <cstdlib>
-#include <map>
-#include <vector>
-
-class SharedMemBuffer {
-   public:
-    SharedMemBuffer();
-    ~SharedMemBuffer();
-
-    void alloc(void* object, std::vector<std::pair<void**, uint64_t>> requests);
-    void dealloc(void* object);
-
-   private:
-    void* buffer_;
-    uint64_t size_;
-    std::map<void*, std::vector<std::vector<std::pair<void**, uint64_t>>>> hist_requests_;
-
-    void arrange(std::vector<std::pair<void**, uint64_t>> requests);
-};
-
-static SharedMemBuffer shared_mem_buffer;
-
-#endif
\ No newline at end of file

doc/en/AMX.md

+# Introduction to AMX Instruction Set
+
+Intel Advanced Matrix Extensions (AMX) are a set of specialized instruction extensions introduced for the x86 architecture starting with Sapphire Rapids (4th generation Xeon Scalable processors) and onward. AMX accelerates large-scale matrix computations at the hardware level, particularly for the compute-intensive parts of deep learning inference and machine learning workloads. By introducing the concept of Tile registers, it loads 2D sub-matrices into dedicated Tile registers and performs matrix multiply-accumulate operations at the register level, significantly improving throughput and energy efficiency.
+
+Each CPU core contains 8 dedicated registers (tmm0–tmm7), with each register capable of holding up to 16 rows × 64 bytes of data to store 2D sub-matrices. Additionally, there is a 64-byte configuration register (TILECFG) used to describe each tmm register's number of rows, columns, and row stride.
+
+The main AMX instructions are summarized as follows:
+
+| Instruction Category | Instruction Names | Description |
+|:---|:---|:---|
+| Configuration Instructions | LDTILECFG, STTILECFG, TILERELEASE, TILEZERO | Configure/reset Tile registers and metadata |
+| Load/Store Instructions | TILELOADD, TILELOADDT1, TILESTORED | Transfer data between memory and Tile registers |
+| INT8 Computation Instructions | TDPBSSD, TDPBUSD, TDPBUUD, TDPBSUD | Perform multiply and accumulate operations on int8 sub-matrices within Tiles |
+| BF16 Computation Instructions | TDPBF16PS | Perform multiply and accumulate operations on bfloat16 sub-matrices within Tiles |
+
+To simplify development, Intel provides corresponding intrinsics, allowing C/C++ developers to leverage AMX's performance benefits without writing lengthy assembly code. For example:
+
+```C++
+#include <immintrin.h>
+
+_tile_loadconfig(cfg_ptr);
+_tile_loadd(tmm0, A_ptr, lda);
+_tile_loadd(tmm1, B_ptr, ldb);
+_tile_zero(tmm2)
+_tile_dpbf16ps(tmm2, tmm0, tmm1);
+_tile_stored(tmm2, C_ptr, ldc);
+_tile_release();
+```
+
+The above code copies sub-matrices from memory (A_ptr, B_ptr) to Tile registers, calls the AMX BF16 compute instruction to multiply two sub-matrices, and then copies the result to memory (C_ptr).
+
+Taking INT8 as an example, AMX can perform the multiplication of two 16×64 sub-matrices (32,768 multiply/add operations) with a single instruction in 16 CPU cycles, enabling each core to complete 2048 multiply/add operations per cycle — 8 times the performance of AVX-512. On an Intel Xeon 4 CPU, a single core can theoretically provide 4 TOPS of compute power, making it highly suitable for compute-intensive tasks on the CPU.
+
+<p align="center">
+  <picture>
+    <img alt="amx_intro" src="../assets/amx_intro.png" width=60%>
+  </picture>
+</p>
+
+
+# AMX Kernel in KTransformers
+
+Before version v0.3, KTransformers performed CPU matrix multiplications based on operators provided by llamafile. Unfortunately, llamafile's implementation had not yet been optimized for the AMX instruction set. This resulted in performance bottlenecks, even in strong hardware environments (such as Xeon 4th Gen + 4090), where inference speeds for large models like DeepSeek-V3 reached only 91 tokens/s during the prefill phase. The CPU thus remained a significant bottleneck. In long prompt scenarios, such performance is clearly unsatisfactory. To fully unleash CPU potential, we introduced a brand-new AMX optimization path along with multiple technical improvements in v0.3.
+
+## 1. AMX Tiling-aware Memory Layout
+
+AMX provides a high-throughput Tile register computation model, reducing instruction count and boosting theoretical throughput through coarse-grained matrix operations. However, to truly exploit AMX's potential, memory access efficiency is critical: because AMX transfers entire Tiles at once, misaligned Tiles and chaotic access patterns can cause severe cache misses, nullifying throughput gains.
+
+Thus, in v0.3, we stopped directly memory-mapping GGUF-format files and introduced AMX Tiling-aware memory preprocessing during model loading. Specifically, expert weight matrices in MoE models are pre-rearranged into Tile-friendly sub-matrices whose shapes precisely match AMX Tile register dimensions, eliminating dynamic transposition overhead during inference. During rearrangement, we strictly align each sub-matrix's start address to 64 bytes to avoid cache line splits, and arrange sub-matrices sequentially according to computation access patterns, maximizing L1/L2 cache hit rates using compiler and hardware sequential prefetch capabilities.
+
+For Int8 quantized formats, we adopted Symmetric Group-wise Quantization, with each column forming a group sharing a scale factor stored separately to maintain memory alignment for Tile data.
+
+This AMX Tiling-aware memory layout design reduces memory latency while providing optimal input conditions for downstream computation kernels.
+
+## 2. Cache-friendly AMX Kernel
+
+During inference, we designed around the CPU’s multi-level cache hierarchy to perform computations in-place in high-speed caches, minimizing DRAM access frequency and overhead.
+
+<p align="center">
+  <picture>
+    <img alt="amx" src="../assets/amx.png" width=60%>
+  </picture>
+</p>
+
+As shown in the figure, 
+- ① Expert weight matrices are first column-wise partitioned into multiple tasks dynamically scheduled across threads. Input activations are shared among tasks and typically reside in the shared L3 cache due to locality.
+- ② Within each task, expert weights are row-wise partitioned into blocks, with block sizes finely tuned to ensure input activations, weights, and intermediate results stay within L2 cache, avoiding DRAM access.
+- ③ ④ ⑤ Each block is treated as a set of sub-matrices matching AMX Tile registers, and during Tile-level computation, input Tiles (tmm0–tmm1) and expert Tiles (tmm2–tmm3) are loaded, and four AMX multiplication instructions directly generate and accumulate products into Tile registers (tmm4–tmm7), with output activations accumulated in Tile registers or L1 cache, avoiding additional data movement.
+
+In short, we leveraged the cache hierarchy: every data element of expert weights and output activations accesses DRAM only once, with the other accesses hitting L2 or higher caches; input activations are accessed from DRAM only once and later hit in L3 or higher caches. This significantly reduces main memory traffic and improves overall execution efficiency.
+
+## 3. AVX-512 Kernel Adaptation for Low Arithmetic Intensity Scenarios
+
+Although AMX is highly efficient for large-scale matrix multiplication, it performs poorly under low arithmetic intensity, such as vector-matrix operations in the decode phase. This is because dispatching AMX Tiles involves fixed instruction overhead, which becomes wasteful when the data volume is insufficient to fill a Tile, causing reduced throughput.
+
+<p align="center">
+  <picture>
+    <img alt="amx_avx" src="../assets/amx_avx.png" width=60%>
+  </picture>
+</p>
+
+To address this, we introduced a lightweight AVX-512 kernel as a complement. This kernel follows the same memory layout as the AMX kernel but replaces heavy AMX matrix-matrix multiplications with fine-grained AVX-512 vector-matrix multiplications, lowering latency for small matrices.
+
+KTransformers dynamically selects between AMX and AVX-512 kernels at runtime based on arithmetic intensity: AMX kernels are automatically selected during long prompt prefill phases (where each expert handles more than 4 tokens on average), while short prompt prefill and decode phases dynamically switch to AVX-512 kernels. This ensures optimal efficiency under different arithmetic intensity conditions.
+
+## 4. MoE Operator Fusion and Dynamic Scheduling
+
+MoE models have many experts per layer, each requiring three matrix multiplications (Gate, Up, Down projections), leading to many small matrix multiplication tasks. Independently scheduling each small task would cause massive synchronization overhead between threads, dragging down overall inference speed.
+
+Thus, we fused the same type of matrix computations for all experts in a layer into large unified tasks. Furthermore, as there are no data dependencies between Gate and Up projections, their computations can also be fused, ultimately consolidating a layer’s matrix multiplications into two major tasks, greatly reducing scheduling overhead.
+
+To address load imbalance — especially during the prefill phase where expert activations can be highly skewed — we introduced a dynamic task scheduling strategy. Each matrix multiplication task is further split into multiple fine-grained sub-tasks, evenly distributed among CPU threads initially. Once a thread completes its assigned tasks, it atomically "steals" tasks from others, greatly mitigating load imbalance and achieving near-optimal CPU resource utilization.
+
+Thanks to these optimizations, our kernel can achieve 21 TFLOPS of BF16 throughput and 35 TOPS of Int8 throughput on Xeon4 CPUs — about 4× faster than PyTorch’s general AMX kernel. For DeepSeek-V3, pairing a Xeon4 CPU with a single RTX 4090 GPU achieves 418 tokens/s end-to-end throughput, close to the performance of multi-machine, multi-GPU setups. KTransformers’ AMX kernel is the first AMX kernel specifically designed for MoE inference scenarios, significantly lowering the hardware barrier for large model deployment and enabling more developers to enjoy GPU cluster level inference experiences at lower cost.
+
+<p align="center">
+  <picture>
+    <img alt="onednn_1" src="../assets/onednn_1.png" width=60%>
+  </picture>
+</p>
+
+# Usage
+
+## Checking AMX Support
+
+Before enabling the AMX-optimized kernels, it is important to verify whether your CPU supports the AMX instruction set. You can check AMX availability with the following command:
+
+```bash
+lscpu | grep -i amx
+```
+
+If your system supports AMX, you should see output similar to:
+
+```bash
+Flags: ... amx-bf16 amx-int8 amx-tile ...
+```
+
+If no amx-related flags are found, your CPU may not support AMX, or AMX may be disabled in BIOS settings. In that case, please ensure that:
+- You are using a Sapphire Rapids (Xeon 4th Gen) or newer CPU.
+- AMX support is enabled in your system BIOS under CPU feature settings.
+
+## Enabling AMX in KTransformers
+
+KTransformers allows users to easily switch between different backends through simple YAML configuration modifications. To enable AMX, modify the injection configuration of your experts by specifying backend as AMXInt8 or AMXBF16:
+
+```YAML
+- match:
+    name: "^model\\.layers\\..*\\.mlp\\.experts$"
+  replace:
+    class: ktransformers.operators.experts.KTransformersExperts    # custom MoE Kernel with expert parallelism
+    kwargs:
+      prefill_device: "cuda"
+      prefill_op: "KExpertsTorch"
+      generate_device: "cpu"
+      generate_op: "KExpertsCPU"
+      out_device: "cuda"
+      backend: "AMXInt8"  # or "AMXBF16" or "llamafile" (default)
+```
+
+**Note:** Currently, using AMXInt8 requires reading weights from a BF16 GGUF file and performing online quantization during model loading. This may cause slightly slower load times. Future versions will provide pre-quantized weights to eliminate this overhead.
\ No newline at end of file

ktransformers/models/configuration_qwen2_moe.py

+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Qwen2MoE model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Qwen2MoeConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Qwen2MoeModel`]. It is used to instantiate a
+    Qwen2MoE model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Qwen1.5-MoE-A2.7B" [Qwen/Qwen1.5-MoE-A2.7B"](https://huggingface.co/Qwen/Qwen1.5-MoE-A2.7B").
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151936):
+            Vocabulary size of the Qwen2MoE model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Qwen2MoeModel`]
+        hidden_size (`int`, *optional*, defaults to 2048):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 5632):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 24):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 16):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        decoder_sparse_step (`int`, *optional*, defaults to 1):
+            The frequency of the MoE layer.
+        moe_intermediate_size (`int`, *optional*, defaults to 1408):
+            Intermediate size of the routed expert.
+        shared_expert_intermediate_size (`int`, *optional*, defaults to 5632):
+            Intermediate size of the shared expert.
+        num_experts_per_tok (`int`, *optional*, defaults to 4):
+            Number of selected experts.
+        num_experts (`int`, *optional*, defaults to 60):
+            Number of routed experts.
+        norm_topk_prob (`bool`, *optional*, defaults to `False`):
+            Whether to normalize the topk probabilities.
+        output_router_logits (`bool`, *optional*, defaults to `False`):
+            Whether or not the router logits should be returned by the model. Enabeling this will also
+            allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
+        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+            The aux loss factor for the total loss.
+        mlp_only_layers (`List[int]`, *optional*, defaults to `[]`):
+            Indicate which layers use Qwen2MoeMLP rather than Qwen2MoeSparseMoeBlock
+            The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
+            If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
+
+    ```python
+    >>> from transformers import Qwen2MoeModel, Qwen2MoeConfig
+
+    >>> # Initializing a Qwen2MoE style configuration
+    >>> configuration = Qwen2MoeConfig()
+
+    >>> # Initializing a model from the Qwen1.5-MoE-A2.7B" style configuration
+    >>> model = Qwen2MoeModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen2_moe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=2048,
+        intermediate_size=5632,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+        decoder_sparse_step=1,
+        moe_intermediate_size=1408,
+        shared_expert_intermediate_size=5632,
+        num_experts_per_tok=4,
+        num_experts=60,
+        norm_topk_prob=False,
+        output_router_logits=False,
+        router_aux_loss_coef=0.001,
+        mlp_only_layers=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+
+        # MoE arguments
+        self.decoder_sparse_step = decoder_sparse_step
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.norm_topk_prob = norm_topk_prob
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
\ No newline at end of file