As mentioned, custom kernels provide a wider range of support for execution scenarios To ensure efficient execution (e,g., to use GPU tensor cores), model configurations need to meet a small number of requirements.  This list of requirements will evolve over time, prospectively relaxing constraints limiting the usage of currently supported custom kernels, or providing additional kernels in the future.
For the most up to date list of custom kernels and dispatch constraints, you can refer to sdp_utils.h.  As of PyTorch 2.0, the existing fused SDPA kernels have the following constraints:

Flash Attention only supports 16 bit floating point data types (float16 and bfloat16).
The head dimension must be a multiple of 8 for 16-bit floating point numbers and a multiple of 4 for 32-bit floating point numbers. At present, the maximum head_dim support for the Flash Attention custom kernel is 128.
