Add more flex attention cases to benchmark.

[chengjunlu]

Add the flex attention shapes which is used by real model to benchmark for tracking performance.

I commented out 4 cases for now for the reason:

1. There is not enough share local memory for the Triton kernel.

• Append shapes of Deepseek-v3 (Nope)
• Decode shapes of Deepseek-v3 (Nope)

2. Flex Attention doesn't support such kind of shapes: Error: LoweringException: ValueError: Number of shared query heads sharing the same KV head must be power of 2.

• Decode shapes of Qwen2-7B

3. Triton kernel block shapes must be power of 2.

• Decode shapes of Phi3-mini-3.8B

We will investigate the first issue on Triton side later.

[chengjunlu]

@liangan1 Please help to comments the expected configuration for the flex attention benchmark.
I will fill up the configuration with reasonable cases from the real use case.

[etiotto]

Note the bmk failure "TypeError: benchmark() got an unexpected keyword argument 'B'"

[chengjunlu]

Only the GEEM computation and inputs are considered for calculating the tflops and gbps.

[chengjunlu]

There is an accuracy issue caused by regression on flex decoding. Need to solve the flex decoding regression issue first.

AssertionError: 
Not equal to tolerance rtol=0.001, atol=0.1

Mismatched elements: 344 / 4096 (8.4%)
Max absolute difference: 1.766
Max relative difference: 0.654
 x: array([[[[ 0.4414  , -0.003366, -0.4326  , ...,  0.0951 ,  0.3289  ,
          -0.614   ]],
...
 y: array([[[[ 1.273  , -0.00971, -1.248  , ...,  0.2744 ,  0.9487 ,
          -1.7705 ]],

[etiotto]

@chengjunlu the test failed due to diff in result. Did you forget to push a local change ?

[whitneywhtsang]

There is an accuracy issue caused by regression on flex decoding. Need to solve the flex decoding regression issue first.

AssertionError: 
Not equal to tolerance rtol=0.001, atol=0.1

Mismatched elements: 344 / 4096 (8.4%)
Max absolute difference: 1.766
Max relative difference: 0.654
 x: array([[[[ 0.4414  , -0.003366, -0.4326  , ...,  0.0951 ,  0.3289  ,
          -0.614   ]],
...
 y: array([[[[ 1.273  , -0.00971, -1.248  , ...,  0.2744 ,  0.9487 ,
          -1.7705 ]],

Similar to other problematic shapes, how about we comment out the decode shape in this PR, and fix it in another PR? (93882ff)

[chengjunlu]

There is an accuracy issue caused by regression on flex decoding. Need to solve the flex decoding regression issue first.

AssertionError: 
Not equal to tolerance rtol=0.001, atol=0.1

Mismatched elements: 344 / 4096 (8.4%)
Max absolute difference: 1.766
Max relative difference: 0.654
 x: array([[[[ 0.4414  , -0.003366, -0.4326  , ...,  0.0951 ,  0.3289  ,
          -0.614   ]],
...
 y: array([[[[ 1.273  , -0.00971, -1.248  , ...,  0.2744 ,  0.9487 ,
          -1.7705 ]],

Similar to other problematic shapes, how about we comment out the decode shape in this PR, and fix it in another PR? (93882ff)

Sounds good to me. Let's add the benchmark first. Let's use other issue to track the decoding regression issue.

[chengjunlu]

There is an accuracy issue caused by regression on flex decoding. Need to solve the flex decoding regression issue first.

AssertionError: 
Not equal to tolerance rtol=0.001, atol=0.1

Mismatched elements: 344 / 4096 (8.4%)
Max absolute difference: 1.766
Max relative difference: 0.654
 x: array([[[[ 0.4414  , -0.003366, -0.4326  , ...,  0.0951 ,  0.3289  ,
          -0.614   ]],
...
 y: array([[[[ 1.273  , -0.00971, -1.248  , ...,  0.2744 ,  0.9487 ,
          -1.7705 ]],

Similar to other problematic shapes, how about we comment out the decode shape in this PR, and fix it in another PR? (93882ff)

The accuracy issue has been fixed by #3999.

I will rebase this PR after the #3999 merged.