export.py: fix custom SDPA type conversion logic & re-enable for bfloat16

q.dtype is always torch.float at this point. Clearly meant to use the input dtype.

ghstack-source-id: 2c002e3
Pull Request resolved: #1171

Author: swolchok

torchchat/export.py

@@ -193,7 +193,7 @@ def forward(self, x, freqs_cis, mask, input_pos=None):
                 input_pos[-1].item(),
                 seqlen,
             )
-            output = output.view(bsz, seqlen, self.dim).to(dtype=q.dtype)
+            output = output.view(bsz, seqlen, self.dim).to(dtype=x.dtype)
             return self.wo(output)
 
     def replace_attention_with_custom_sdpa_attention(module: nn.Module):
@@ -285,11 +285,7 @@ def export_for_et(model, device, output_path) -> str:
             model = model.to(dtype=target_precision)
             state_dict_dtype = target_precision
 
-        # Custom SDPA does not work with bfloat16 on CPU currently. (The op doesn't
-        # support anything but bfloat32, and our attempt to use it anyway by converting
-        # to and from float causes other errors.)
-        if target_precision != torch.bfloat16:
-            replace_attention_with_custom_sdpa_attention(model)
+        replace_attention_with_custom_sdpa_attention(model)
 
         with torch.nn.attention.sdpa_kernel(
             [torch.nn.attention.SDPBackend.MATH]