@triton.autotune(configs=[triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64,
    'waves_per_eu': 0, 'pre_load_v': True}, num_stages=1, num_warps=4),
    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 1,
    'pre_load_v': True}, num_stages=1, num_warps=4), triton.Config({
    'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 2, 'pre_load_v': True},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    64, 'waves_per_eu': 3, 'pre_load_v': True}, num_stages=1, num_warps=4),
    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 4,
    'pre_load_v': True}, num_stages=1, num_warps=4), triton.Config({
    'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 0, 'pre_load_v': False},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    64, 'waves_per_eu': 1, 'pre_load_v': False}, num_stages=1, num_warps=4),
    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 2,
    'pre_load_v': False}, num_stages=1, num_warps=4), triton.Config({
    'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'pre_load_v': False},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    64, 'waves_per_eu': 4, 'pre_load_v': False}, num_stages=1, num_warps=4)
    ], key=['seqlen_q', 'seqlen_k', 'STAGE'])
@triton.jit
def tuned_attn_fwd(Q, K, V, sm_scale, M, Out, stride_qz, stride_qh,
    stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh,
    stride_om, stride_on, seqlen_q, seqlen_k, dropout_p, philox_seed,
    philox_offset_base, encoded_softmax, STAGE: tl.constexpr, BLOCK_M: tl.
    constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
    pre_load_v: tl.constexpr, ENABLE_DROPOUT: tl.constexpr,
    RETURN_ENCODED_SOFTMAX: tl.constexpr):
    bare_attn_fwd(Q, K, V, sm_scale, M, Out, stride_qz, stride_qh,
        stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk,
        stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh,
        stride_om, stride_on, seqlen_q, seqlen_k, dropout_p, philox_seed,
        philox_offset_base, encoded_softmax, STAGE, BLOCK_M, BLOCK_DMODEL,
        BLOCK_N, pre_load_v, ENABLE_DROPOUT, RETURN_ENCODED_SOFTMAX)
