Flash Attention doesn't support attention mask

[clownrat6]

Test code from this:

import time
import torch
import torch.nn.functional as F


N = 32
L = 2048
dims = 64
n_heads = 8
q = torch.randn(N, n_heads, L, dims, dtype=torch.float16).cuda()
k = torch.randn(N, n_heads, L, dims, dtype=torch.float16).cuda()
v = torch.randn(N, n_heads, L, dims, dtype=torch.float16).cuda()

dropout_rate = 0.2
num_trials = 10

with torch.backends.cuda.sdp_kernel(
    enable_flash=True, enable_math=False, enable_mem_efficient=False
):

    attn_bias = torch.zeros(N, n_heads, L, L, dtype=q.dtype).to(q.device)
    torch.cuda.synchronize()
    start = time.time()
    for i in range(num_trials):
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias, dropout_p=dropout_rate)
    torch.cuda.synchronize()
    end = time.time()
    print('Flash attention took {} seconds for {} trials'.format(end - start, num_trials))

Errors:

Related issues:
Dao-AILab/flash-attention#352

[LinB203]

Thanks! We fixed that by turning to xformers.
See https://github.com/PKU-YuanGroup/Open-Sora-Plan/blob/main/opensora/models/diffusion/latte/latte.py#L109-L118

[clownrat6]

Bin God 🙏🙏