[Feature] Add attention mask support

[ShunyuYao]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/mit-han-lab/nunchaku/discussions/new/choose. Otherwise, it will be closed.
• 2. I will do my best to describe the issue in English.

Motivation

Is it possible to support attention mask in the attention process? attention mask is now widely supported in diffusers and comfyui and can be really helpful in lots of circumustances.

Related resources

In the implentation of attention process in Diffuerses, scaled_dot_product_attention function support attention mask option.
def scaled_dot_product_attention(query: Tensor, key: Tensor, value: Tensor, attn_mask: Optional[Tensor] = None, dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False) -> Tensor: ...

[ShunyuYao]

The attention mask is actually add the mask tensor to the results after matrix caculation of Q and K, If the mask tensor is bool type, it will set the False values to negative inf and set the True values to zero and then do the addition caculation.

[felixslu]

Our product also need this feature. For examples:
Image to Image Flux pipeline, the condition branch may use a special scale to enhance image control factor, then it will product a special mask(not causual mask, not Identify mask, in softmax(q*Kt + logM) formula).
Nunchaku currently not support self-defined attn-mask, It's really requirement for us.

[gonggp]

attention mask is a commonly used feature, and hope it can be supported.

[github-actions]

This issue has been automatically closed due to 30-day inactivity. Please feel free to reopen it with /reopen if needed.