Add KleiDI BF16 GEMM and SDPA dispatch (pytorch#22886)

Summary:

Add KleiDI NEON and SME2 BF16 GEMM dispatch for Arm64 and route BF16 SDPA prefill through it. Keep single-column decode on BFDOT and retain fallbacks for unsupported CPUs or older KleiDI builds. Scalar NEON packing and an adapter for the existing SME RHS packer allow this integration to land before the new upstream packers.

AI-assisted: Codex.

Reviewed By: digantdesai, Gasoonjia

Differential Revision: D119433385

Pull Request resolved: pytorch#22886

Author: JakeStevens

CMakeLists.txt

@@ -1195,6 +1195,22 @@ if(EXECUTORCH_BUILD_KERNELS_TORCHAO)
 
 endif()
 
+if(TARGET cpublas
+   AND TARGET kleidiai
+   AND TARGET cpuinfo
+)
+  if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64"
+     OR ANDROID_ABI STREQUAL "arm64-v8a"
+     OR CMAKE_OSX_ARCHITECTURES MATCHES "arm64"
+  )
+    target_compile_definitions(cpublas PRIVATE ET_BUILD_WITH_KLEIDIAI)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
+      target_compile_definitions(cpublas PRIVATE ET_KLEIDIAI_DISABLE_NEON_BF16)
+    endif()
+    target_link_libraries(cpublas PRIVATE kleidiai)
+  endif()
+endif()
+
 # The shared build ships the profiler as one of its libraries, and the Python
 # extension records a hard dependency on it, and the Apple ETDump extension
 # links it, so the target has to exist whenever any of those is being built

extension/llm/custom_ops/op_sdpa_impl.h

@@ -955,6 +955,12 @@ void cpu_flash_attention(
   int64_t qSplitSize = q_split_size > qSize ? qSize : q_split_size;
   int64_t kvSplitSize = kv_split_size > kvSize ? kvSize : kv_split_size;
   int64_t qSlice = (qSize - 1) / qSplitSize + 1;
+  const auto can_use_kleidiai_bfloat16_prefill = [&](int64_t q_block_size) {
+    return seq_dim == SeqDim::TWO && qSize > 1 &&
+        std::is_same<scalar_t, ::executorch::aten::BFloat16>::value &&
+        ::executorch::cpublas::gemm_uses_kleidiai_bfloat16(
+               ::executorch::cpublas::TransposeType::NoTranspose, q_block_size);
+  };
 #ifdef ET_USE_THREADPOOL
   int64_t num_thread =
       ::executorch::extension::threadpool::get_threadpool()->get_thread_count();
@@ -1005,11 +1011,15 @@ void cpu_flash_attention(
   // Scratch for widening q@K.T to fp32 (see _q_at_k_gemm): one K block plus one
   // q block. qBlockSize cannot exceed qSplitSize, so include the runtime bounds
   // that determine whether any block can use the widened path.
+  const auto can_widen_qk = [&](int64_t q_block_size) {
+    return !can_use_kleidiai_bfloat16_prefill(q_block_size) &&
+        std::is_same<scalar_t, ::executorch::aten::BFloat16>::value &&
+        ::executorch::cpublas::gemm_uses_blas() &&
+        headSize <= kMaxHeadSizeForWidenedQK &&
+        q_block_size >= kMinQBlockForWidenedQK;
+  };
   const bool widen_reduced_qk =
-      std::is_same<scalar_t, ::executorch::aten::BFloat16>::value &&
-      ::executorch::cpublas::gemm_uses_blas() &&
-      headSize <= kMaxHeadSizeForWidenedQK &&
-      qSplitSize >= kMinQBlockForWidenedQK;
+      can_widen_qk(qSplitSize) || can_widen_qk(qSize % qSplitSize);
 #if defined(__APPLE__)
   const bool dequantize_qk = is_quantized_sdpa &&
       ::executorch::cpublas::gemm_uses_blas() && qSplitSize > 4;
@@ -1082,6 +1092,9 @@ void cpu_flash_attention(
     for (int64_t z = begin; z < end; z++) {
       int64_t m = k * qSplitSize;
       int64_t qBlockSize = std::min(qSplitSize, qSize - m);
+      const bool use_kleidiai_bfloat16_prefill =
+          can_use_kleidiai_bfloat16_prefill(qBlockSize);
+      const bool widen_qk_block = can_widen_qk(qBlockSize);
       // Initialize max and sum
       fill_stub(
           qk_max_data, -std::numeric_limits<accum_t>::infinity(), qBlockSize);
@@ -1181,10 +1194,8 @@ void cpu_flash_attention(
             k_sub_matrix_data,
             kStrideN,
             qk_data,
-            ((widen_reduced_qk && qBlockSize >= kMinQBlockForWidenedQK) ||
-             (dequantize_qk && qBlockSize > 4))
-                ? widen_ptr
-                : nullptr);
+            (widen_qk_block || (dequantize_qk && qBlockSize > 4)) ? widen_ptr
+                                                                  : nullptr);
 
         // Update coefficients with scaling, attention mask, and softmax.
         accum_t tmp_max = 0, tmp_sum = 0, exp_tmp = 0;
@@ -1298,7 +1309,8 @@ void cpu_flash_attention(
         // them in accum_t, widen V and let BLAS multiply -- also one rounding
         // step fewer. Below it the widening stops amortizing.
         constexpr int64_t kMinQBlockForWidenedAV = 64;
-        const bool widen_v = is_reduced_type && !is_quantized_sdpa &&
+        const bool widen_v = !use_kleidiai_bfloat16_prefill &&
+            is_reduced_type && !is_quantized_sdpa &&
             std::is_same<scalar_t, ::executorch::aten::BFloat16>::value &&
             qBlockSize >= kMinQBlockForWidenedAV;
         const bool use_fp32_qk_weights = is_reduced_type &&

extension/llm/custom_ops/op_sdpa_test.cpp

@@ -839,6 +839,60 @@ TEST(OpScaledDotProductAttentionTest, BFloat16MatchesFloat) {
       2e-2, 2e-2);
 }
 
+TEST(OpScaledDotProductAttentionTest, BFloat16PrefillTrailingQueryBlocks) {
+  using executorch::aten::BFloat16;
+  TensorFactory<executorch::aten::ScalarType::BFloat16> tf_bfloat16;
+  TensorFactory<executorch::aten::ScalarType::Float> tf_float;
+  constexpr int32_t kHeadSize = 8;
+  constexpr int32_t kKeys = 513;
+
+  // These sequences use 64-query blocks and leave 1-4 queries in the tail.
+  // A second KV block also exercises accumulation into the first block's
+  // result.
+  for (const int32_t queries : {193, 194, 195, 196}) {
+    SCOPED_TRACE(::testing::Message() << "queries=" << queries);
+    auto query = tf_bfloat16.zeros({1, 1, queries, kHeadSize});
+    auto key = tf_bfloat16.zeros({1, 1, kKeys, kHeadSize});
+    auto value = tf_bfloat16.zeros({1, 1, kKeys, kHeadSize});
+    auto out = tf_bfloat16.zeros({1, 1, queries, kHeadSize});
+    auto query_float = tf_float.zeros({1, 1, queries, kHeadSize});
+    auto key_float = tf_float.zeros({1, 1, kKeys, kHeadSize});
+    auto value_float = tf_float.zeros({1, 1, kKeys, kHeadSize});
+    auto expected = tf_float.zeros({1, 1, queries, kHeadSize});
+    for (int32_t i = 0; i < queries * kHeadSize; ++i) {
+      const float v = static_cast<float>((i * 7) % 31 - 15) / 16.0f;
+      query.mutable_data_ptr<BFloat16>()[i] = BFloat16(v);
+      query_float.mutable_data_ptr<float>()[i] = v;
+    }
+    for (int32_t i = 0; i < kKeys * kHeadSize; ++i) {
+      const float k = static_cast<float>((i * 11) % 29 - 14) / 16.0f;
+      const float v = static_cast<float>((i * 13) % 37 - 18) / 16.0f;
+      key.mutable_data_ptr<BFloat16>()[i] = BFloat16(k);
+      key_float.mutable_data_ptr<float>()[i] = k;
+      value.mutable_data_ptr<BFloat16>()[i] = BFloat16(v);
+      value_float.mutable_data_ptr<float>()[i] = v;
+    }
+    op_scaled_dot_product_attention(
+        query, key, value, std::nullopt, 0.0, false, std::nullopt, out);
+    op_scaled_dot_product_attention(
+        query_float,
+        key_float,
+        value_float,
+        std::nullopt,
+        0.0,
+        false,
+        std::nullopt,
+        expected);
+    for (int32_t i = 0; i < queries * kHeadSize; ++i) {
+      EXPECT_NEAR(
+          static_cast<float>(out.const_data_ptr<BFloat16>()[i]),
+          expected.const_data_ptr<float>()[i],
+          2e-3f)
+          << "index=" << i;
+    }
+  }
+}
+
 TEST(OpScaledDotProductAttentionTest, HalfMatchesFloat) {
   test_reduced_precision_matches_float<executorch::aten::ScalarType::Half>(
       1e-2, 1e-2);

kernels/optimized/blas/CPUBlas.cpp

@@ -8,7 +8,10 @@
 
 #include <executorch/kernels/optimized/blas/CPUBlas.h>
 
-#include <limits.h>
+#include <cstddef>
+#include <cstdint>
+
+#include <executorch/kernels/optimized/blas/KleidiBlas.h>
 
 #ifdef ET_BUILD_WITH_BLAS
 #ifdef ET_BUILD_FOR_APPLE
@@ -267,6 +270,12 @@ void gemm(
     const float beta,
     float *c, int64_t ldc) {
   normalize_last_dims(transa, transb, m, n, k, &lda, &ldb, &ldc);
+#if defined(ET_KLEIDIAI_HAS_NEON_BF16) || defined(ET_KLEIDIAI_HAS_SME2_BF16)
+  if (kleidiai_bfloat16_gemm(
+          transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)) {
+    return;
+  }
+#endif
   gemm_impl<BFloat16, float, float>(
       transa, transb,
       m, n, k,

kernels/optimized/blas/CPUBlas.h

@@ -50,6 +50,9 @@ inline char to_blas(TransposeType trans) {
 // this library.
 bool gemm_uses_blas();
 
+// Whether the BF16-input, float-output gemm can use KleidiAI for this shape.
+bool gemm_uses_kleidiai_bfloat16(TransposeType transb, int64_t n);
+
 // Column-major c = beta * c + alpha * (a @ b), where a is BFloat16 and b and
 // c are float.
 void gemv(