Don't disable uneven k to support more headdims

[njhill]

Originally disabled by @WoosukKwon in eee8e47

[WoosukKwon]

Adding the flag increased the wheel size of vllm-flash-attn from 107 MB to 143MB (not sure whether this is before compression).

[njhill]

Thanks @WoosukKwon, I guess that's fairly significant.

The motivation here is that we depend of flash attention currently for certain features like multi-step scheduling, so it's a big performance downside for models which have these other head sizes (in particular some of the IBM Granite models use head_dim 80)

[WoosukKwon]

@njhill We can just add the head size 80 to

flash-attention/csrc/flash_attn/src/generate_kernels.py

Line 18 in 344c988

HEAD_DIMENSIONS = [32, 64, 96, 128, 160, 192, 224, 256]

Would that work for you?

[njhill]

@WoosukKwon yes, I think for our immediate needs that would be great, if you're sure that will be sufficient!

[njhill]

@WoosukKwon I've opened #22 for this

[WoosukKwon]

@njhill Since the idea above doesn't work, can you check how much this PR affects the vllm wheel size?