Preserve BF16 in custom SDPA (#22160)

examples/models/llama/BUCK

@@ -270,6 +270,7 @@ fbcode_target(_kind = runtime.python_test,
         ":sdpa",
         "//caffe2:torch",
         "//executorch/examples/models/llama:llama_transformer",
+        "//executorch/extension/llm/custom_ops:custom_ops_aot_py",
     ],
 )
 

examples/models/llama/source_transformation/sdpa.py

@@ -42,12 +42,12 @@ def forward(
         k = k.transpose(1, 2)
         v = v.transpose(1, 2)
 
-        # Custom op only supports float32 currently. Converting to/from float32 is
-        # faster than not having the op.
         input_dtype = q.dtype
-        q = q.to(dtype=torch.float)
-        k = k.to(dtype=torch.float)
-        v = v.to(dtype=torch.float)
+        use_fp32_fallback = input_dtype not in (torch.float, torch.bfloat16)
+        if use_fp32_fallback:
+            q = q.to(dtype=torch.float)
+            k = k.to(dtype=torch.float)
+            v = v.to(dtype=torch.float)
 
         if self.use_attention_mask:
             output = torch.ops.llama.custom_sdpa(
@@ -69,7 +69,10 @@ def forward(
                 0,  # dropout probability. Ignored by the code
                 True,  # is_causal
             )
-        return output.view(bsz, seqlen, self.dim).to(dtype=input_dtype)
+        output = output.view(bsz, seqlen, self.dim)
+        if use_fp32_fallback:
+            output = output.to(dtype=input_dtype)
+        return output
 
 
 def _replace_sdpa_with_custom_op(

examples/models/llama/source_transformation/test_sdpa_with_quantized_kv_cache.py

@@ -7,6 +7,7 @@
 import unittest
 
 import torch
+import torch.nn.functional as F
 
 from executorch.examples.models.llama.attention import KVCache
 
@@ -18,6 +19,8 @@
 
 from executorch.examples.models.llama.source_transformation.sdpa import SDPACustom
 
+from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
+
 
 class SDPAWithQuantizedKVCacheTest(unittest.TestCase):
     def _init_cache(self):
@@ -95,3 +98,41 @@ def test_simple(self, is_dynamic_shape=False):
             rtol=1e-03,
             atol=1e-03,
         )
+
+    def test_bfloat16_custom_sdpa_export_has_no_dtype_conversions(self):
+        bsz = 1
+        seqlen = 3
+        n_heads = 4
+        head_dim = 16
+        dim = n_heads * head_dim
+        input_pos = torch.tensor([0], dtype=torch.int64)
+        q = torch.rand((bsz, n_heads, seqlen, head_dim), dtype=torch.bfloat16)
+        k = torch.rand((bsz, n_heads, seqlen, head_dim), dtype=torch.bfloat16)
+        v = torch.rand((bsz, n_heads, seqlen, head_dim), dtype=torch.bfloat16)
+        mask = torch.triu(torch.full((seqlen, seqlen), float("-inf")), diagonal=1)
+        sdpa = SDPACustom(dim, use_attention_mask=True)
+
+        output = sdpa(input_pos, q, k, v, bsz, seqlen, mask)
+        reference = (
+            F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+            .transpose(1, 2)
+            .contiguous()
+            .view(bsz, seqlen, dim)
+        )
+        self.assertEqual(output.dtype, torch.bfloat16)
+        torch.testing.assert_close(
+            output.float(), reference.float(), atol=5e-2, rtol=5e-2
+        )
+
+        exported = torch.export.export(
+            sdpa,
+            (input_pos, q, k, v, bsz, seqlen, mask),
+        )
+        call_targets = [
+            node.target
+            for node in exported.graph_module.graph.nodes
+            if node.op == "call_function"
+        ]
+
+        self.assertIn(torch.ops.llama.custom_sdpa.default, call_targets)
+        self.assertNotIn(torch.ops.aten._to_copy.default, call_targets)

extension/llm/custom_ops/custom_ops.py

@@ -291,7 +291,7 @@ def custom_sdpa(
         scale,
     )
 
-    return torch.empty_like(query)
+    return torch.empty(query.shape, dtype=query.dtype, device=query.device)
 
 
 def _validate_update_cache_params(

extension/llm/custom_ops/op_sdpa_aot.cpp

@@ -247,7 +247,7 @@ at::Tensor custom_sdpa_aten(
     const bool is_causal,
     // @lint-ignore CLANGTIDY facebook-hte-ParameterMightThrowOnCopy
     const std::optional<double> scale) {
-  auto output = at::empty(q.sizes());
+  auto output = at::empty(q.sizes(), q.options());
   WRAP_TO_ATEN(custom_sdpa_out_no_context, 8)
   (q, k, v, start_pos, attn_mask, dropout_p, is_causal, scale, output);
   return output;