Question about flash attention dropout

Hi

I noticed that with attention dropout, besides dropout for P, in the final write, there is code as follows

for (int jj = 0; jj < Gmem_tile_o::STGS_PER_LOOP; jj++) {
            float sum = p_sum_o[jj][0];
            float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum;
            if (Is_dropout && is_final_write) {
                inv_sum *= params.rp_dropout;
            }
            out[jj] = fmha::fmul4(out[jj], inv_sum);
        }

https://github.com/HazyResearch/flash-attention/blob/main/csrc/flash_attn/src/fmha_fprop_kernel_1xN.h#L629

same for backward https://github.com/HazyResearch/flash-attention/blob/main/csrc/flash_attn/src/fmha_dgrad_kernel_1xN_loop.h#L689

I did not find it in the paper

[tridao]

There’s a trick in the implementation here to save storage: This line generates the dropout mask and stores the dropout mask as sign of the elements in P (e.g. we store 0.1 if attention prob is 0.1 and the corresponding dropout mask is 1, we store -0.2 if attention prob is 0.2 and the corresponding dropout mask is 0).

This block of code corresponds to line 18 and 20: if the value we store in p is negative, then we get ds = p * d (i.e., -attn_prob * d), because it corresponds to dropout mask being zero. If the value we store in p is positive, then we get ds = p * acc_dp (i.e., attn_prob * (dp - d)), since it corresponds to dropout mask being 1.

Later on when we need to use P we take relu to remove the dropout mask that was being stored as sign of P.

If you want to implement it it’s probably simpler to store the dropout mask separately without this trick.

[SeaOfOcean]

Another question, what’s the code corresponding line:18 in algorithm?