@@ -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
 
@@ -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 &&
 
@@ -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);
 
@@ -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,
 
@@ -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(