fix flex attention

Author: xuzhao9

benchmarks/nightly/autogen.yaml

@@ -155,3 +155,7 @@ vector_exp_bwd:
 welford_fwd:
   args: --op welford --baseline eager_layer_norm --metrics latency,speedup --only
     test_no_welford,triton_welford,eager_layer_norm
+bf16_flex_attention_fwd:
+  args: --op flex_attention --metrics latency,tflops --only compiled
+bf16_flex_attention_bwd:
+  args: --op flex_attention --metrics latency,tflops --only compiled

benchmarks/nightly/manual.yaml

@@ -22,3 +22,7 @@ extra_args:
   # flash_attention triton_tutorial_flash_v2 impl only supports causal in backward
   bf16_flash_attention_bwd:
     args: --op flash_attention --baseline flash_v3 --metrics tflops,speedup --bwd --only triton_tutorial_flash_v2,flash_v3 --causal
+  bf16_flex_attention_fwd:
+    args: --op flex_attention --metrics latency,tflops --only compiled
+  bf16_flex_attention_bwd:
+    args: --op flex_attention --metrics latency,tflops --only compiled