align the inputs with flash attention
Summary:
Each run will go through a range of N_CTX and we will have D_HEAD set via arg.
Also add an additional key for autotuning.
Reviewed By: xuzhao9, sijiac
Differential Revision: D57979174
fbshipit-source-id: e415acdea7f5f912574dc227a0e52ff13bf2b261
