[perf]: add K2 RAWINT4 prefill mat-mat dispatch

kt-kernel/ext_bindings.cpp

@@ -791,6 +791,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
   bind_moe_module<AMX_MOE_TP<amx::GemmKernel224Int4_1>>(moe_module, "AMXInt4_1_MOE");
   bind_moe_module<AMX_AWQ_MOE_TP<amx::GemmKernel224Int4_1_LowKGroup>>(moe_module, "AMXInt4_1KGroup_MOE");
   bind_moe_module<AMX_K2_MOE_TP<amx::GemmKernel224Int4SmallKGroup>>(moe_module, "AMXInt4_KGroup_MOE");
+  bind_moe_module<AMX_K2_MOE_TP<amx::GemmKernel224Int4SmallKGroupBlocked>>(moe_module, "AMXInt4_KGroupBlocked_MOE");
 #if defined(__AVX512F__)
   bind_moe_module<AMX_BF16_MOE_TP<amx::GemmKernel224BF16>>(moe_module, "AMXBF16_MOE");
   bind_moe_module<AMX_FP8_MOE_TP<amx::GemmKernel224FP8>>(moe_module, "AMXFP8_MOE");

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

@@ -241,6 +241,91 @@ class AMX_K2_MOE_TP : public AMX_MOE_BASE<T, AMX_K2_MOE_TP<T>> {
     }
   }
 
+
+  void write_weights_to_buffer_blocked(int gpu_tp_count, [[maybe_unused]] int cpu_tp_count, int expert_id,
+                                       const GeneralMOEConfig& full_config,
+                                       const std::vector<uintptr_t>& w13_weight_ptrs,
+                                       const std::vector<uintptr_t>& w13_scale_ptrs,
+                                       const std::vector<uintptr_t>& w2_weight_ptrs,
+                                       const std::vector<uintptr_t>& w2_scale_ptrs) const {
+    const int group_size = config_.quant_config.group_size;
+    auto pool = config_.pool->get_subpool(tp_part_idx);
+
+    constexpr int NUM_W13_TASKS = 32;
+    constexpr int NUM_W2_TASKS = 32;
+    const int total_tasks = NUM_W13_TASKS + NUM_W2_TASKS;
+
+    const int cpu_n_w13 = config_.intermediate_size;
+    const int cpu_k_w13 = config_.hidden_size;
+    const int gpu_n_w13 = full_config.intermediate_size / gpu_tp_count;
+    const int gpu_k_w13 = full_config.hidden_size;
+    const int global_n_offset_w13 = tp_part_idx * cpu_n_w13;
+    const size_t gpu_w13_weight_per_mat = static_cast<size_t>(gpu_n_w13) * gpu_k_w13 / 2;
+    const size_t gpu_w13_scale_per_mat = static_cast<size_t>(gpu_n_w13) * (gpu_k_w13 / group_size);
+
+    const int cpu_n_w2 = config_.hidden_size;
+    const int cpu_k_w2 = config_.intermediate_size;
+    const int gpu_k_w2 = full_config.intermediate_size / gpu_tp_count;
+    const int global_k_offset_w2 = tp_part_idx * cpu_k_w2;
+
+    pool->do_work_stealing_job(
+        total_tasks, nullptr,
+        [=, &w13_weight_ptrs, &w13_scale_ptrs, &w2_weight_ptrs, &w2_scale_ptrs, this](int task_id) {
+          if (task_id < NUM_W13_TASKS) {
+            const int rows_per_task = (cpu_n_w13 + NUM_W13_TASKS - 1) / NUM_W13_TASKS;
+            const int row_start = task_id * rows_per_task;
+            const int row_end = std::min(row_start + rows_per_task, cpu_n_w13);
+            for (int local_n = row_start; local_n < row_end; local_n++) {
+              const int global_n = global_n_offset_w13 + local_n;
+              const int target_gpu = global_n / gpu_n_w13;
+              const int n_in_gpu = global_n % gpu_n_w13;
+
+              uint8_t* w13_weight_base = reinterpret_cast<uint8_t*>(w13_weight_ptrs[target_gpu]);
+              ggml_bf16_t* w13_scale_base = reinterpret_cast<ggml_bf16_t*>(w13_scale_ptrs[target_gpu]);
+              const size_t weight_row_offset = static_cast<size_t>(n_in_gpu) * gpu_k_w13 / 2;
+              const size_t scale_row_offset = static_cast<size_t>(n_in_gpu) * (gpu_k_w13 / group_size);
+
+              gate_bb_[expert_id]->copy_weight_rows_to(w13_weight_base + weight_row_offset, local_n, 1, 0, cpu_k_w13,
+                                                        gpu_k_w13 / 2);
+              up_bb_[expert_id]->copy_weight_rows_to(w13_weight_base + gpu_w13_weight_per_mat + weight_row_offset,
+                                                      local_n, 1, 0, cpu_k_w13, gpu_k_w13 / 2);
+
+              gate_bb_[expert_id]->copy_scale_rows_to(w13_scale_base + scale_row_offset, local_n, 1, 0,
+                                                       cpu_k_w13 / group_size, gpu_k_w13 / group_size);
+              up_bb_[expert_id]->copy_scale_rows_to(w13_scale_base + gpu_w13_scale_per_mat + scale_row_offset, local_n,
+                                                     1, 0, cpu_k_w13 / group_size, gpu_k_w13 / group_size);
+            }
+            return;
+          }
+
+          const int w2_task_id = task_id - NUM_W13_TASKS;
+          const int rows_per_task = (cpu_n_w2 + NUM_W2_TASKS - 1) / NUM_W2_TASKS;
+          const int row_start = w2_task_id * rows_per_task;
+          const int row_end = std::min(row_start + rows_per_task, cpu_n_w2);
+          for (int row = row_start; row < row_end; row++) {
+            int k_local = 0;
+            while (k_local < cpu_k_w2) {
+              const int global_k = global_k_offset_w2 + k_local;
+              const int target_gpu = global_k / gpu_k_w2;
+              const int k_in_gpu = global_k % gpu_k_w2;
+              const int k_count = std::min(cpu_k_w2 - k_local, gpu_k_w2 - k_in_gpu);
+
+              uint8_t* w2_weight_base = reinterpret_cast<uint8_t*>(w2_weight_ptrs[target_gpu]);
+              ggml_bf16_t* w2_scale_base = reinterpret_cast<ggml_bf16_t*>(w2_scale_ptrs[target_gpu]);
+              uint8_t* weight_dst = w2_weight_base + static_cast<size_t>(row) * gpu_k_w2 / 2 + k_in_gpu / 2;
+              ggml_bf16_t* scale_dst =
+                  w2_scale_base + static_cast<size_t>(row) * (gpu_k_w2 / group_size) + k_in_gpu / group_size;
+
+              down_bb_[expert_id]->copy_weight_rows_to(weight_dst, row, 1, k_local, k_count, gpu_k_w2 / 2);
+              down_bb_[expert_id]->copy_scale_rows_to(scale_dst, row, 1, k_local / group_size, k_count / group_size,
+                                                       gpu_k_w2 / group_size);
+              k_local += k_count;
+            }
+          }
+        },
+        nullptr);
+  }
+
   // Write a single expert's weights to the output buffers
   // The caller provides pointers that already point to the target expert's location (no offset needed)
   // expert_id: the index of the expert to write
@@ -250,6 +335,12 @@ class AMX_K2_MOE_TP : public AMX_MOE_BASE<T, AMX_K2_MOE_TP<T>> {
                                const std::vector<uintptr_t>& w13_scale_ptrs,
                                const std::vector<uintptr_t>& w2_weight_ptrs,
                                const std::vector<uintptr_t>& w2_scale_ptrs) const {
+    if constexpr (T::BLOCKED_B_LAYOUT) {
+      write_weights_to_buffer_blocked(gpu_tp_count, cpu_tp_count, expert_id, full_config, w13_weight_ptrs,
+                                      w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs);
+      return;
+    }
+
     const int group_size = config_.quant_config.group_size;
     auto pool = config_.pool->get_subpool(tp_part_idx);
 

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

@@ -1141,6 +1141,121 @@ struct BufferBInt4KGroupImpl {
   }
 };
 
+// Block-major BufferB for signed RAWINT4 with KGroup scale. The external
+// format remains row-major packed int4; internally weights are stored by
+// N_BLOCK and 64-K-element tiles for prefill-friendly access.
+template <typename K>
+struct BufferBInt4KGroupBlockedImpl {
+  using dt = typename K::dt;
+  dt* b;
+  float* d;
+  int n, k, k_group_size, k_group_count;
+
+  static constexpr int N_STEP = K::N_STEP;
+  static constexpr int K_STEP = K::K_STEP;
+  static constexpr int N_BLOCK = K::N_BLOCK;
+  static constexpr int K_TILE = 64;
+  static constexpr int K_TILE_BYTES = K_TILE / 2;
+  static constexpr bool SCALE = true;
+
+  static size_t required_size(int n, int k, int k_group_size) {
+    return sizeof(int8_t) * n * k / 2 + sizeof(float) * n * (k / k_group_size);
+  }
+
+  BufferBInt4KGroupBlockedImpl(int n, int k, int k_group_size, void* ptr) : n(n), k(k), k_group_size(k_group_size) {
+    assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
+    if (n % N_STEP || k % K_TILE || k % k_group_size) {
+      printf("BufferBInt4KGroupBlockedImpl: n: %d, k: %d, N_STEP: %d, K_TILE: %d, k_group_size: %d\n", n, k,
+             N_STEP, K_TILE, k_group_size);
+      throw std::runtime_error("n or k is not aligned to blocked RAWINT4 layout");
+    }
+    k_group_count = k / k_group_size;
+    b = reinterpret_cast<dt*>(ptr);
+    d = reinterpret_cast<float*>(offset_pointer(b, n * k / 2));
+  }
+
+  size_t block_offset_bytes(int n_begin, int k_begin) const {
+    const int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
+    const int n_in_block = n_begin - n_block_begin;
+    const int n_block_size = std::min(N_BLOCK, n - n_block_begin);
+    const int k_tile = k_begin / K_TILE;
+    return static_cast<size_t>(n_block_begin) * k / 2 + static_cast<size_t>(k_tile) * n_block_size * K_TILE_BYTES +
+           static_cast<size_t>(n_in_block) * K_TILE_BYTES;
+  }
+
+  void from_raw_mat(uint8_t* proj, int ith, int nth) {
+    auto [n_start, n_end] = K::split_range_n(n, ith, nth);
+    if (n_start >= n_end) return;
+    const size_t src_row_bytes = static_cast<size_t>(k) / 2;
+    for (int row = n_start; row < n_end; row++) {
+      const uint8_t* src_row = proj + static_cast<size_t>(row) * src_row_bytes;
+      for (int k_begin = 0; k_begin < k; k_begin += K_TILE) {
+        uint8_t* dst = reinterpret_cast<uint8_t*>(b) + block_offset_bytes(row, k_begin);
+        std::memcpy(dst, src_row + k_begin / 2, K_TILE_BYTES);
+      }
+    }
+  }
+
+  dt* get_submat(int n_, int k_, int n_begin, int k_begin) {
+    (void)n_;
+    (void)k_;
+    return reinterpret_cast<dt*>(reinterpret_cast<uint8_t*>(b) + block_offset_bytes(n_begin, k_begin));
+  }
+
+  const uint8_t* get_kblock(int n_begin, int k_begin) const {
+    return reinterpret_cast<const uint8_t*>(b) + block_offset_bytes(n_begin, k_begin);
+  }
+
+  uint8_t* get_kblock(int n_begin, int k_begin) {
+    return reinterpret_cast<uint8_t*>(b) + block_offset_bytes(n_begin, k_begin);
+  }
+
+  float* get_scale(int n_, int n_begin, int k_, int k_begin) {
+    (void)n_;
+    int k_group_idx = k_begin / k_group_size;
+    return d + n_begin * (k_ / k_group_size) + k_group_idx;
+  }
+
+  const float* get_scale(int n_, int n_begin, int k_, int k_begin) const {
+    (void)n_;
+    int k_group_idx = k_begin / k_group_size;
+    return d + n_begin * (k_ / k_group_size) + k_group_idx;
+  }
+
+  void copy_weight_rows_to(uint8_t* dst, int row_start, int row_count, int k_start, int k_count,
+                           size_t dst_row_stride_bytes) const {
+    assert(k_start % K_TILE == 0);
+    assert(k_count % K_TILE == 0);
+    for (int r = 0; r < row_count; r++) {
+      uint8_t* dst_row = dst + static_cast<size_t>(r) * dst_row_stride_bytes;
+      const int src_row = row_start + r;
+      for (int kk = 0; kk < k_count; kk += K_TILE) {
+        std::memcpy(dst_row + kk / 2, get_kblock(src_row, k_start + kk), K_TILE_BYTES);
+      }
+    }
+  }
+
+  void copy_scale_rows_to(ggml_bf16_t* dst, int row_start, int row_count, int kg_start, int kg_count,
+                          size_t dst_row_stride_elems) const {
+    for (int r = 0; r < row_count; r++) {
+      const float* src = d + static_cast<size_t>(row_start + r) * k_group_count + kg_start;
+      ggml_bf16_t* dst_row = dst + static_cast<size_t>(r) * dst_row_stride_elems;
+      for (int kg = 0; kg < kg_count; kg++) {
+        dst_row[kg] = ggml_fp32_to_bf16(src[kg]);
+      }
+    }
+  }
+
+  static std::pair<int, int> split_range_n(int n, int ith, int nth) {
+    int n_per_thread = (n + nth - 1) / nth;
+    n_per_thread = (n_per_thread + N_STEP - 1) / N_STEP * N_STEP;
+    int n_start = std::min(ith * n_per_thread, n);
+    int n_end = std::min(n_start + n_per_thread, n);
+    return {n_start, n_end};
+  }
+
+};
+
 template <typename K>
 struct BufferBInt4WithZeroKGroupImpl {
   using dt = typename K::dt;