[NVIDIA XLA:GPU] Fused MHA Support in XLA GPU: SE
@cheshire @akuegel Opened this new PR > I attached comments to the places I had to adjust because of ClangTidy and Linter. It might be easier for you if you revert the revert PR, then you have all those fixes already and can add additional fixes for the build time issues on top.

The reopened PR seemed to have just the last commit and I don't think I can revert the reverted PR. Are u able to do that on your end? Otherwise, I'll manually make the changes. > > I attached comments to the places I had to adjust because of ClangTidy and Linter. It might be easier for you if you revert the revert PR, then you have all those fixes already and can add additional fixes for the build time issues on top.
> 
> The reopened PR seemed to have just the last commit and I don't think I can revert the reverted PR. Are u able to do that on your end? Otherwise, I'll manually make the changes.

I could do that internally, and export it as a XLA PR. Would that help?
But I guess you are working in a git checkout of tensorflow? Then you can do

git revert 7a4d2e8b1a503a45ff86d40b0387fc832302379f

And use that as a base commit for your new PR.
cuDNN currently supports the following patterns for Multi-headed attention:

BMM1 - BMM2
BMM1 - Scale - Bias - Mask - Softmax - BMM2
BMM1 - Scale - Bias - Mask - Softmax - Dropout - BMM2
BMM1 - Scale - Mask - Softmax - BMM2
BMM1 - Scale - Mask - Softmax - Dropout - BMM2
BMM1 - Softmax - Dropout - BMM2

This PR adds support for the stream executor for these patterns.