Support head dim 256

[woct0rdho]

Head dim 256 is needed for models like AuraFlow and Ideogram 4. This PR adds it in the CUDA kernels. Later we can also add it in the Triton kernels if needed.

I've tested on RTX 3080 (sm86) and RTX 4090 (sm89). When seq len is large, they reach the ideal 2x speedup compared to FlashAttention.

As always, pv_accum_dtype="fp16+fp32" (or "fp32+fp16" in some interfaces) is faster than pv_accum_dtype="fp32", and pv_accum_dtype="fp16" is even faster, but more likely to cause black/noise/degraded output. Maybe the time has come that we have to tune it for each model.

Currently sm90 is not supported due to constraint of WGMMA size. If you have sm90, you can help make it work.

Implementation details

When head_dim == 256, in transpose_pad_permute_cuda, we reduce CTA_SIZE on device to 32 so the assert CTA_SIZE * HEAD_DIM <= 8192 (1024 threads per block) is satisfied. In all attn kernels, we reduce WARP_Q to 16 because of register pressure.

By the way, I've applied a patch although it's not related to head dim 256: In sm80 attn kernels, when head_dim == 128 and pv_accum_dtype == "fp16+fp32" and non-causal, we should set (CTA_Q, CTQ_K, WARP_Q, WARP_K) to (128, 32, 32, 32), because of how the inst buf works. This gives 2x or even 3x speedup and actually makes pv_accum_dtype="fp16+fp32" usable. I found it when benchmarking all configs with common ComfyUI workloads on sm86. In summary, the benchmarked workloads are:

batch_size = 1
num_heads = [16, 32]
seq_len = [1024, 2048, 4096, 8192]
head_dim = [64, 128, 256]
pv_accum_dtype = [fp16, fp16+fp32, fp32]

And the heuristic configs (CTA_Q, CTQ_K, WARP_Q, WARP_K) on sm86 are:

pv_accum_dtype=fp16 or fp32:
head_dim= 64: (128, 64, 32, 64)
head_dim=128: (128, 64, 32, 64)
head_dim=256: (128, 64, 16, 64)

pv_accum_dtype=fp16+fp32:
head_dim= 64: (128, 64, 32, 64)
head_dim=128: (128, 64, 16, 64) if causal else (128, 32, 32, 32)
head_dim=256: (128, 64, 16, 64)

Maybe there are similar cases in the seven sm89 attn kernels but I haven't looked.