Flash attention has no effect on inference
Hi, I was exploring the benefits of using flash attention 2 with Mistral and Mixtral during inference. Yet, I can see no memory reduction & no speed acceleration. Some number under different attention implementations:Mixtral (mistralai/Mixtral-8x7B-Instruct-v0.1):attn_implementation=‘flash_attention_2’:27.335Gb, 15.8 secondsattn_implementation=‘eager’:27.335Gb, 16.1 secondsattn_implementation=‘sdpa’:27.335Gb, 15.4 secondsMistral (mistralai/Mistral-7B-Instruct-v0.2):attn_implementation=‘flash_attention_2’:6.407Gb, 31.5 secondsattn_implementation=‘eager’:6.407Gb, 30.7 secondsattn_implementation=‘sdpa’:6.405Gb, 28.9 seconds(With Mistral it took much more in terms of speed compared to Mixtral because I tested on 20 examples with smallermax_new_tokens).The code I used to test Mixtral is taken fromHuggingFace page:import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
device = "cuda" # the device to load the model onto

model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1", torch_dtype=torch.float16, attn_implementation=attn_implementation, cache_dir="cache", load_in_4bit=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1")

prompt = "My favourite condiment is"

model_inputs = tokenizer([prompt], return_tensors="pt").to(device)

%time generated_ids = model.generate(**model_inputs, max_new_tokens=100, do_sample=True)
tokenizer.batch_decode(generated_ids)[0]Library versions:transformers==4.37.2
torch==2.2.0
flash-attn==2.5.3My GPU is NVIDIA A100-SXM4-40GB with cuda release 12.3, V12.3.107.I know that the major benefit of flash-attention-2 blossoms out during training, yet, it is reported to be beneficial during inference as well:HuggingFace Page. Hence, I wonder whether I may have any problems with my environment or there are any problems with the current implementations.Thanks in advance!
2024-02-21T09:02:51Z
@ybelkadahope you can shed some light here
2024-02-22T18:58:00Z
Also experiencing this using Mistral-7B-v0.1 on a 4090. Running for:4096 tokensbatch size of 1No paddingMulti-head attention takes ~144 seconds.Flash-attention takes ~154.
2024-02-28T01:31:04Z
Hi everyone !@Aktsvigunto see the effect of FA2 you need to run inference on a large context length. In the benchmarks I ran where I saw interesting speedups I used ~2048 tokensI used this script to benchmark everything:Benchmark FA2 + transformers integration · GitHub@kreascan you try with transformers==4.35.2 ?
2024-02-29T01:27:41Z
Changing transformers version doesn’t seem to affect anything;I found out that I had some NVTX calls which added overhead, but even after removing, Flash attention is slower on mistral. I’ve slightly modified the script you linked in your comment. These are the results:
2024-04-06T15:04:12Z
It speeds up other models such as  phi-2 for the same text around 3x, while it actually slows down mistral v0.2 instruct.I think mistral even say in some of descriptions of the model it is very fast because uses flash attention. Not sure. Maybe it it used by default if available.
2024-08-22T12:37:26Z
@Aktsviguntrytorch_dtype=torch.bfloat16
2024-09-04T18:43:41Z
Hi! I thought I would chime in here since I’m also trying to figure out how much flash_attention is improving inference speed. I modified Younes Belkada’s script here and got some results.gist.github.comhttps://gist.github.com/younesbelkada/02f35734da906cc0f2389ae4f665c58f?permalink_comment_id=5178402#gistcomment-5178402bench-fa-2.pyimport torch
import os
import argparse
import matplotlib.pyplot as plt
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
import seaborn as sns


def get_parser():This file has been truncated.show originalI found that FA improvements were only realized after context length went past 4k and when use_cache was set to True.