Add key to FA Autotuner (#4199)

Add "STAGE" as a key for the autotuner to optimize the configuration for
CAUSAL FA as well.

---------

Signed-off-by: Maxime France-Pillois <maxime.francepillois@codeplay.com>

Author: Maxime France-Pillois

benchmarks/triton_kernels_benchmark/flash_attention_benchmark.py

@@ -160,12 +160,12 @@ def _attn_fwd_with_block_pointers(Q, K, V, sm_scale, M, Out,  #
 configs = [
     triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN, 'grf_mode': 'large', 'one_matrix_per_load_for_bt': True}, num_stages=s, num_warps=w) \
     for BM in [128, 256] \
-    for BN in [32, 64, 128] \
+    for BN in [32, 64] \
     for s in [2, 3, 4] \
     for w in [8, 16, 32] \
     ]
 
-tuner = triton.autotune(configs, key=['N_CTX', 'BLOCK_DMODEL'])
+tuner = triton.autotune(configs, key=['N_CTX', 'BLOCK_DMODEL', 'STAGE'])
 
 
 @triton.jit