Can flash attention be used for inference?

[Shannen3206]

I tried inference with and without flash attention in the megatron-deepspeed code and found a difference in inference speed of just 0.2 seconds.

In addition, in huggingface's openllama model structure, flash attention is also limited to training.

Can flash attention be used for inference acceleration?

[tridao]

FlashAttention right now isn't very fast for iterative decoding where Q has seqlen=1 (you can use it for prompt processing). It's ongoing work, we'll eventually make inference fast (idk when yet).

[77h2l]

@tridao hello，can flash_attn only be used for inference, in other word, when a model was trained by standard attention, can it use flash_attn to replace during infernce,thx

[tridao]

@tridao hello，can flash_attn only be used for inference, in other word, when a model was trained by standard attention, can it use flash_attn to replace during infernce,thx

Yes, FlashAttention computes the same attention (up to usual numerical difference). You can also train with FlashAttention and do inference with standard attention.

[jmzeng]

Want to follow up on this. I tried to do inference on longchat-v1.5-7b-32k with the flash attention patch and was getting OOM on a single inference run with accelerate. Do you know if there is a way flash attention 2 can be used?

[tridao]

Yes you can use FlashAttention 2, for both prompt processing and iterative decoding. It's now optimized for both.

[jmzeng]

Thanks for the update. I opened a new issue for my OOM, see here:

#590

[thusinh1969]

@tridao hello，can flash_attn only be used for inference, in other word, when a model was trained by standard attention, can it use flash_attn to replace during infernce,thx

Yes, FlashAttention computes the same attention (up to usual numerical difference). You can also train with FlashAttention and do inference with standard attention.

Love you my man (a VN hard core here...!)

Cheers,
Steve