Uncertainty in FlashAttention2 backward

[ezioliao]

When I ran the unit test "test_flash_atten.py" for the function "flash_attn_varlen_func", I noticed that the values of dq in "flash_attn_varlen_func" during the backward process has slight distinct in different test(with a fixed random seed and the same input), but dk and dv are exactly same. During the forward pass, the output for the same input remained consistent. Does this indicate that in FA2 the forward is deterministic, but the backward is not? What could be the possible reasons for this uncertainty?

[tridao]

Yes, backward pass is not deterministic because we use atomic adds.

[ezioliao]

Thanks for your answer! :)