Support 256 Head Dim

[lin-ht]

Hi,
Is head dim 256 supported in flexible_flash_attention/magi_attention?

I noticed that get_max_headdim() could be 256 if FLASHATTENTION_DISABLE_HDIM256 is not defined as shown here.
However, when I actually run the attention with 256 head dim inputs, I hit this sanity_check failure (code here):
assert head_dim <= 128, "head_dim must be <= 128 for now"

# Sample code:
out_, lse_ = flex_flash_attn_func(
    q_[i],
    k_[i],
    v_[i],
    q_ranges,
    k_ranges,
    attn_type_map=attn_type_map,
    softmax_scale=softmax_scale,
    disable_fwd_atomic_reduction=False,
)

[Strivin0311]

Hi~For now, we only support head_dim in {64, 128}. And we can easily extend to {64, 128, 192}, but it's non-trivial to scale to 256 or 512 for FFA.

If you have a compelling and general use-case, you are very welcome to share the details with us and we’ll look into it.

[lin-ht]

Thanks for the quick reply. We are working on a video generation model that demands 256 head dim. We'd like to try sparse attention through magi_attention.
Could you elaborate on the non-trivial issue you mentioned above so that we have a better understanding on the complexity of the issue? May I know if your team has any plan to upgrade the support to 256? Thank you~

[littsk]

Apologies for the delay in responding.

We currently do not support headdim=256 due to excessive register pressure on Hopper architectures. Specifically, Tensor Core outputs for matrix multiplications must be stored in registers. With headdim=256 and a minimum q_block=64, the output shape causes significant register spilling, which severely degrades performance. Supporting this configuration would require a dedicated algorithm/kernel implementation, which is not currently on our roadmap.

Regarding sparse attention, if you are using the flex_flash_attn_func API, we highly recommend enabling the auto_range_merge flag to significantly boost performance. Furthermore, we are actively optimizing sparse attention and enhancing its native distributed capabilities.