perf(cpu): productize MiniCPM5 native KV-head GQA

Author: Aharrypotter

docs/assets/minicpm5-pr0-performance-strategy.png

@@ -143,9 +143,9 @@ if(MLLM_KERNEL_USE_THREADS AND MLLM_KERNEL_THREADS_VENDOR_OPENMP)
     # Apple should not use OpenMP
     message(ERROR "Apple platform should not use OpenMP. Pls set MLLM_KERNEL_THREADS_VENDOR_APPLE_GCD=ON")
   else()
-    target_link_libraries(MllmRT PUBLIC ${OpenMP_CXX_FLAGS})
-    target_compile_options(MllmRT PRIVATE ${OpenMP_CXX_FLAGS})
-    target_include_directories(MllmRT PUBLIC ${OpenMP_CXX_INCLUDE_DIR})
+    target_link_libraries(MllmCPUBackend PUBLIC ${OpenMP_CXX_FLAGS})
+    target_compile_options(MllmCPUBackend PRIVATE ${OpenMP_CXX_FLAGS})
+    target_include_directories(MllmCPUBackend PUBLIC ${OpenMP_CXX_INCLUDE_DIR})
   endif()
 endif()
 

mllm/backends/cpu/CMakeLists.txt

@@ -24,6 +24,7 @@
 #include "mllm/backends/cpu/ops/FlashAttn2WithSinkAndSwaOp.hpp"
 #include "mllm/backends/cpu/ops/GELUOp.hpp"
 #include "mllm/backends/cpu/ops/GatherOp.hpp"
+#include "mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp"
 #include "mllm/backends/cpu/ops/InterpolateOp.hpp"
 #include "mllm/backends/cpu/ops/LayerNorm2DOp.hpp"
 #include "mllm/backends/cpu/ops/MaskedScatterOp.hpp"
@@ -83,7 +84,7 @@ CPUBackend::CPUBackend() : Backend(kCPU, createCPUAllocator()) {
                CPUConv2DOpFactory, CPULayerNorm2DOpFactory, CPUInterpolateOpFactory, CPUPadOpFactory, CPUMaskedScatterOpFactory,
                CPUArgsortOpFactory, CPUCloneOpFactory, CPUAvgPool1dOpFactory, CPUFlashAttention2SwaSinkOpFactory,
                CPURadixAttnRelaxOpFactory, CPURadixAttnSwaSinkOpFactory, CPUEqualOpFactory, CPUWhereOpFactory,
-               CPUGatherOpFactory>();
+               CPUGatherOpFactory, CPUGroupedQueryAttentionDecodeOpFactory>();
 }
 
 CPUBackend::~CPUBackend() {

mllm/backends/cpu/CPUBackend.cpp

@@ -394,14 +394,13 @@ void KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul(float* __restrict__ dst, con
     const int m_step = ukernels_[tile_cfg].get_m_step();  // Scheduling along M
     const int n_step = ukernels_[tile_cfg].get_n_step();  // Scheduling along N
 
-    std::vector<std::pair<int, int>> tile_splits;
-    for (int i_m_step = 0; i_m_step < M; i_m_step += m_step) {
-      for (int i_n_step = 0; i_n_step < N; i_n_step += n_step) { tile_splits.emplace_back(i_m_step, i_n_step); }
-    }
-    auto tile_sizes = tile_splits.size();
+    const int m_tiles = (M + m_step - 1) / m_step;
+    const int n_tiles = (N + n_step - 1) / n_step;
+    const int tile_count = m_tiles * n_tiles;
 
-    MLLM_CONDITIONAL_PARALLEL_FOR(thread_count > 1, thread_count, tile_idx, 0, tile_sizes, 1, {
-      auto [i_m_step, i_n_step] = tile_splits[tile_idx];
+    MLLM_CONDITIONAL_PARALLEL_FOR(thread_count > 1, thread_count, tile_idx, 0, tile_count, 1, {
+      const int i_m_step = (tile_idx / n_tiles) * m_step;
+      const int i_n_step = (tile_idx % n_tiles) * n_step;
 
       // Support functions return offset in bytes
       const void* lhs_ptr = (const void*)((const char*)workspace + (ukernels_[tile_cfg].get_lhs_packed_offset(i_m_step, K)));

mllm/backends/cpu/kernels/arm/linear/kai.cpp

@@ -0,0 +1,124 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+
+#include "mllm/backends/cpu/kernels/common/fa2_1/arch.hpp"
+#include "mllm/backends/cpu/kernels/common/fa2_1/impl-any.hpp"
+#include "mllm/core/DataTypes.hpp"
+#include "mllm/utils/CPUArchHelper.hpp"
+
+#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM)
+#include "mllm/backends/cpu/kernels/common/fa2_1/impl-arm.hpp"
+#elif defined(MLLM_HOST_ARCH_X86_64) || defined(MLLM_HOST_ARCH_X86)
+#include "mllm/backends/cpu/kernels/common/fa2_1/impl-any-simd.hpp"
+#endif
+
+namespace mllm::cpu::gqa_decode {
+
+struct BhsdStrides {
+  int32_t batch;
+  int32_t head;
+  int32_t sequence;
+  int32_t dimension;
+};
+
+namespace detail {
+
+#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM)
+using attention_arch_tag = flash_attn2::details::arm_arch_tag;
+#elif defined(MLLM_HOST_ARCH_X86_64) || defined(MLLM_HOST_ARCH_X86)
+using attention_arch_tag = flash_attn2::details::x86_arch_tag;
+#else
+using attention_arch_tag = flash_attn2::details::any_arch_tag;
+#endif
+
+inline bool validStrides(const BhsdStrides& strides) {
+  return strides.batch > 0 && strides.head > 0 && strides.sequence > 0 && strides.dimension == 1;
+}
+
+}  // namespace detail
+
+// Single-token float32 GQA for native KV-head [B, H, S, D] cache views.
+// QK and softmax preserve C10's complete per-query-head traversal. P@V groups
+// the query heads sharing one KV head so each V token is reused while hot, but
+// every output head still accumulates in increasing key-index order.
+// Scratch: group_size * kv_sequence floats.
+inline bool fwdBhsdFp32(int32_t batch_size, int32_t query_heads, int32_t kv_heads, int32_t kv_sequence, int32_t qk_dim,
+                        int32_t value_dim, const mllm_fp32_t* query, BhsdStrides query_strides, const mllm_fp32_t* key,
+                        BhsdStrides key_strides, const mllm_fp32_t* value, BhsdStrides value_strides, mllm_fp32_t* output,
+                        BhsdStrides output_strides, mllm_fp32_t* grouped_scratch) {
+  if (batch_size <= 0 || query_heads <= 0 || kv_heads <= 0 || kv_sequence <= 0 || qk_dim <= 0 || value_dim <= 0
+      || query_heads % kv_heads != 0 || query == nullptr || key == nullptr || value == nullptr || output == nullptr
+      || grouped_scratch == nullptr || !detail::validStrides(query_strides) || !detail::validStrides(key_strides)
+      || !detail::validStrides(value_strides) || !detail::validStrides(output_strides)) {
+    return false;
+  }
+
+  using ArchTag = detail::attention_arch_tag;
+  const int32_t group_size = query_heads / kv_heads;
+  const float scale = 1.0F / std::sqrt(static_cast<float>(qk_dim));
+
+  for (int32_t batch = 0; batch < batch_size; ++batch) {
+    for (int32_t kv_head = 0; kv_head < kv_heads; ++kv_head) {
+      const int32_t first_query_head = kv_head * group_size;
+      const auto* query_group = query + static_cast<size_t>(batch) * query_strides.batch
+                                + static_cast<size_t>(first_query_head) * query_strides.head;
+      const auto* key_head = key + static_cast<size_t>(batch) * key_strides.batch
+                             + static_cast<size_t>(kv_head) * key_strides.head;
+      const auto* value_head = value + static_cast<size_t>(batch) * value_strides.batch
+                               + static_cast<size_t>(kv_head) * value_strides.head;
+      auto* output_group = output + static_cast<size_t>(batch) * output_strides.batch
+                           + static_cast<size_t>(first_query_head) * output_strides.head;
+
+      for (int32_t group_index = 0; group_index < group_size; ++group_index) {
+        const auto* query_token = query_group + static_cast<size_t>(group_index) * query_strides.head;
+        auto* probabilities = grouped_scratch + static_cast<size_t>(group_index) * kv_sequence;
+        float maximum = std::numeric_limits<float>::lowest();
+        for (int32_t key_index = 0; key_index < kv_sequence; ++key_index) {
+          const auto* key_token = key_head + static_cast<size_t>(key_index) * key_strides.sequence;
+          float score = 0.0F;
+          flash_attn2::details::VectorDotProduct<ArchTag, mllm_fp32_t, mllm_fp32_t, mllm_fp32_t>::run(
+              query_token, key_token, &score, static_cast<size_t>(qk_dim));
+          probabilities[key_index] = score * scale;
+          maximum = std::max(maximum, probabilities[key_index]);
+        }
+
+        float denominator = 0.0F;
+        for (int32_t key_index = 0; key_index < kv_sequence; ++key_index) {
+          probabilities[key_index] = std::exp(probabilities[key_index] - maximum);
+          denominator += probabilities[key_index];
+        }
+        const float inverse_denominator = 1.0F / denominator;
+        for (int32_t key_index = 0; key_index < kv_sequence; ++key_index) {
+          probabilities[key_index] *= inverse_denominator;
+        }
+      }
+
+      for (int32_t group_index = 0; group_index < group_size; ++group_index) {
+        auto* output_token = output_group + static_cast<size_t>(group_index) * output_strides.head;
+        flash_attn2::details::FilledWithConst<ArchTag, mllm_fp32_t>::run(output_token, 0.0F,
+                                                                       static_cast<size_t>(value_dim));
+      }
+      for (int32_t key_index = 0; key_index < kv_sequence; ++key_index) {
+        const auto* value_token = value_head + static_cast<size_t>(key_index) * value_strides.sequence;
+        auto* probability = grouped_scratch + key_index;
+        auto* output_token = output_group;
+        for (int32_t group_index = 0; group_index < group_size; ++group_index) {
+          flash_attn2::details::FMAConstArray<ArchTag, mllm_fp32_t, mllm_fp32_t, mllm_fp32_t>::run(
+              output_token, *probability, value_token, static_cast<size_t>(value_dim));
+          probability += kv_sequence;
+          output_token += output_strides.head;
+        }
+      }
+    }
+  }
+  return true;
+}
+
+}  // namespace mllm::cpu::gqa_decode

mllm/backends/cpu/kernels/common/gqa_decode/fwd_bhsd.hpp

@@ -0,0 +1,104 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#include "mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <vector>
+
+#include "mllm/backends/cpu/kernels/common/gqa_decode/fwd_bhsd.hpp"
+
+namespace mllm::cpu {
+namespace {
+
+void groupedQueryAttentionDecodeFloat32Reference(const Tensor& query, const Tensor& key, const Tensor& value, Tensor& output) {
+  const auto q_shape = query.shape();
+  const auto k_shape = key.shape();
+  const auto v_shape = value.shape();
+  const auto q_stride = query.stride();
+  const auto k_stride = key.stride();
+  const auto v_stride = value.stride();
+  const int32_t groups = q_shape[1] / k_shape[1];
+  const float scale = 1.0F / std::sqrt(static_cast<float>(q_shape[3]));
+
+  static thread_local std::vector<float> probabilities;
+  probabilities.resize(static_cast<size_t>(k_shape[2]));
+
+  for (int32_t batch = 0; batch < q_shape[0]; ++batch) {
+    for (int32_t query_head = 0; query_head < q_shape[1]; ++query_head) {
+      const int32_t kv_head = query_head / groups;
+      const auto* q_head = query.coffsettedPtr<float>({batch, query_head, 0, 0});
+      const auto* k_head = key.coffsettedPtr<float>({batch, kv_head, 0, 0});
+      const auto* v_head = value.coffsettedPtr<float>({batch, kv_head, 0, 0});
+      auto* output_head = output.offsettedPtr<float>({batch, query_head, 0, 0});
+
+      // Android release builds use -ffast-math; a finite sentinel keeps stable
+      // softmax valid under finite-math assumptions.
+      float maximum = std::numeric_limits<float>::lowest();
+      for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) {
+        const auto* key_token = k_head + static_cast<size_t>(key_index) * k_stride[2];
+        float score = 0.0F;
+        for (int32_t dim = 0; dim < q_shape[3]; ++dim) {
+          score += q_head[static_cast<size_t>(dim) * q_stride[3]] * key_token[static_cast<size_t>(dim) * k_stride[3]];
+        }
+        probabilities[static_cast<size_t>(key_index)] = score * scale;
+        maximum = std::max(maximum, probabilities[static_cast<size_t>(key_index)]);
+      }
+
+      float denominator = 0.0F;
+      for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) {
+        auto& probability = probabilities[static_cast<size_t>(key_index)];
+        probability = std::exp(probability - maximum);
+        denominator += probability;
+      }
+      const float inverse_denominator = 1.0F / denominator;
+      for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) {
+        probabilities[static_cast<size_t>(key_index)] *= inverse_denominator;
+      }
+
+      for (int32_t value_dim = 0; value_dim < v_shape[3]; ++value_dim) {
+        float result = 0.0F;
+        for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) {
+          const auto* value_token = v_head + static_cast<size_t>(key_index) * v_stride[2];
+          result += probabilities[static_cast<size_t>(key_index)] * value_token[static_cast<size_t>(value_dim) * v_stride[3]];
+        }
+        output_head[value_dim] = result;
+      }
+    }
+  }
+}
+
+}  // namespace
+
+CPUGroupedQueryAttentionDecodeOp::CPUGroupedQueryAttentionDecodeOp(const aops::GroupedQueryAttentionDecodeOpOptions& options)
+    : aops::GroupedQueryAttentionDecodeOp(options) {}
+
+void CPUGroupedQueryAttentionDecodeOp::forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) {
+  const auto& query = inputs[0];
+  const auto& key = inputs[1];
+  const auto& value = inputs[2];
+  auto& output = outputs[0];
+  const auto q_shape = query.shape();
+  const auto k_shape = key.shape();
+  const auto v_shape = value.shape();
+  const auto q_stride = query.stride();
+  const auto k_stride = key.stride();
+  const auto v_stride = value.stride();
+  const auto output_stride = output.stride();
+
+  static thread_local std::vector<float> probabilities;
+  const int32_t group_size = q_shape[1] / k_shape[1];
+  probabilities.resize(static_cast<size_t>(group_size) * k_shape[2]);
+
+  const bool completed = cpu::gqa_decode::fwdBhsdFp32(
+      q_shape[0], q_shape[1], k_shape[1], k_shape[2], q_shape[3], v_shape[3], query.ptr<float>(),
+      {q_stride[0], q_stride[1], q_stride[2], q_stride[3]}, key.ptr<float>(),
+      {k_stride[0], k_stride[1], k_stride[2], k_stride[3]}, value.ptr<float>(),
+      {v_stride[0], v_stride[1], v_stride[2], v_stride[3]}, output.ptr<float>(),
+      {output_stride[0], output_stride[1], output_stride[2], output_stride[3]}, probabilities.data());
+  if (!completed) { groupedQueryAttentionDecodeFloat32Reference(query, key, value, output); }
+}
+
+}  // namespace mllm::cpu

mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.cpp

@@ -0,0 +1,26 @@
+// Copyright (c) MLLM Team.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "mllm/core/BaseOp.hpp"
+#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp"
+
+namespace mllm::cpu {
+
+class CPUGroupedQueryAttentionDecodeOp final : public aops::GroupedQueryAttentionDecodeOp {
+ public:
+  explicit CPUGroupedQueryAttentionDecodeOp(const aops::GroupedQueryAttentionDecodeOpOptions& options);
+
+  void forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
+};
+
+class CPUGroupedQueryAttentionDecodeOpFactory
+    : public TypedOpFactory<OpTypes::kGroupedQueryAttentionDecode, aops::GroupedQueryAttentionDecodeOpOptions> {
+ protected:
+  std::shared_ptr<BaseOp> createOpImpl(const aops::GroupedQueryAttentionDecodeOpOptions& options) override {
+    return std::make_shared<CPUGroupedQueryAttentionDecodeOp>(options);
+  }
+};
+
+}  // namespace mllm::cpu