'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
# Preserve all cache keys without interpreting them as automatic tiling axes.
@triton.runtime.autotuner.autotune(
configs=get_autotune_config(multibuffer_list=(True, False)),
key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'IS_VARLEN'],
)
@triton.jit(do_not_specialize=['T'])