bf16 fa3

fa3.py

@@ -10,8 +10,7 @@ _flash_attn_func = get_kernel("kernels-community/vllm-flash-attn3").flash_attn_f
 
 @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
 def flash_attn_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    dtype = torch.float8_e4m3fn
-    outputs, lse = _flash_attn_func(q.to(dtype), k.to(dtype), v.to(dtype))
+    outputs, lse = _flash_attn_func(q, k, v)
     return outputs
 
 @flash_attn_func.register_fake