[BUG] Attention mask in example is incorrect

[yiyousong]

Currently, in the flexattention variant/modeling_doge, the "KV Score" or 'zoh' is based on the whole sequence instead of causal sequence.
https://github.com/SmallDoges/flash-dmattn/blob/801e8161e30eae64d49da1fb9d5139f14bb5aa2f/examples/modeling/modeling_doge.py#L326
attn_bias = dt_states[:, :, None, :].expand( -1, -1, hidden_states.shape[1], -1 ) # [batch_size, num_heads, query_len, key_len]
the causal mask is missing here (after this step). Therefore, the topK KV-selection is the same across all queries. This causes reduced attention range for earlier tokens.

I am unfamiliar with cuda, so I don't know if this issue persists in other versions.

[yiyousong]

BTW，the flex_attention_forward https://github.com/SmallDoges/flash-dmattn/blob/801e8161e30eae64d49da1fb9d5139f14bb5aa2f/examples/modeling/modeling_doge.py#L186 seems to use the wrong mask format (1,0) instead of (0,-inf) for addition mask in FlexAttention's score_mod

[yiyousong]

Nevermind, the whole modeling_doge example is full of bugs, I will report a summary when I'am done.

[LoserCheems]

My bad. After completing the utils integration, I forgot to update the example modeling.