How to get attention score? "return_attn_probs=True" is not work.

[UnableToUseGit]

version: flash-attn 2.7.0.post2

When I use the function "flash_attn_func" and I set return_attn_probs=True, it not returns a correct S_dmask

code:

attn_output, softmax_lse, S_dmask = flash_attn_func(query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal, return_attn_probs=True)
print(S_dmask)
>> tensor([], device='cuda:0', dtype=torch.float16)

[tridao]

As the docstring says, it's for testing only.

        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
           testing only. The returned probabilities are not guaranteed to be correct
           (they might not have the right scaling).

In fact for dropout=0.0 we will return None.