Starcoder2 has +10% inference latency when flash attention 2 is enabled
Author: lidingsnyk
Pull request: no

### System Info

transformers version: 4.39.3
flash-attn: Tried both 2.5.2 and 2.5.8  (flash_attn-2.5.2+cu118torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64)
Platform: linux_x86_64 cp310 ubuntu-22.04
Python version: 3.10
Huggingface_hub version: 0.20.3
Safetensors version: 0.4.2
Accelerate version: 0.26.1
Accelerate config: not found
PyTorch version (GPU): 2.1.2 (True) torch-2.1.2-cu118-cp310-cp310-linux_x86_64.whl
Flax version (CPU?/GPU?/TPU?): not installed
Jax version: not installed
JaxLib version: not installed
Using GPU in script?: yes. A100
CUDA_VERSION: 11.8.0

### Who can help?

@amyeroberts

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

We currently have a generation task of fixing bad code, with a max length of 512.  We've noticed that when flash attention 2 is enabled, the inference latency is almost uniformly +10% across different input length. For instance, latency to our max length increases from 19.9 seconds to 22.6s. 

With everything else being equal, starcoder1 sees latency reduction instead.

```
 def make_inference_pipeline() # Redacted for now

  pipeline: transformers.Pipeline = make_inference_pipeline(
        model,
        tokenizer,
        truncate_input=USE_TRUNCATION,
        max_input_size=MAX_INPUT_SIZE,
    )

    pipeline_iterator = pipeline(
        test_ds_wrapper,
        num_workers=args.training_args.dataloader_num_workers,
        batch_size=args.training_args.per_device_eval_batch_size,
        **make_generation_config(args.inference_args, tokenizer),
    )

    # `predictions` is a list (with size `num_samples`) of lists (with size `num_return_seqs`) that contains
    # the beam search predictions for every sample.
    predictions = [
        [out["generated_text"] for out in cast(list[dict], predictions_for_one_sample)]
        # `predictions_for_one_sample` is a list that contains the `num_return_seqs` predictions for a sample.
        for predictions_for_one_sample in tqdm.tqdm(
            pipeline_iterator,
            desc=f"Generating predictions {data_id} for shard: {worker_id}",
            position=worker_id,
        )
    ]
```
bf16: "True"

Side note: We also noticed a latency increase when attempting to upgrade from starcoder1 to starcoder2. With the latency of max length sequence being 16.5s vs 19.9s. ORT is not supported for starcoder2. flash-attn-2 behavior is mentioned above, which prompted me to report this bug.



### Expected behavior

flash attention 2 provides latency reduction for starcoder2 compared to the  defaults.

Comments:
FYI @agronskiy @BBerabi


Maybe cc @ArthurZucker, FA expert

Hey! The starcoder2 is pretty much copied from our mistral model. Thus I would refer to the benchmarks and test we do on mistral. #30041 #28130 are good reference points for this 😉 

Thanks! The existing benchmark and discussion suggest that for low input length and low batch size, FA2 is known to run slower. Looks like this issue is the same as the existing one then.