PR #10668: [XLA:GPU] Fix cuDNN FMHA rewriter sequence length checks and Bump up minimum flash attn cuDNN version to 8.9.4
PR #10668: [XLA:GPU] Fix cuDNN FMHA rewriter sequence length checks and Bump up minimum flash attn cuDNN version to 8.9.4

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

* If only seqlen_q or seqlen_kv is larger than 512, it will be lowered to fused attn instead of flash attn. This is incorrect as fused attn does not support this case. Add checks to lower to flash attn if one seqlen is larger than 512.
* While fixing the seqlen checks, also relax the seqlen/head dim constraints to support any seqlen % 2 == 0, head_dim <= 128 and head_dim % 8 = 0. (Since 8.9.4).
* Consolidate multiple get seqlen/head dim from fused attn/flash attn checks into one `getBHSD` function and rewrite both fused attn/flash attn checks to make it easier to understand.
* Bump up minimum cuDNN version require for flash attn from 8.9.3 to 8.9.4. (More seqlen/head dim support and cross attn support).
* Remove cross attn checks from rewriter and rewriter test since it is now default supported with 8.9.4.
* Add one testcase to cover the first bullet point.
Copybara import of the project:

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

rewrite seqlen checks

--
8a0e8b465cc1790953b1dc5f8740adfa7cf1124a by cjkkkk <ske@nvidia.com>:

return vector directly

--
72abf544d2f962c5554fe00ffea8d3167a9cde81 by cjkkkk <ske@nvidia.com>:

use struct for qkv_layout && add () for extra &&

Merging this change closes #10668

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/10668 from Cjkkkk:rewrite_seqlen_checks 72abf544d2f962c5554fe00ffea8d3167a9cde81
