Pallas flash attention benchmarks

[jenkspt]

I'm very excited to see the new documentation for Pallas. I'm also excited to see a flash attention implementation. I'm wondering if there are any TPU / GPU benchmarks for the pallas flash attention kernel?

Related to #18121

[MoFHeka]

Any update？