[Question] How can I make sure Tensorflow has already used cuDNN Flash Fused Attention?
Hi @MoFHeka ,

To enable XLA we need to decorate respective function with `@tf.function(jit_compile=True)`. There is this method called  [experimental_get_compiler_ir](https://github.com/tensorflow/tensorflow/blob/51cda66dedc2db71709efc6e68998be53c94a533/tensorflow/python/types/core.py#L251) which will give IR for of the input code. Could you check this whether useful for what you are looking for ? @SuryanarayanaY Here is the IR, I can't see any fused kernel but only matmul
```bash
HloModule a_inference_build_tf_xla_graph_104__.107, entry_computation_layout={()->f16[4096,2048,8,16]{3,2,1,0}}

%region_0.57 (Arg_0.58: f16[], Arg_1.59: f16[]) -> f16[] {
  %Arg_0.58 = f16[] parameter(0)
  %Arg_1.59 = f16[] parameter(1)
  ROOT %maximum.60 = f16[] maximum(f16[] %Arg_0.58, f16[] %Arg_1.59), metadata={op_type="Softmax" op_name="attention/Softmax"}
}

%region_1.69 (Arg_0.70: f32[], Arg_1.71: f32[]) -> f32[] {
  %Arg_0.70 = f32[] parameter(0)
  %Arg_1.71 = f32[] parameter(1)
  ROOT %add.72 = f32[] add(f32[] %Arg_0.70, f32[] %Arg_1.71), metadata={op_type="Softmax" op_name="attention/Softmax"}
}

ENTRY %a_inference_build_tf_xla_graph_104__.107 () -> f16[4096,2048,8,16] {
  %constant.3 = s32[4]{0} constant({4096, 2048, 8, 8}), metadata={op_type="StridedSlice" op_name="attention/strided_slice" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.1 = s32[4]{0} constant({4096, 2048, 8, 8}), metadata={op_type="StridedSlice" op_name="attention/strided_slice" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.2 = s32[2]{0} slice(s32[4]{0} %constant.1), slice={[0:2]}, metadata={op_type="StridedSlice" op_name="attention/strided_slice" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.4 = s32[2]{0} reshape(s32[2]{0} %slice.2), metadata={op_type="StridedSlice" op_name="attention/strided_slice" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.8 = s32[4]{0} constant({4096, 2048, 8, 16}), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.9 = s32[1]{0} slice(s32[4]{0} %constant.8), slice={[0:1]}, metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.10 = s32[] reshape(s32[1]{0} %slice.9), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.11 = s32[] convert(s32[] %reshape.10), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.12 = s32[1]{0} slice(s32[4]{0} %constant.8), slice={[1:2]}, metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.13 = s32[] reshape(s32[1]{0} %slice.12), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.14 = s32[] convert(s32[] %reshape.13), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.15 = s32[1]{0} slice(s32[4]{0} %constant.8), slice={[2:3]}, metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.16 = s32[] reshape(s32[1]{0} %slice.15), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.17 = s32[] convert(s32[] %reshape.16), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.18 = s32[1]{0} slice(s32[4]{0} %constant.8), slice={[3:4]}, metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.19 = s32[] reshape(s32[1]{0} %slice.18), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.20 = s32[] convert(s32[] %reshape.19), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.30 = s32[4]{0} constant({4096, 2048, 8, 16}), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.31 = s32[1]{0} slice(s32[4]{0} %constant.30), slice={[0:1]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.32 = s32[] reshape(s32[1]{0} %slice.31), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.33 = s32[] convert(s32[] %reshape.32), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.34 = s32[1]{0} slice(s32[4]{0} %constant.30), slice={[1:2]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.35 = s32[] reshape(s32[1]{0} %slice.34), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.36 = s32[] convert(s32[] %reshape.35), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.37 = s32[1]{0} slice(s32[4]{0} %constant.30), slice={[2:3]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.38 = s32[] reshape(s32[1]{0} %slice.37), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.39 = s32[] convert(s32[] %reshape.38), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.40 = s32[1]{0} slice(s32[4]{0} %constant.30), slice={[3:4]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.41 = s32[] reshape(s32[1]{0} %slice.40), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.42 = s32[] convert(s32[] %reshape.41), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.83 = s32[4]{0} constant({4096, 2048, 8, 16}), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.84 = s32[1]{0} slice(s32[4]{0} %constant.83), slice={[0:1]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.85 = s32[] reshape(s32[1]{0} %slice.84), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.86 = s32[] convert(s32[] %reshape.85), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.87 = s32[1]{0} slice(s32[4]{0} %constant.83), slice={[1:2]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.88 = s32[] reshape(s32[1]{0} %slice.87), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.89 = s32[] convert(s32[] %reshape.88), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.90 = s32[1]{0} slice(s32[4]{0} %constant.83), slice={[2:3]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.91 = s32[] reshape(s32[1]{0} %slice.90), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.92 = s32[] convert(s32[] %reshape.91), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %slice.93 = s32[1]{0} slice(s32[4]{0} %constant.83), slice={[3:4]}, metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.94 = s32[] reshape(s32[1]{0} %slice.93), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %convert.95 = s32[] convert(s32[] %reshape.94), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.6 = f16[] constant(0), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.5 = f16[] constant(1), metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %rng.7 = f16[4096,2048,8,16]{3,2,1,0} rng(f16[] %constant.6, f16[] %constant.5), distribution=rng_normal, metadata={op_type="RandomStandardNormal" op_name="random_normal/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.21 = f16[] constant(1), metadata={op_type="Mul" op_name="random_normal/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.22 = f16[4096,2048,8,16]{3,2,1,0} broadcast(f16[] %constant.21), dimensions={}, metadata={op_type="Mul" op_name="random_normal/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %multiply.23 = f16[4096,2048,8,16]{3,2,1,0} multiply(f16[4096,2048,8,16]{3,2,1,0} %rng.7, f16[4096,2048,8,16]{3,2,1,0} %broadcast.22), metadata={op_type="Mul" op_name="random_normal/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.24 = f16[] constant(0), metadata={op_type="AddV2" op_name="random_normal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.25 = f16[4096,2048,8,16]{3,2,1,0} broadcast(f16[] %constant.24), dimensions={}, metadata={op_type="AddV2" op_name="random_normal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %add.26 = f16[4096,2048,8,16]{3,2,1,0} add(f16[4096,2048,8,16]{3,2,1,0} %multiply.23, f16[4096,2048,8,16]{3,2,1,0} %broadcast.25), metadata={op_type="AddV2" op_name="random_normal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.28 = f16[] constant(0), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.27 = f16[] constant(1), metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %rng.29 = f16[4096,2048,8,16]{3,2,1,0} rng(f16[] %constant.28, f16[] %constant.27), distribution=rng_normal, metadata={op_type="RandomStandardNormal" op_name="random_normal_1/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.43 = f16[] constant(1), metadata={op_type="Mul" op_name="random_normal_1/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.44 = f16[4096,2048,8,16]{3,2,1,0} broadcast(f16[] %constant.43), dimensions={}, metadata={op_type="Mul" op_name="random_normal_1/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %multiply.45 = f16[4096,2048,8,16]{3,2,1,0} multiply(f16[4096,2048,8,16]{3,2,1,0} %rng.29, f16[4096,2048,8,16]{3,2,1,0} %broadcast.44), metadata={op_type="Mul" op_name="random_normal_1/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.46 = f16[] constant(0), metadata={op_type="AddV2" op_name="random_normal_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.47 = f16[4096,2048,8,16]{3,2,1,0} broadcast(f16[] %constant.46), dimensions={}, metadata={op_type="AddV2" op_name="random_normal_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %add.48 = f16[4096,2048,8,16]{3,2,1,0} add(f16[4096,2048,8,16]{3,2,1,0} %multiply.45, f16[4096,2048,8,16]{3,2,1,0} %broadcast.47), metadata={op_type="AddV2" op_name="random_normal_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %dot.49 = f16[4096,2048,8,8]{3,2,1,0} dot(f16[4096,2048,8,16]{3,2,1,0} %add.26, f16[4096,2048,8,16]{3,2,1,0} %add.48), lhs_batch_dims={0,1}, lhs_contracting_dims={3}, rhs_batch_dims={0,1}, rhs_contracting_dims={3}, metadata={op_type="BatchMatMulV2" op_name="attention/MatMul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %transpose.50 = f16[4096,2048,8,8]{3,2,1,0} transpose(f16[4096,2048,8,8]{3,2,1,0} %dot.49), dimensions={0,1,2,3}, metadata={op_type="BatchMatMulV2" op_name="attention/MatMul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.51 = f16[1,1,8,8]{3,2,1,0} constant({...}), metadata={op_type="Sub" op_name="attention/sub" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.52 = f16[8,8]{1,0} reshape(f16[1,1,8,8]{3,2,1,0} %constant.51), metadata={op_type="Sub" op_name="attention/sub" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.53 = f16[4096,2048,8,8]{3,2,1,0} broadcast(f16[8,8]{1,0} %reshape.52), dimensions={2,3}, metadata={op_type="Sub" op_name="attention/sub" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %subtract.54 = f16[4096,2048,8,8]{3,2,1,0} subtract(f16[4096,2048,8,8]{3,2,1,0} %transpose.50, f16[4096,2048,8,8]{3,2,1,0} %broadcast.53), metadata={op_type="Sub" op_name="attention/sub" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.56 = f16[] constant(-inf), metadata={op_type="Softmax" op_name="attention/Softmax" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reduce.61 = f16[4096,2048,8]{2,1,0} reduce(f16[4096,2048,8,8]{3,2,1,0} %subtract.54, f16[] %constant.56), dimensions={3}, to_apply=%region_0.57, metadata={op_type="Softmax" op_name="attention/Softmax"}
  %reshape.62 = f16[4096,2048,8,1]{3,2,1,0} reshape(f16[4096,2048,8]{2,1,0} %reduce.61), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %broadcast.63 = f16[4096,2048,8,1]{3,2,1,0} broadcast(f16[4096,2048,8,1]{3,2,1,0} %reshape.62), dimensions={0,1,2,3}, metadata={op_type="Softmax" op_name="attention/Softmax"}
  %reshape.64 = f16[4096,2048,8]{2,1,0} reshape(f16[4096,2048,8,1]{3,2,1,0} %broadcast.63), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %broadcast.65 = f16[4096,2048,8,8]{3,2,1,0} broadcast(f16[4096,2048,8]{2,1,0} %reshape.64), dimensions={0,1,2}, metadata={op_type="Softmax" op_name="attention/Softmax"}
  %subtract.66 = f16[4096,2048,8,8]{3,2,1,0} subtract(f16[4096,2048,8,8]{3,2,1,0} %subtract.54, f16[4096,2048,8,8]{3,2,1,0} %broadcast.65), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %exponential.67 = f16[4096,2048,8,8]{3,2,1,0} exponential(f16[4096,2048,8,8]{3,2,1,0} %subtract.66), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %convert.68 = f32[4096,2048,8,8]{3,2,1,0} convert(f16[4096,2048,8,8]{3,2,1,0} %exponential.67), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %constant.55 = f32[] constant(-0), metadata={op_type="Softmax" op_name="attention/Softmax" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reduce.73 = f32[4096,2048,8]{2,1,0} reduce(f32[4096,2048,8,8]{3,2,1,0} %convert.68, f32[] %constant.55), dimensions={3}, to_apply=%region_1.69, metadata={op_type="Softmax" op_name="attention/Softmax"}
  %convert.74 = f16[4096,2048,8]{2,1,0} convert(f32[4096,2048,8]{2,1,0} %reduce.73), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %reshape.75 = f16[4096,2048,8,1]{3,2,1,0} reshape(f16[4096,2048,8]{2,1,0} %convert.74), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %broadcast.76 = f16[4096,2048,8,1]{3,2,1,0} broadcast(f16[4096,2048,8,1]{3,2,1,0} %reshape.75), dimensions={0,1,2,3}, metadata={op_type="Softmax" op_name="attention/Softmax"}
  %reshape.77 = f16[4096,2048,8]{2,1,0} reshape(f16[4096,2048,8,1]{3,2,1,0} %broadcast.76), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %broadcast.78 = f16[4096,2048,8,8]{3,2,1,0} broadcast(f16[4096,2048,8]{2,1,0} %reshape.77), dimensions={0,1,2}, metadata={op_type="Softmax" op_name="attention/Softmax"}
  %divide.79 = f16[4096,2048,8,8]{3,2,1,0} divide(f16[4096,2048,8,8]{3,2,1,0} %exponential.67, f16[4096,2048,8,8]{3,2,1,0} %broadcast.78), metadata={op_type="Softmax" op_name="attention/Softmax"}
  %constant.81 = f16[] constant(0), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.80 = f16[] constant(1), metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %rng.82 = f16[4096,2048,8,16]{3,2,1,0} rng(f16[] %constant.81, f16[] %constant.80), distribution=rng_normal, metadata={op_type="RandomStandardNormal" op_name="random_normal_2/RandomStandardNormal" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.96 = f16[] constant(1), metadata={op_type="Mul" op_name="random_normal_2/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.97 = f16[4096,2048,8,16]{3,2,1,0} broadcast(f16[] %constant.96), dimensions={}, metadata={op_type="Mul" op_name="random_normal_2/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %multiply.98 = f16[4096,2048,8,16]{3,2,1,0} multiply(f16[4096,2048,8,16]{3,2,1,0} %rng.82, f16[4096,2048,8,16]{3,2,1,0} %broadcast.97), metadata={op_type="Mul" op_name="random_normal_2/mul" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %constant.99 = f16[] constant(0), metadata={op_type="AddV2" op_name="random_normal_2" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %broadcast.100 = f16[4096,2048,8,16]{3,2,1,0} broadcast(f16[] %constant.99), dimensions={}, metadata={op_type="AddV2" op_name="random_normal_2" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %add.101 = f16[4096,2048,8,16]{3,2,1,0} add(f16[4096,2048,8,16]{3,2,1,0} %multiply.98, f16[4096,2048,8,16]{3,2,1,0} %broadcast.100), metadata={op_type="AddV2" op_name="random_normal_2" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %dot.102 = f16[4096,2048,8,16]{3,2,1,0} dot(f16[4096,2048,8,8]{3,2,1,0} %divide.79, f16[4096,2048,8,16]{3,2,1,0} %add.101), lhs_batch_dims={0,1}, lhs_contracting_dims={3}, rhs_batch_dims={0,1}, rhs_contracting_dims={2}, metadata={op_type="BatchMatMulV2" op_name="attention/MatMul_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %transpose.103 = f16[4096,2048,8,16]{3,2,1,0} transpose(f16[4096,2048,8,16]{3,2,1,0} %dot.102), dimensions={0,1,2,3}, metadata={op_type="BatchMatMulV2" op_name="attention/MatMul_1" source_file="/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py" source_line=1161}
  %reshape.104 = f16[4096,2048,8,16]{3,2,1,0} reshape(f16[4096,2048,8,16]{3,2,1,0} %transpose.103), metadata={op_name="XLA_Retvals"}
  %tuple.105 = (f16[4096,2048,8,16]{3,2,1,0}) tuple(f16[4096,2048,8,16]{3,2,1,0} %reshape.104), metadata={op_name="XLA_Retvals"}
  ROOT %get-tuple-element.106 = f16[4096,2048,8,16]{3,2,1,0} get-tuple-element((f16[4096,2048,8,16]{3,2,1,0}) %tuple.105), index=0, metadata={op_name="XLA_Retvals"}
}
``` @SuryanarayanaY @AyanmoI @Tixxx  Here is a little "fmha-bmm-bmm" test, which I enable VLOG in cudnn_fused_mha_rewriter. But still, there is nothing log about cudnn_fused_mha.
```python
import os
os.environ['TF_CPP_VMODULE'] = 'cudnn_fused_mha_rewriter=2'

import tensorflow as tf
import numpy as np


batch_size = 4096
seq_len_q = 2048
num_heads = 8
head_dim = 16
mat_dtypes = tf.float16


# @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)
``` Enable XLA CudnnFusedMHARewriter successfully by export XLA_FLAGS=--xla_gpu_enable_cudnn_fmha, but trigger new bugs. #62666  Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62661">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62661">No</a>

### Issue type

Support

### 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

### 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?

In the [commit](https://github.com/tensorflow/tensorflow/commit/81a1946366e57edd1029153caa8f826c19717b75), we can see Nvidia people have already added the XLA supporting to convert separated kernel into one fused kernel when calculating attention(mha). 
So how can I make sure it really happens. For example when I use tf.keras.layers.Attention.

### Standalone code to reproduce the issue

```shell
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmark for FlashAttn and grad of FlashAttn."""

import numpy as np
import math

from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session as session_lib
from tensorflow.python.framework import config as framework_config
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import variable_scope
from tensorflow.python.platform import benchmark
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging

import tensorflow as tf


def build_origin_graph(batch_size, seq_len_q, num_heads, head_dim, dtype=dtypes.float32):
  """Build a graph containing a sequence of FlashAttn operations.

  Args:

  Returns:
    
  """
  seq_len_v = seq_len_q
  queries = tf.random.normal(
      (batch_size, seq_len_q, num_heads, head_dim), dtype=dtype)
  keys = tf.random.normal(
      (batch_size, seq_len_v, num_heads, head_dim), dtype=dtype)
  values = tf.random.normal(
      (batch_size, seq_len_v, num_heads, head_dim), dtype=dtype)
  scores = tf.einsum('bthd,bshd->bhts', queries / math.sqrt(queries.shape[-1]), keys)
  attention = tf.nn.softmax(scores)
  outputs = tf.einsum('bhts,bshd->bthd', attention, values)
  return outputs


def build_tf_plus_graph(batch_size, seq_len_q, num_heads, head_dim, dtype=dtypes.float32):
  """Build a graph containing a sequence of FlashAttn operations.

  Args:

  Returns:
    
  """
  seq_len_v = seq_len_q
  queries = tf.random.normal(
      (batch_size, seq_len_q, num_heads, head_dim), dtype=dtype)
  keys = tf.random.normal(
      (batch_size, seq_len_v, num_heads, head_dim), dtype=dtype)
  values = tf.random.normal(
      (batch_size, seq_len_v, num_heads, head_dim), dtype=dtype)
  flash_attention_layer = FlashAttentionLayer(
      seq_len_q, seq_len_v, num_heads, head_dim)
  flash_attention_output = flash_attention_layer(
      queries, keys, values)
  return flash_attention_output

@tf.function
def build_tf_xla_graph(batch_size, seq_len_q, num_heads, head_dim, dtype=dtypes.float32):
  """Build a graph containing a sequence of FlashAttn operations.

  Args:

  Returns:
    
  """
  seq_len_v = seq_len_q
  with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
    queries = tf.random.normal(
        (batch_size, seq_len_q, num_heads, head_dim), dtype=dtype)
    keys = tf.random.normal(
        (batch_size, seq_len_v, num_heads, head_dim), dtype=dtype)
    values = tf.random.normal(
        (batch_size, seq_len_v, num_heads, head_dim), dtype=dtype)
    keras_attention_layer = tf.keras.layers.Attention(use_scale=False)
    flash_attention_output = keras_attention_layer(inputs=[queries, values, keys])
  return flash_attention_output


class OriginAttnBenchmark(test.Benchmark):
  """Benchmark FlashAttn!"""
  _build_graph_fn = staticmethod(build_origin_graph)

  def _run_graph(self, batch_size, seq_len_q, num_heads, head_dim, dtype):
    """Run the graph and print its execution time.

    Args:

    Returns:
      
    """
    graph = ops.Graph()
    with graph.as_default():
      outputs = self._build_graph_fn(batch_size, seq_len_q, num_heads, head_dim, dtype)
    config = config_pb2.ConfigProto(graph_options=config_pb2.GraphOptions(
        optimizer_options=config_pb2.OptimizerOptions(
            opt_level=config_pb2.OptimizerOptions.ON_2)))
    with session_lib.Session(graph=graph, config=config) as session:
      logging.set_verbosity(0)
      variables.global_variables_initializer().run()
      bench = benchmark.TensorFlowBenchmark()
      attention_matrix_computation = batch_size * 2 * seq_len_q * seq_len_q * num_heads * head_dim
      attention_over_values = batch_size * 2 * seq_len_q * seq_len_q * num_heads * head_dim
      dtype_byte = dtype.size
      bench.run_op_benchmark(
          session,
          outputs,
          mbs=(attention_matrix_computation + attention_over_values) * dtype_byte * 100 / 1e6,
          store_trace=True,
          extras={
              "batch_size": batch_size,
              "seq_len_q": seq_len_q,
              "num_heads": num_heads,
              "head_dim": batch_size,
              "dtype": seq_len_q,
          })

  def benchmark_FlashAttn(self):
    print("Attention Forward")
    # batch_size = [4096, 8192]
    # seq_len_q = [1, 7, 30]
    # num_heads = [2, 4, 8]
    # head_dim = [128, 512, 1024]
    # dtypes = [dtypes.float32, dtypes.float16]
    batch_size = [4096]
    seq_len_q = [30]
    num_heads = [4]
    head_dim = [128]
    mat_dtypes = [dtypes.float16]
    for b in batch_size:
      for s in seq_len_q:
        for h in num_heads:
          for d in head_dim:
            for dtype in mat_dtypes:
              self._run_graph(b, s, h, d, dtype)


class TFXLAAttnBenchmark(OriginAttnBenchmark):
  _build_graph_fn = staticmethod(build_tf_xla_graph)


if __name__ == "__main__":
  physical_devices = framework_config.list_physical_devices('GPU')
  if len(physical_devices):
    framework_config.set_visible_devices(physical_devices[0], 'GPU')
  test.main()
```


### Relevant log output

```shell
2023-12-18 16:20:00.138356: I tensorflow/core/util/port.cc:111] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-12-18 16:20:00.180184: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9360] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-12-18 16:20:00.180220: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-12-18 16:20:00.180265: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1537] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-12-18 16:20:00.188748: I tensorflow/core/platform/cpu_feature_guard.cc:183] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Attention Forward
2023-12-18 16:20:02.332430: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x563f6ff5d7c0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2023-12-18 16:20:02.332459: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2023-12-18 16:20:02.332501: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-18 16:20:02.466545: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1883] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 79005 MB memory:  -> device: 0, name: NVIDIA A800-SXM4-80GB, pci bus id: 0000:d0:00.0, compute capability: 8.0
2023-12-18 16:20:02.478813: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x563f7158d210 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2023-12-18 16:20:02.478843: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): NVIDIA A800-SXM4-80GB, Compute Capability 8.0
2023-12-18 16:20:02.526593: E ./tensorflow/compiler/xla/stream_executor/stream_executor_internal.h:124] SetPriority unimplemented for this stream.
2023-12-18 16:20:02.532643: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1883] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 79005 MB memory:  -> device: 0, name: NVIDIA A800-SXM4-80GB, pci bus id: 0000:d0:00.0, compute capability: 8.0
2023-12-18 16:20:02.542548: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:382] MLIR V1 optimization pass is not enabled
2023-12-18 16:20:02.628241: I tensorflow/compiler/jit/xla_device.cc:460] XLA_GPU and XLA_CPU devices are deprecated and will be removed in subsequent releases. Instead, use either @tf.function(jit_compile=True) for must-compile semantics, or run with TF_XLA_FLAGS=--tf_xla_auto_jit=2 for auto-clustering best-effort compilation.
2023-12-18 16:20:02.692172: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2023-12-18 16:20:02.746939: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:442] Loaded cuDNN version 8906
2023-12-18 16:20:05.145859: I ./tensorflow/compiler/jit/device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2023-12-18 16:20:06.211684: I tensorflow/tsl/profiler/lib/profiler_session.cc:104] Profiler session initializing.
2023-12-18 16:20:06.211730: I tensorflow/tsl/profiler/lib/profiler_session.cc:119] Profiler session started.
2023-12-18 16:20:06.211783: I tensorflow/compiler/xla/backends/profiler/gpu/cupti_tracer.cc:1694] Profiler found 1 GPUs
2023-12-18 16:20:06.320320: I tensorflow/tsl/profiler/lib/profiler_session.cc:70] Profiler session collecting data.
2023-12-18 16:20:06.322079: I tensorflow/compiler/xla/backends/profiler/gpu/cupti_tracer.cc:1828] CUPTI activity buffer flushed
2023-12-18 16:20:06.347580: I tensorflow/compiler/xla/backends/profiler/gpu/cupti_collector.cc:541]  GpuTracer has collected 24 callback api events and 17 activity events.
2023-12-18 16:20:06.348167: I tensorflow/tsl/profiler/lib/profiler_session.cc:131] Profiler session tear down.
```
