Flash2 and supports cross attention and dropout

[hanzhi713]

This PR updates pallas implementation to FlashAttention 2. Additionally, we add supports for the following features:

1. q_seq_len != kv_seq_len
2. dropout for triton kernels
   3. cudnn dropout.

Additional changes:

1. Added NoPopDict to workaround a jax bug. See GPU pallas_call loses compiler params during second call when double jit-wrapped jax-ml/jax#25714.
2. Separated dKdV and dQ into two kernels in flash backward. This improved performance by 10~15% when head_dim >= 128. Note that technically fusing dKdVdQ into a single loop and use atomic add for dQ is the fastest solution, but pallas atomics are extremely slow according to empirical testing.

This PR requires Jax >= 0.4.34 and can only be merged after Jax upgrade.

Update: cudnn dropout is disabled in this PR to allow it to be merged before jax upgrade. A follow-up PR will be created to enable cudnn dropout after jax upgrade.

[hanzhi713]

Hi, @ruomingp can you re-approve this PR? I removed cudnn dropout support (but kept the unit tests) so that it can be merged before the jax upgrade and used by our users.