fp32 or even fp64 support for cudnn flashattention

[wangleiphy]

Now, jax.nn.dot_product_attentiononly supports fp16 or bf16 when implementation='cudnn', see

jax/jax/_src/cudnn/fused_attention_stablehlo.py

Line 240 in e3c4b20

if q_dtype not in [jnp.bfloat16, jnp.float16]:

I am wondering whether is it possible to support fp32 or even fp64, that would be great for some application which requires numerical accuracy.

[superbobry]

CC @Cjkkkk