make torch use flash-attn

[pass-lin]

from #2143
I found that I can't use flash-attn under torch. This doesn't seem to be a reasonable setting. According to some previous reports, this might be because there are some errors under torch. Such as #2145 and keras-team/keras#20459
So I first submitted a change to see where the test will report an error, and then adjust his code.

[mattdangerw]

@divyashreepathihalli any thoughts on the overall goal with this PR?

In terms of implementation, let's wait till #2184 goes in and add this on top. We should add some unit tests to at the very least make sure the fused attention op is getting hit only when it should. We are in general trying to make the flash attention code less spaghetti, since there's so many cases to handle.

[divyashreepathihalli]

Can you please also add a mock test similar to Here

[pass-lin]

Can you please also add a mock test similar to Here

Of course, I have no problem. But before that, please run a GPU test to see where there are problems under the current code.

[divyashreepathihalli]

The branch has conflicts please resolve them for the tests to run.

[pass-lin]

@divyashreepathihalli
I can run the test locally, so I don't quite understand what went wrong.

[pass-lin]

I've been quite busy lately, and I think I should focus on these PRs for now

So I'll close this PR first.I will restart and submit this PR in the near future.

keras-team/keras#21160
#2192
#2177
@divyashreepathihalli