PR #11249: [XLA:GPU] Update cuDNN flash attention backward to use new cuDNN frontend (>1.0)
PR #11249: [XLA:GPU] Update cuDNN flash attention backward to use new cuDNN frontend (>1.0)

Imported from GitHub PR https://github.com/openxla/xla/pull/11249

* Replace complicated node by node flash attention cuDNN graph construction with simple cudnn sdpa call provided in new cuDNN frontend (>1.0).
* Remove softmax_sum/dq_accum tensors from the custom call outputs since it is not required anymore in new cuDNN frontend.
* Add a pass to query workspace from cuDNN graph instead of manual calculation as workspace size is subject to change.
Copybara import of the project:

--
2a1ee169fb2c12fbbd4e791ff406b3272490cf09 by cjkkkk <ske@nvidia.com>:

init

--
e84388a1859ef26acad7ba68cebedf546098d99b by cjkkkk <ske@nvidia.com>:

fix stat datatype to be fp32

--
3a0b75ccfa88791d5cc437a9cc86333849d5fee3 by cjkkkk <ske@nvidia.com>:

use temp value for workspace & remove softmax buffer/dq_accum

--
77e6db53f4f4d0b5344e3492da21e871ceb0bb0b by cjkkkk <ske@nvidia.com>:

add workspace upperbound

--
11dd85db2ebea5cbd12b7108bdc995d8f8f5f1ba by cjkkkk <ske@nvidia.com>:

use tighter upperbound

--
2f245303893f0aef5dc48557f8254d16ae373d9d by cjkkkk <ske@nvidia.com>:

remove unused param

--
5e648ae07c743e95955f23fbd1c4fa50a13be5de by cjkkkk <ske@nvidia.com>:

add pass to fix workspace size

--
368e51879c10266a3ecc200c105afc37dcba544c by cjkkkk <ske@nvidia.com>:

remove comments

--
2fb7b3854f57c80e0c2e7805349ff5b637ba2007 by cjkkkk <ske@nvidia.com>:

rm header order change

--
d0cb184a32a5489f17ca5a20cd11901544e88ba1 by cjkkkk <ske@nvidia.com>:

move def close to use and remove comments

--
f09eb59328d2416374c31c5336744231b67954cd by cjkkkk <ske@nvidia.com>:

move use to def

--
1c56f1fe670cfc43cc1bfcbc16ed75cf9ce5357f by cjkkkk <ske@nvidia.com>:

add missing dependency in BUILD

--
efd8ff45b5cf13e0ae820c03150068a6412204d0 by cjkkkk <ske@nvidia.com>:

pass in correct descriptor to activation

Merging this change closes #11249

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11249 from Cjkkkk:flash-attn-backward efd8ff45b5cf13e0ae820c03150068a6412204d0
