Support bf16/fp16 activations in CPU SDPA

Activate the dormant mixed-precision path in cpu_flash_attention, using
float accumulation for reduced-precision (bf16/fp16) activations, and add
a reduced-in/float-out gemm to cpublas to back the two matmuls. Validation
and dispatch now accept Half/BFloat16 alongside Float, and the Python meta
accepts bf16/fp16.

Review order: start with kernels/optimized/blas (the new bf16/fp16 -> fp32
gemm primitive and the out_t threading), then op_sdpa_impl.h (the kernel
that uses it), then the validation/dispatch in op_sdpa.cpp and the Python
meta in custom_ops.py, then the C++ and Python tests.

Authored with Claude Code.

Author: GregoryComer

extension/llm/custom_ops/custom_ops.py

@@ -75,13 +75,16 @@ def _validate_params(
         value.dim() == 4
     ), f"Expected value to be 4 dimensional but got {value.dim()} dimensions."
 
+    supported_dtypes = (torch.float32, torch.bfloat16, torch.float16)
     assert (
-        query.dtype == torch.float32
-    ), f"Expected query to be float32 but got {query.dtype}"
-    assert key.dtype == torch.float32, f"Expected key to be float32 but got {key.dtype}"
+        query.dtype in supported_dtypes
+    ), f"Expected query to be float32, bfloat16, or float16 but got {query.dtype}"
     assert (
-        value.dtype == torch.float32
-    ), f"Expected value to be float32 but got {value.dtype}"
+        key.dtype == query.dtype
+    ), f"Expected key to have the same dtype as query but got {key.dtype}"
+    assert (
+        value.dtype == query.dtype
+    ), f"Expected value to have the same dtype as query but got {value.dtype}"
 
     assert (
         key_cache.dim() == 4
@@ -91,11 +94,11 @@ def _validate_params(
     ), f"Expected value_cache to be 4 dimensional but got {value_cache.dim()}"
 
     assert (
-        key_cache.dtype == torch.float32
-    ), f"Expected key_cache to be float32 but got {key_cache.dtype}"
+        key_cache.dtype == query.dtype
+    ), f"Expected key_cache to have the same dtype as query but got {key_cache.dtype}"
     assert (
-        value_cache.dtype == torch.float32
-    ), f"Expected value_cache to be float32 but got {value_cache.dtype}"
+        value_cache.dtype == query.dtype
+    ), f"Expected value_cache to have the same dtype as query but got {value_cache.dtype}"
 
     assert (
         key_cache.size() == value_cache.size()

extension/llm/custom_ops/op_sdpa.cpp

@@ -45,8 +45,10 @@ bool validate_flash_attention_args(
 
   ET_CHECK_OR_RETURN_FALSE(
       (query.scalar_type() == ScalarType::Float) ||
+          (query.scalar_type() == ScalarType::Half) ||
+          (query.scalar_type() == ScalarType::BFloat16) ||
           (query.scalar_type() == ScalarType::Char),
-      "Query must be Float type");
+      "Query must be Float, Half, BFloat16, or Char type");
 
   ET_CHECK_OR_RETURN_FALSE(
       (query.scalar_type() == key.scalar_type()) &&
@@ -266,7 +268,7 @@ Tensor& flash_attention_kernel_out(
 
   auto seq_len = query.size(2);
 
-  ET_SWITCH_FLOAT_TYPES(
+  ET_SWITCH_FLOATHBF16_TYPES(
       query.scalar_type(), ctx, "flash_attention", CTYPE, [&] {
         // TODO we need to re-evaluate this for ARM CPUs
         // And there can be many so instead of templatizing
@@ -414,7 +416,7 @@ Tensor& custom_sdpa_out_impl(
 
   // TODO(task): replace the template param selection logic
   // with whatever apprpriately makes more sense for
-  ET_SWITCH_FLOAT_TYPES(
+  ET_SWITCH_FLOATHBF16_TYPES(
       output.scalar_type(), ctx, "flash_attention", CTYPE, [&] {
         // TODO we need to re-evaluate this for ARM CPUs
         // And there can be many so instead of templatizing

extension/llm/custom_ops/op_sdpa_impl.h

@@ -74,8 +74,10 @@ void _q_at_k_gemm(
     accum_t* qk_data) {
   ET_CHECK_MSG(q_data.dtype == k_data.dtype, "q and k must have same dtype");
   ET_CHECK_MSG(
-      q_data.dtype == ScalarType::Char || q_data.dtype == ScalarType::Float,
-      "q and k must be either int8 or float");
+      q_data.dtype == ScalarType::Char || q_data.dtype == ScalarType::Float ||
+          q_data.dtype == ScalarType::Half ||
+          q_data.dtype == ScalarType::BFloat16,
+      "q and k must be int8, float, half, or bfloat16");
   if (q_data.dtype == ScalarType::Char) {
     if constexpr (std::is_same<accum_t, float>::value) {
       int a_stride_m_tmp, b_stride_n_tmp;
@@ -103,6 +105,46 @@ void _q_at_k_gemm(
       ET_CHECK_MSG(
           false, "Accumulation in dtype other than float not supported yet");
     }
+  } else if (
+      q_data.dtype == ScalarType::BFloat16 ||
+      q_data.dtype == ScalarType::Half) {
+    if constexpr (std::is_same<accum_t, float>::value) {
+      if (q_data.dtype == ScalarType::BFloat16) {
+        using rt = ::executorch::aten::BFloat16;
+        ::executorch::cpublas::gemm(
+            ::executorch::cpublas::TransposeType::Transpose,
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            k_n,
+            q_m,
+            qk_k,
+            1.0f,
+            static_cast<const rt*>(k_data.data),
+            k_stride_n,
+            static_cast<const rt*>(q_data.data),
+            q_stride_m,
+            0.0f,
+            qk_data,
+            k_n);
+      } else {
+        using rt = ::executorch::aten::Half;
+        ::executorch::cpublas::gemm(
+            ::executorch::cpublas::TransposeType::Transpose,
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            k_n,
+            q_m,
+            qk_k,
+            1.0f,
+            static_cast<const rt*>(k_data.data),
+            k_stride_n,
+            static_cast<const rt*>(q_data.data),
+            q_stride_m,
+            0.0f,
+            qk_data,
+            k_n);
+      }
+    } else {
+      ET_CHECK_MSG(false, "Reduced-precision q@k requires float accumulation");
+    }
   } else {
     ::executorch::cpublas::gemm(
         ::executorch::cpublas::TransposeType::Transpose,
@@ -251,7 +293,7 @@ void _qk_at_v_gemm(
     const int64_t m,
     const int64_t n,
     const int64_t k,
-    const accum_t* qk_data,
+    const void* qk_data,
     const int64_t qk_stride_m,
     const MaybeQuantizedMatrixData& v_data,
     const int64_t v_stride_n,
@@ -261,14 +303,15 @@ void _qk_at_v_gemm(
     accum_t* buf_qdq_ptr) {
   if (v_data.dtype == ScalarType::Char) {
     if constexpr (std::is_same<accum_t, float>::value) {
+      const float* qk = static_cast<const float*>(qk_data);
       if (m > 4) {
         // For larger batch sizes, dequantize and use BLAS for better
         // performance
         dequant_and_gemm(
             m,
             n,
             k,
-            const_cast<float*>(qk_data),
+            const_cast<float*>(qk),
             qk_stride_m,
             v_data,
             v_stride_n,
@@ -286,7 +329,7 @@ void _qk_at_v_gemm(
             m,
             n,
             k,
-            qk_data,
+            qk,
             qk_stride_m /*lhs_stride_m*/,
             static_cast<const int8_t*>(v_data.data),
             v_stride_n /*rhs_stride_n*/,
@@ -301,6 +344,48 @@ void _qk_at_v_gemm(
       ET_CHECK_MSG(
           false, "Accumulation in dtype other than float not supported yet");
     }
+  } else if (
+      v_data.dtype == ScalarType::BFloat16 ||
+      v_data.dtype == ScalarType::Half) {
+    // qk has been cast to the activation dtype (see qk_reduced_data); both
+    // operands are reduced precision and accumulate into the float output.
+    if constexpr (std::is_same<accum_t, float>::value) {
+      if (v_data.dtype == ScalarType::BFloat16) {
+        using rt = ::executorch::aten::BFloat16;
+        ::executorch::cpublas::gemm(
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            n,
+            m,
+            k,
+            1.0f,
+            static_cast<const rt*>(v_data.data),
+            v_stride_n,
+            static_cast<const rt*>(qk_data),
+            qk_stride_m,
+            beta,
+            o_data,
+            o_stride_m);
+      } else {
+        using rt = ::executorch::aten::Half;
+        ::executorch::cpublas::gemm(
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            ::executorch::cpublas::TransposeType::NoTranspose,
+            n,
+            m,
+            k,
+            1.0f,
+            static_cast<const rt*>(v_data.data),
+            v_stride_n,
+            static_cast<const rt*>(qk_data),
+            qk_stride_m,
+            beta,
+            o_data,
+            o_stride_m);
+      }
+    } else {
+      ET_CHECK_MSG(false, "Reduced-precision qk@v requires float accumulation");
+    }
   } else {
     ::executorch::cpublas::gemm(
         ::executorch::cpublas::TransposeType::NoTranspose,
@@ -311,7 +396,7 @@ void _qk_at_v_gemm(
         static_cast<accum_t>(1),
         static_cast<const accum_t*>(v_data.data),
         v_stride_n,
-        qk_data,
+        static_cast<const accum_t*>(qk_data),
         qk_stride_m,
         beta,
         o_data,
@@ -572,11 +657,9 @@ void cpu_flash_attention(
   constexpr bool is_reduced_type =
       ::executorch::runtime::is_reduced_floating_point_v<scalar_t>;
 
-  ET_CHECK_MSG(
-      !is_reduced_type, "FlashAttention does not support reduced types.");
-  // Figure out mixed precision a little later
-  // using accum_t = at::opmath_type<scalar_t>;
-  using accum_t = scalar_t;
+  // Reduced-precision (bf16/fp16) activations accumulate in float: the two
+  // matmuls run as reduced-in/float-out gemms and the softmax stays in float.
+  using accum_t = std::conditional_t<is_reduced_type, float, scalar_t>;
   using Vec = vec::Vectorized<accum_t>;
   accum_t scaling_factor = static_cast<accum_t>(calculate_scale(query, scale));
 
@@ -774,7 +857,19 @@ void cpu_flash_attention(
   } else {
     buf = scratch.get();
   }
+  std::unique_ptr<char[]> allocated_buf_reduced;
   void* buf_reduced = nullptr;
+  if (is_reduced_type) {
+    int64_t size_reduced_bytes =
+        qSplitSize * kvSplitSize * num_thread * sizeof(scalar_t);
+    Result<void*> scratch_reduced = ctx.allocate_temp(size_reduced_bytes, 64);
+    if (!scratch_reduced.ok()) {
+      allocated_buf_reduced = std::make_unique<char[]>(size_reduced_bytes);
+      buf_reduced = allocated_buf_reduced.get();
+    } else {
+      buf_reduced = scratch_reduced.get();
+    }
+  }
   int64_t size_per_thread_qdq_vec = kvSplitSize * headSize;
   // Lets align size_per_thread_qdq_vec to 64 bytes, for coalesced cache reads,
   // by padding with right number of per thread elements
@@ -1015,18 +1110,16 @@ void cpu_flash_attention(
           if (tmp_max == -std::numeric_limits<accum_t>::infinity()) {
             // to avoid `nan = exp2f(-inf - (-inf))`
             fill_stub(
-                conditional_data_ptr(qk_data, qk_reduced_data) +
-                    row * kvBlockSize,
-                static_cast<scalar_t>(0),
+                qk_data + row * kvBlockSize,
+                static_cast<accum_t>(0),
                 kvBlockSize);
           } else {
             // qk <- exp(qk - max) and sum per row
             tmp_sum = tmp_max;
             _exp_reduce_sum_fusion_kernel(
                 qk_data + row * kvBlockSize,
                 kvBlockSize,
-                conditional_data_ptr(qk_data, qk_reduced_data) +
-                    row * kvBlockSize,
+                qk_data + row * kvBlockSize,
                 tmp_sum);
             // exp_tmp <- exp(max[row] - max)
             exp_tmp = std::exp(qk_max_data[row] - tmp_max);
@@ -1068,12 +1161,24 @@ void cpu_flash_attention(
             headSize,
             v_quant_params_StrideN,
             value.scalar_type());
+        // For reduced-precision activations the attention weights are cast to
+        // the activation dtype so that Softmax(q @ k.T) @ v runs as a
+        // reduced-in/float-out gemm matching the value matrix.
+        const void* qk_gemm_data = qk_data;
+        if constexpr (is_reduced_type) {
+          if (!is_quantized_sdpa) {
+            for (int64_t idx = 0; idx < qBlockSize * kvBlockSize; ++idx) {
+              qk_reduced_data[idx] = static_cast<scalar_t>(qk_data[idx]);
+            }
+            qk_gemm_data = qk_reduced_data;
+          }
+        }
         // Calculate Softmax(q @ k.T) @ v
         _qk_at_v_gemm<accum_t>(
             qBlockSize,
             headSize,
             kvBlockSize,
-            qk_data,
+            qk_gemm_data,
             kvBlockSize,
             v_sub_matrix_data,
             vStrideN,
@@ -1086,12 +1191,20 @@ void cpu_flash_attention(
       // reorder MHA output with strides
       for (int64_t row = 0; row < qBlockSize; ++row) {
         accum_t sum_reciprocal = 1 / qk_sum_data[row];
-        vec::map<scalar_t>(
-            [sum_reciprocal](Vec x) { return x * Vec(sum_reciprocal); },
-            out_data + i * oStrideB + j * oStrideH + m * oStrideM +
-                row * oStrideM,
-            dst_data + row * headSize,
-            headSize);
+        scalar_t* out_row = out_data + i * oStrideB + j * oStrideH +
+            m * oStrideM + row * oStrideM;
+        const accum_t* dst_row = dst_data + row * headSize;
+        if constexpr (is_reduced_type) {
+          for (int64_t d = 0; d < headSize; ++d) {
+            out_row[d] = static_cast<scalar_t>(dst_row[d] * sum_reciprocal);
+          }
+        } else {
+          vec::map<scalar_t>(
+              [sum_reciprocal](Vec x) { return x * Vec(sum_reciprocal); },
+              out_row,
+              dst_row,
+              headSize);
+        }
       }
       // Move to the next query
       data_index_step(i, batchSize, j, num_head, k, qSlice);