Support bf16/fp16 activations in CPU SDPA (#20611)

Summary:
Add support for reduced-precision (bf16 + f16) activations in the LLM SDPA op. Bf16 support is the main goal, but I wired up f16 at the same time, as it's straightforward.

This initial implementation just exposes the actual operator-level support and does not yet wire e2e through ET-LLM or anything else. It also initially only supports our simple CPU BLAS implementation. I'll wire up MKL as a follow up. I'm probably going to leverage Kleidi for the ARM GEMM kernels. Apple's Accelerate doesn't support BF16 gemms as far as I can tell. Kleidi will be fast on SME2 (M4/A18+) but won't hit AMX on M1-M3. I think that's fine, though. The i8mm/dot kernels should be sufficient. We can use BNNS if we need to close the gap.

## Perplexity
Measured on wikitext, M4 Max, ~40k token sanity check, 1000 token windows
| f32 | bf16 | delta |
| -- | -- | -- |
| 26.453 | 26.671 | +0.22 (+0.82%) |

## Performance
It is slow because it uses the fallback GEMM (as described above). Will wire Kleidi + MKL as a follow-up. Nothing uses bf16 SDPA yet, so this is fine.

(SDPA op timing - qwen3-1.7b)
  | Device | Context | Phase   | f32 (ms) | bf16 (ms) | bf16/f32 |
  |--------|--------:|---------|---------:|----------:|---------:|
  | M4 Max | 128     | Prefill |    1.855 |    10.180 |    5.5×  |
  | M4 Max | 128     | Decode  |    0.112 |     0.074 |    0.66× |
  | M4 Max | 1024    | Prefill |   76.930 |  1254.800 |   16.3×  |
  | M4 Max | 1024    | Decode  |    0.800 |     1.201 |    1.5×  |
  | S25    | 128     | Prefill |    3.924 |    43.760 |   11.2×  |
  | S25    | 128     | Decode  |    0.245 |     0.329 |    1.3×  |
  | S25    | 1024    | Prefill |  168.340 |  2790.700 |   16.6×  |
  | S25    | 1024    | Decode  |    1.614 |     2.698 |    1.7×  |


Differential Revision: D110246161

Pulled By: GregoryComer

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,35 @@ 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) {
+      auto do_gemm = [&](auto rt_tag) {
+        using rt = decltype(rt_tag);
+        ::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);
+      };
+      if (q_data.dtype == ScalarType::BFloat16) {
+        do_gemm(::executorch::aten::BFloat16{});
+      } else {
+        do_gemm(::executorch::aten::Half{});
+      }
+    } else {
+      ET_CHECK_MSG(false, "Reduced-precision q@k requires float accumulation");
+    }
   } else {
     ::executorch::cpublas::gemm(
         ::executorch::cpublas::TransposeType::Transpose,
@@ -251,7 +282,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 +292,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 +318,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 +333,37 @@ 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) {
+      auto do_gemm = [&](auto rt_tag) {
+        using rt = decltype(rt_tag);
+        ::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);
+      };
+      if (v_data.dtype == ScalarType::BFloat16) {
+        do_gemm(::executorch::aten::BFloat16{});
+      } else {
+        do_gemm(::executorch::aten::Half{});
+      }
+    } else {
+      ET_CHECK_MSG(false, "Reduced-precision qk@v requires float accumulation");
+    }
   } else {
     ::executorch::cpublas::gemm(
         ::executorch::cpublas::TransposeType::NoTranspose,
@@ -311,7 +374,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 +635,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 +835,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
@@ -1012,18 +1085,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);
@@ -1065,12 +1136,23 @@ 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) {
+            vec::convert<accum_t, scalar_t>(
+                qk_data, qk_reduced_data, qBlockSize * kvBlockSize);
+            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,
@@ -1083,12 +1165,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);