[Bug] XLA fail to compile cudnn_fmha by CudnnFusedMHARewriter
@SuryanarayanaY @AyanmoI @Tixxx topic move from #62661 to here. @sachinprasadhs @akuegel @Cheshire Any progress? If you don't have time for the moment, can you give me an idea to fix this bug? I'm really in a hurry. @akuegel  @MoFHeka apologies for the late response, I have been out of office. I will take a look at this issue. I'm getting this error when running your repro:
ValueError: Invalid stage selected: 'hoo'. Valid values are: 'hlo', 'hlo_serialized', 'optimized_hlo', 'optimized_hlo_serialized', 'optimized_hlo_dot'

Is your TF a custom build?
Also to make it more efficient, it'd better if you could provide the compiled HLO graph directly. Can you run your repro with these xla flags?
`XLA_FLAGS="--xla_gpu_enable_cudnn_fmha=true --xla_dump_hlo_as_text --xla_dump_to=OUTPUT_LOCATION --xla_dump_hlo_pass_re=.*" ` @Tixxx 
Here I ran:
```bash
TF_DUMP_GRAPH_PREFIX=generated_test/ TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" XLA_FLAGS="--xla_gpu_enable_cudnn_fmha --xla_dump_hlo_as_text --xla_dump_hlo_pass_re=.* --xla_dump_to=generated_test/" python test_xla_mha.py
```
Here is the zip file:
[generated_test_hlo.zip](https://github.com/tensorflow/tensorflow/files/13774793/generated_test_hlo.zip)

One more thing, may I ask how could I trigger MHA_Layout::QKV_INTERLEAVED type when using TF XLA with user interface python code? Here is the [example](https://github.com/NVIDIA/cudnn-frontend/blob/9f82dda5c029d15a5f371f0fe003dc0c74a0c987/samples/legacy_samples/helpers.cpp#L250C27-L250C54) showed in cud-frontend repo.
When I check tensor slice in HLO, it's a ordinary scheme like:
```python
strideA[seqlen_dim_idx] = h * d;
strideA[hidden_dim_idx] = 1;
strideA[head_dim_idx]   = d;
strideA[batch_dim_idx]  = s_q * h * d;
```
Not a QKV 3X form.

 @MoFHeka Hi, right now you can't trigger QKV_INTERLEAVED in XLA since MHA rewriter in XLA does not pattern match this case. This feature will be added in the future. @MoFHeka Also this zip file seems invalid. Could you verify if it is valid?
 One thing I noticed in your mlir is the fmha call is not lowered properly to a function call,
in your mlir:
`"lmhlo_gpu.fMHA"(%view, %view_1, %view_3, %view_5, %view_6) {......} : (memref<2x4096x16x64xbf16>, memref<2x4096x16x64xbf16>, memref<2x4096x16x64xbf16>, memref<2x4096x16x64xbf16>, memref<0xui8>) -> () "fmha-bmm-bmm"`

the lowered mlir should look something like:
`call @xla.gpu.fused.attention.bmm.bmm`

This would suggest that either the runtime lowering logic has a bug or your xla build doesn't have the runtime lowering code we merged a few months ago.

what's the HEAD of your xla build?
 > @MoFHeka Also this zip file seems invalid. Could you verify if it is valid?

zip file seems work well. I try to upload it again.
[generated_test.zip](https://github.com/tensorflow/tensorflow/files/13782675/generated_test.zip)
 > One thing I noticed in your mlir is the fmha call is not lowered properly to a function call, in your mlir: `"lmhlo_gpu.fMHA"(%view, %view_1, %view_3, %view_5, %view_6) {......} : (memref<2x4096x16x64xbf16>, memref<2x4096x16x64xbf16>, memref<2x4096x16x64xbf16>, memref<2x4096x16x64xbf16>, memref<0xui8>) -> () "fmha-bmm-bmm"`
> 
> the lowered mlir should look something like: `call @xla.gpu.fused.attention.bmm.bmm`
> 
> This would suggest that either the runtime lowering logic has a bug or your xla build doesn't have the runtime lowering code we merged a few months ago.
> 
> what's the HEAD of your xla build?

I use TF2.4 from image nvcr.io/nvidia/tensorflow:23.11-tf2-py3. When I update TF to latest stable 2.5, it seems all work well. Thank you!

But there is another question. Why is FMHA almost no faster than a regular XLA? Also which time we can use FMHA with dim 128 and so on? It should be faster if you have softmax and training. Right now, XLA(cuDNN) has two fmha kernels implemented. The first one is fused attention and the second one is flash attention. Fused attention only support sequence length up to 512 and head dim 64 and it is not likely to support 128 in the future. Flash attention on the other hand, supports all any sequence length and head dim 128. Right now, XLA choose between these 2 kernels based on whether sequence length is larger than 512 or not. @MoFHeka  have you tried with softmax and other intermediate nodes to get some perf improvement? if you dont see other errors during compilation, maybe we can close this issue? @Tixxx The improvement is very small, even within 10%, I don't know if it is a measurement error. @Tixxx One more question, when using Keras Attention, query shape is (batch_size, seq_len_q, num_heads * head_dim). But FMHA only support shape (batch_size, seq_len_q, num_heads, head_dim). 
But if use Keras MultiHeadAttention, it use einsum not matmul to calculate.
FMHA will never work.!!!
[Keras Attention]<img width="594" alt="Keras Attention" src="https://github.com/tensorflow/tensorflow/assets/90189118/b94826bb-aa80-4804-a28e-9a3c24518bad">
 > @Tixxx One more question, when using Keras Attention, query shape is (batch_size, seq_len_q, num_heads * head_dim). But FMHA only support shape (batch_size, seq_len_q, num_heads, head_dim). But if use Keras MultiHeadAttention, it use einsum not matmul to calculate. FMHA will never work.!!! [Keras Attention]<img alt="Keras Attention" width="594" src="https://private-user-images.githubusercontent.com/90189118/294462790-b94826bb-aa80-4804-a28e-9a3c24518bad.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MDQ0NzcwNDUsIm5iZiI6MTcwNDQ3Njc0NSwicGF0aCI6Ii85MDE4OTExOC8yOTQ0NjI3OTAtYjk0ODI2YmItYWE4MC00ODA0LWEyOGUtOWEzYzI0NTE4YmFkLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDAxMDUlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwMTA1VDE3NDU0NVomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTRiNWJiMzE3MDY4NTA4YWZiMGZiM2U0ZDI2MGJjNmEwMjM1NTJmMzA3ZDlmMDIyODBlNDM0MDcxMTRiMmUyM2EmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.XwVZyt41uooLS6IZcbgDifVv8SqkUVWZCJ8axgt8cbE">

Yea i think the tf.keras.layers.Attention is mainly for seq2seq models where one layer can only attend to 1 position, it's similar to an MHA with a single head, CUDNN lib supports single-headed MHAs but without hacking this keras api to give us the correct shape, XLA won't lower it. I think the keras.mha API might be a better option to develop transformer models, https://github.com/keras-team/keras/blob/v2.14.0/keras/layers/attention/multi_head_attention.py#L130-L731
Einsum will be lowered to a gemm eventually in xla so we are able to match it(Jax uses einsum too and we are able to match most of the jax MHA blocks). I checked the source code above for keras MHA, the shape is in the correct form. Maybe give it a try and let us know if that works? @Tixxx Unfortunately, Einsum will be converted to triton kernel instead of cudnn matmul in both TF and Jax(Praxis or Lax).
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

tf 2.14

### Custom code

No

### OS platform and distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

CUDA 12.3/cuDNN 8.9.6

### GPU model and memory

A800 80GB

### Current behavior?

set xla_gpu_enable_cudnn_fmha, try using XLA to jit a simple bmm1-bmm2 function by CudnnFusedMHARewriter.

### Standalone code to reproduce the issue
```python
import os

import tensorflow as tf
import numpy as np


batch_size = 16
seq_len_q = 16
num_heads = 256
head_dim = 64
mat_dtypes = tf.bfloat16

# @tf.function(jit_compile=True)
# def build_tf_xla_graph(queries, keys, values):
#   keras_attention_layer = tf.keras.layers.Attention(causal=True, use_scale=False, dtype=mat_dtypes)
#   flash_attention_output = keras_attention_layer(inputs=[queries, values, keys])
#   return flash_attention_output

@tf.function(jit_compile=True)
def build_tf_xla_graph(queries, keys, values):
  score = tf.matmul(queries, keys, transpose_b=True)
  flash_attention_output = tf.matmul(score, values)
  return flash_attention_output

seq_len_v = seq_len_q
queries = tf.random.normal(
    (batch_size, seq_len_q, num_heads, head_dim), dtype=mat_dtypes)
keys = tf.random.normal(
    (batch_size, seq_len_v, num_heads, head_dim), dtype=mat_dtypes)
values = tf.random.normal(
    (batch_size, seq_len_v, num_heads, head_dim), dtype=mat_dtypes)

ir = build_tf_xla_graph.experimental_get_compiler_ir(queries, keys, values)(stage='hoo')
print(ir)

output = build_tf_xla_graph(queries, keys, values)
print(output)
```
and run:
```shell
XLA_FLAGS="--xla_gpu_enable_cudnn_fmha" python test_xla_mha.py
```


### Relevant log output

```shell
2023-12-19 16:03:11.235376: I tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter.cc:645] Before CudnnFusedMHARewriter:
HloModule a_inference_build_tf_xla_graph_26__XlaMustCompile_true_config_proto_6001324581131673121_executor_type_11160318154034397263_.14, alias_passthrough_params=true, entry_computation_layout={(bf16[16,16,256,64]{3,2,1,0}, bf16[16,16,256,64]{3,2,1,0}, bf16[16,16,256,64]{3,2,1,0})->bf16[16,16,256,64]{3,2,1,0}}

ENTRY %a_inference_build_tf_xla_graph_26__XlaMustCompile_true_config_proto_6001324581131673121_executor_type_11160318154034397263_.14 (arg0.1: bf16[16,16,256,64], arg1.2: bf16[16,16,256,64], arg2.3: bf16[16,16,256,64]) -> bf16[16,16,256,64] {
  %arg0.1 = bf16[16,16,256,64]{3,2,1,0} parameter(0), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %arg1.2 = bf16[16,16,256,64]{3,2,1,0} parameter(1), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %dot.7 = bf16[16,16,256,256]{3,2,1,0} dot(bf16[16,16,256,64]{3,2,1,0} %arg0.1, bf16[16,16,256,64]{3,2,1,0} %arg1.2), lhs_batch_dims={0,1}, lhs_contracting_dims={3}, rhs_batch_dims={0,1}, rhs_contracting_dims={3}, metadata={op_type="BatchMatMulV2" op_name="MatMul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %arg2.3 = bf16[16,16,256,64]{3,2,1,0} parameter(2), parameter_replication={false}, metadata={op_name="XLA_Args"}
  ROOT %dot.9 = bf16[16,16,256,64]{3,2,1,0} dot(bf16[16,16,256,256]{3,2,1,0} %dot.7, bf16[16,16,256,64]{3,2,1,0} %arg2.3), lhs_batch_dims={0,1}, lhs_contracting_dims={3}, rhs_batch_dims={0,1}, rhs_contracting_dims={2}, metadata={op_type="BatchMatMulV2" op_name="MatMul_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
}


2023-12-19 16:03:11.236008: I tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter.cc:841] After CudnnFusedMHARewriter:
HloModule a_inference_build_tf_xla_graph_26__XlaMustCompile_true_config_proto_6001324581131673121_executor_type_11160318154034397263_.14, alias_passthrough_params=true, entry_computation_layout={(bf16[16,16,256,64]{3,2,1,0}, bf16[16,16,256,64]{3,2,1,0}, bf16[16,16,256,64]{3,2,1,0})->bf16[16,16,256,64]{3,2,1,0}}

ENTRY %a_inference_build_tf_xla_graph_26__XlaMustCompile_true_config_proto_6001324581131673121_executor_type_11160318154034397263_.14 (arg0.1: bf16[16,16,256,64], arg1.2: bf16[16,16,256,64], arg2.3: bf16[16,16,256,64]) -> bf16[16,16,256,64] {
  %arg0.1 = bf16[16,16,256,64]{3,2,1,0} parameter(0), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %arg1.2 = bf16[16,16,256,64]{3,2,1,0} parameter(1), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %arg2.3 = bf16[16,16,256,64]{3,2,1,0} parameter(2), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %fmha-bmm-bmm = (bf16[16,16,256,64]{3,2,1,0}, u8[0]{0}) custom-call(bf16[16,16,256,64]{3,2,1,0} %arg0.1, bf16[16,16,256,64]{3,2,1,0} %arg1.2, bf16[16,16,256,64]{3,2,1,0} %arg2.3), custom_call_target="__cudnn$fhmaBmmBmm", metadata={op_type="BatchMatMulV2" op_name="MatMul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}, backend_config={"algorithm":{"algo_id":"0","math_type":"TENSOR_OP_MATH","tuning_knobs":{"17":"1","24":"0"},"is_cudnn_frontend":true,"workspace_size":"0"},"fmha_scale":1,"dropout_rate":0,"bmm1_dot_dimension_numbers":{"lhs_contracting_dimensions":["3"],"rhs_contracting_dimensions":["3"],"lhs_batch_dimensions":["0","1"],"rhs_batch_dimensions":["0","1"]},"bmm2_dot_dimension_numbers":{"lhs_contracting_dimensions":["3"],"rhs_contracting_dimensions":["2"],"lhs_batch_dimensions":["0","1"],"rhs_batch_dimensions":["0","1"]},"intermediate_tensor_shape":{"element_type":"BF16","dimensions":["16","16","256","256"],"tuple_shapes":[],"layout":{"dim_level_types":[],"dim_unique":[],"dim_ordered":[],"minor_to_major":["3","2","1","0"],"tiles":[],"element_size_in_bits":"0","memory_space":"0","index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dynamic_shape_metadata_prefix_bytes":"0"},"is_dynamic_dimension":[false,false,false,false]},"seed":"42"}
  ROOT %get-tuple-element = bf16[16,16,256,64]{3,2,1,0} get-tuple-element((bf16[16,16,256,64]{3,2,1,0}, u8[0]{0}) %fmha-bmm-bmm), index=0, metadata={op_type="BatchMatMulV2" op_name="MatMul_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
}


2023-12-19 16:03:11.256692: W tensorflow/core/framework/op_kernel.cc:1839] OP_REQUIRES failed at xla_ops.cc:624 : INTERNAL: Failed to compile XLA Runtime program: failed to run compilation pipeline:
xla.program:5:13: error: failed to legalize operation 'builtin.unrealized_conversion_cast' that was explicitly marked illegal
    %view = memref.view %arg0[%c0][] : memref<8388608xi8> to memref<16x16x256x64xbf16>
            ^
xla.program:5:13: note: see current operation: %70 = "builtin.unrealized_conversion_cast"(%69) : (!llvm.struct<(ptr, ptr, i64, array<4 x i64>, array<4 x i64>)>) -> memref<16x16x256x64xbf16>

Traceback (most recent call last):
  File "/code/hej/deepray/test_tf_vmodule.py", line 56, in <module>
    output = build_tf_xla_graph(queries, keys, values)
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InternalError: Failed to compile XLA Runtime program: failed to run compilation pipeline:
xla.program:5:13: error: failed to legalize operation 'builtin.unrealized_conversion_cast' that was explicitly marked illegal
    %view = memref.view %arg0[%c0][] : memref<8388608xi8> to memref<16x16x256x64xbf16>
            ^
xla.program:5:13: note: see current operation: %70 = "builtin.unrealized_conversion_cast"(%69) : (!llvm.struct<(ptr, ptr, i64, array<4 x i64>, array<4 x i64>)>) -> memref<16x16x256x64xbf16>
 [Op:__inference_build_tf_xla_graph_26]
```
