[fix](kt-kernel): fix AMX BF16 Full-FT expert gradients

kt-kernel/CMakeLists.txt

@@ -395,6 +395,7 @@ if(HOST_IS_X86)
                 # 获取不带扩展名的文件名作为 target 名
                 get_filename_component(test_name ${test_src} NAME_WE)
                 add_executable(${test_name} ${test_src} ${CMAKE_CURRENT_SOURCE_DIR}/cpu_backend/shared_mem_buffer.cpp)
+                target_compile_options(${test_name} PRIVATE ${ARCH_FLAGS})
                 target_link_libraries(${test_name} llama OpenMP::OpenMP_CXX numa)
             endforeach()
         endif()

kt-kernel/ext_bindings.cpp

@@ -325,22 +325,29 @@ class MOESFTBindings {
       intptr_t grad_down_lora_a;
       intptr_t grad_down_lora_b;
       intptr_t grad_weights;
+      intptr_t grad_gate_proj;
+      intptr_t grad_up_proj;
+      intptr_t grad_down_proj;
     };
     static void inner(void* args) {
       Args* args_ = (Args*)args;
       args_->cpuinfer->enqueue(&TP_MOE_SFT<T>::backward_binding, args_->moe, args_->grad_output, args_->grad_input,
                                args_->grad_gate_lora_a, args_->grad_gate_lora_b, args_->grad_up_lora_a,
                                args_->grad_up_lora_b, args_->grad_down_lora_a, args_->grad_down_lora_b,
-                               args_->grad_weights);
+                               args_->grad_weights, args_->grad_gate_proj, args_->grad_up_proj,
+                               args_->grad_down_proj);
     }
     static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE_SFT<T>> moe, intptr_t grad_output,
                                                             intptr_t grad_input, intptr_t grad_gate_lora_a,
                                                             intptr_t grad_gate_lora_b, intptr_t grad_up_lora_a,
                                                             intptr_t grad_up_lora_b, intptr_t grad_down_lora_a,
-                                                            intptr_t grad_down_lora_b, intptr_t grad_weights) {
+                                                            intptr_t grad_down_lora_b, intptr_t grad_weights,
+                                                            intptr_t grad_gate_proj, intptr_t grad_up_proj,
+                                                            intptr_t grad_down_proj) {
       Args* args = new Args{nullptr,          moe.get(),        grad_output,    grad_input,
                             grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b,
-                            grad_down_lora_a, grad_down_lora_b, grad_weights};
+                            grad_down_lora_a, grad_down_lora_b, grad_weights,
+                            grad_gate_proj,   grad_up_proj,     grad_down_proj};
       return std::make_pair((intptr_t)&inner, (intptr_t)args);
     }
   };
@@ -403,7 +410,15 @@ void bind_moe_sft_module(py::module_& moe_module, const char* name) {
              self.prepare_and_save_bwd((void*)gate, (void*)up, (void*)down, path);
            })
       .def("submit_backward_repack", &MoeClass::submit_backward_repack)
-      .def("wait_backward_repack", &MoeClass::wait_backward_repack);
+      .def("wait_backward_repack", &MoeClass::wait_backward_repack)
+      .def("get_profile_stats", &MoeClass::get_profile_stats, py::arg("reset") = false)
+      .def("reset_profile_stats", &MoeClass::reset_profile_stats)
+      // Update base weight BF16 pointers for reload_base_weights (full mode training)
+      // After calling this, call load_weights_task() to re-quantize BF16->AMX
+      .def("set_base_weight_pointers",
+           [](MoeClass& self, intptr_t gate, intptr_t up, intptr_t down) {
+             self.set_base_weight_pointers((void*)gate, (void*)up, (void*)down);
+           });
 }
 #endif  // defined(__x86_64__) && defined(USE_AMX_AVX_KERNEL)
 
@@ -779,7 +794,11 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
       .DEF_PTR_PROPERTY(MOESFTConfig, up_lora_a)
       .DEF_PTR_PROPERTY(MOESFTConfig, up_lora_b)
       .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_a)
-      .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_b);
+      .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_b)
+      .def_readwrite("full_weight_grad", &MOESFTConfig::full_weight_grad)
+      .DEF_PTR_PROPERTY(MOESFTConfig, grad_gate_proj)
+      .DEF_PTR_PROPERTY(MOESFTConfig, grad_up_proj)
+      .DEF_PTR_PROPERTY(MOESFTConfig, grad_down_proj);
 
   py::class_<MoE_Interface, std::shared_ptr<MoE_Interface>>(moe_module, "MoE_Interface");
 
@@ -800,7 +819,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
 #endif
 #if defined(__AVX512BF16__)
   // SFT MoE with LoRA support (BF16, INT8, INT4, AWQ, K2)
-  bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF>>(moe_module, "AMXBF16_SFT_MOE");
+  bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF16, AMX_BF16_MOE_TP>>(moe_module, "AMXBF16_SFT_MOE");
   bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int8>>(moe_module, "AMXInt8_SFT_MOE");
   bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4>>(moe_module, "AMXInt4_SFT_MOE");
   // bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4_1>>(moe_module, "AMXInt4_1_SFT_MOE");
@@ -809,7 +828,8 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
   // bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4SmallKGroup, AMX_K2_MOE_TP>>(moe_module,
   //                                                                                       "AMXInt4_KGroup_SFT_MOE");
   // SFT MoE with SkipLoRA=true (skip all LoRA computation in backward, only compute base weight grad_input)
-  bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF, AMX_MOE_TP, true>>(moe_module, "AMXBF16_SFT_MOE_SkipLoRA");
+  bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF16, AMX_BF16_MOE_TP, true>>(moe_module,
+                                                                                     "AMXBF16_SFT_MOE_SkipLoRA");
   bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int8, AMX_MOE_TP, true>>(moe_module, "AMXInt8_SFT_MOE_SkipLoRA");
   bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4, AMX_MOE_TP, true>>(moe_module, "AMXInt4_SFT_MOE_SkipLoRA");
   // bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4_1, AMX_MOE_TP, true>>(moe_module,

kt-kernel/operators/amx/bf16-moe.hpp

@@ -30,6 +30,8 @@
 template <class T = amx::GemmKernel224BF16>
 class AMX_BF16_MOE_TP : public AMX_MOE_BASE<T, AMX_BF16_MOE_TP<T>> {
   using Base = AMX_MOE_BASE<T, AMX_BF16_MOE_TP<T>>;
+
+ protected:
   using Base::config_;
   using Base::down_ba_;
   using Base::down_bb_;

kt-kernel/operators/amx/la/amx_raw_buffers.hpp

@@ -133,16 +133,13 @@ struct BufferBBF16Impl {
   }
   void set_data(void* new_ptr) { b = reinterpret_cast<ggml_bf16_t*>(new_ptr); }
 
-  void from_mat(ggml_bf16_t* src, int ith, int nth) {
-    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
-    int n_block_begin = n_start;
-    int n_block_size = n_end - n_block_begin;
+  void pack_block(ggml_bf16_t* src, int src_stride, int n_block_begin, int n_block_size) {
     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* s = (__m512i*)(src + (n_begin + i) * src_stride + 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);
@@ -155,6 +152,116 @@ struct BufferBBF16Impl {
       }
     }
   }
+
+  void from_mat(ggml_bf16_t* src, int ith, int nth) {
+    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+    int n_block_begin = n_start;
+    int n_block_size = n_end - n_block_begin;
+    pack_block(src + n_block_begin * k, k, n_block_begin, n_block_size);
+  }
+
+  void from_mat_strided(ggml_bf16_t* src, int src_stride, int ith, int nth) {
+    assert(src_stride >= k);
+    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+    int n_block_begin = n_start;
+    int n_block_size = n_end - n_block_begin;
+    if (n_block_size <= 0) return;
+    pack_block(src + (size_t)n_block_begin * src_stride, src_stride, n_block_begin, n_block_size);
+  }
+
+  void from_mat_transposed(ggml_bf16_t* src, int src_n, int src_k, int ith, int nth) {
+    assert(n == src_k && k == src_n);
+    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+    int n_block_begin = n_start;
+    int n_block_size = n_end - n_block_begin;
+    if (n_block_size <= 0) return;
+
+    thread_local std::vector<ggml_bf16_t> strip;
+    strip.resize((size_t)n_block_size * k);
+    constexpr int TILE = 32;
+    for (int c_tile = 0; c_tile < k; c_tile += TILE) {
+      int c_end = std::min(c_tile + TILE, k);
+      for (int r_tile = 0; r_tile < n_block_size; r_tile += TILE) {
+        int r_end = std::min(r_tile + TILE, n_block_size);
+        for (int c = c_tile; c < c_end; c++) {
+          for (int r = r_tile; r < r_end; r++) {
+            strip[(size_t)r * k + c] = src[(size_t)c * src_k + n_block_begin + r];
+          }
+        }
+      }
+    }
+    pack_block(strip.data(), k, n_block_begin, n_block_size);
+  }
+
+  void to_mat(ggml_bf16_t* dst, int ith, int nth) const {
+    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+    int n_block_begin = n_start;
+    int n_block_size = n_end - n_block_begin;
+    if (n_block_size <= 0) return;
+
+    alignas(64) ggml_bf16_t tile_copy[N_STEP * K_STEP];
+    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) {
+          const ggml_bf16_t* tile_src =
+              b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP;
+          memcpy(tile_copy, tile_src, sizeof(tile_copy));
+          transpose_16x16_32bit((__m512i*)tile_copy);
+          transpose_16x16_32bit((__m512i*)(tile_copy + TILE_N * K_STEP));
+          for (int i = 0; i < N_STEP; i++) {
+            __m512i* s = (__m512i*)(tile_copy + i * K_STEP);
+            __m512i* d = (__m512i*)(dst + (size_t)(n_block_begin + n_begin + i) * k + k_block_begin + k_begin);
+            avx512_copy_32xbf16(s, d);
+          }
+        }
+      }
+    }
+  }
+
+  void from_bb_transposed(const BufferBBF16Impl& src, int ith, int nth) {
+    assert(n == src.k && k == src.n);
+    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+    int dst_n_block_begin = n_start;
+    int dst_n_block_size = n_end - dst_n_block_begin;
+    if (dst_n_block_size <= 0) return;
+
+    auto tile_ptr = [](ggml_bf16_t* base, int total_n, int total_k, int abs_n, int abs_k) {
+      int n_block_begin = abs_n / N_BLOCK * N_BLOCK;
+      int n_within = abs_n - n_block_begin;
+      int n_block_size = std::min(N_BLOCK, total_n - n_block_begin);
+      int k_block_begin = abs_k / K_BLOCK * K_BLOCK;
+      int k_within = abs_k - k_block_begin;
+      return base + (size_t)n_block_begin * total_k + (size_t)k_block_begin * n_block_size +
+             (size_t)n_within * std::min(K_BLOCK, total_k - k_block_begin) + (size_t)k_within * N_STEP;
+    };
+
+    alignas(64) ggml_bf16_t src_tile[N_STEP * K_STEP];
+    alignas(64) ggml_bf16_t dst_tile[N_STEP * K_STEP];
+    for (int dst_n = 0; dst_n < dst_n_block_size; dst_n += N_STEP) {
+      for (int dst_k_block = 0; dst_k_block < k; dst_k_block += K_BLOCK) {
+        int dst_k_block_size = std::min(K_BLOCK, k - dst_k_block);
+        for (int dst_k = 0; dst_k < dst_k_block_size; dst_k += K_STEP) {
+          int abs_dst_n = dst_n_block_begin + dst_n;
+          int abs_dst_k = dst_k_block + dst_k;
+          ggml_bf16_t* src_ptr = tile_ptr(src.b, src.n, src.k, abs_dst_k, abs_dst_n);
+          memcpy(src_tile, src_ptr, sizeof(src_tile));
+          transpose_16x16_32bit((__m512i*)src_tile);
+          transpose_16x16_32bit((__m512i*)(src_tile + TILE_N * K_STEP));
+
+          for (int i = 0; i < N_STEP; i++) {
+            for (int j = 0; j < K_STEP; j++) {
+              dst_tile[j * K_STEP + i] = src_tile[i * K_STEP + j];
+            }
+          }
+          transpose_16x16_32bit((__m512i*)dst_tile);
+          transpose_16x16_32bit((__m512i*)(dst_tile + TILE_N * K_STEP));
+          ggml_bf16_t* dst_ptr = tile_ptr(b, n, k, abs_dst_n, abs_dst_k);
+          memcpy(dst_ptr, dst_tile, sizeof(dst_tile));
+        }
+      }
+    }
+  }
   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;