Some questions

Hello,

Thank you for your awesome work! Flash attention is going to be used everywhere!

I have a few questions please:

1. To use flash attention in an existing PyTorch transformer, it suffices to replace torch.nn.MultiheadAttention with flash_attn.flash_attention.FlashMHA, is that correct?
2. Training is also supported out-of-the-box I guess? The question also includes mixed precision training, i.e., compatibility with torch.autocast() context manager.
3. I see that you also provide a Fused Softmax implementation. According to the docstrings, this layer is used for auto-regressive models. If I only use the transformer encoder, e.g., vision transformers, then it’s not worth using it. Is that correct?

Thank you very much in advance for your answers.

[tridao]

1. To use flash attention in an existing PyTorch transformer, it suffices to replace torch.nn.MultiheadAttention with flash_attn.flash_attention.FlashMHA, is that correct?

Yes. The two modules do the same thing, though they might have different APIs and arguments. You should read the arguments and documents to pass in the right things.

1. Training is also supported out-of-the-box I guess? The question also includes mixed precision training, i.e., compatibility with torch.autocast() context manager.

Yes, training works, and mixed-precision training works. torch.autocast() will do the right thing and cast the q, k, v to either fp16 or bf16. FlashAttention does not support fp32.

1. I see that you also provide a Fused Softmax implementation. According to the docstrings, this layer is used for auto-regressive models. If I only use the transformer encoder, e.g., vision transformers, then it’s not worth using it. Is that correct?

The Fused Softmax was taken from apex/megatron purely for benchmarking. They’re only useful if either you have causal mask before softmax (e.g. autoregressive models) or key padding mask before softmax (e.g. BERT where sequences in a batch have different lengths). If you’re using transformer encoder and all sequences in the batch have the same length, then they’re won’t apply to your case.

[netw0rkf10w]

@tridao Great, thanks! Will create a PR soon!