Whisper models appear to be broken with Flash Attention 2
Author: Anjum48
Pull request: no
Labels: name: bug

name: Audio

name: Whisper

### System Info

- `transformers` version: 4.52.4
- Platform: Linux-6.8.0-60-generic-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.32.4
- Safetensors version: 0.4.5
- Accelerate version: 1.7.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.7.1+cu128 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA GeForce RTX 5090

flash-attn: Built from main

### Who can help?

@eustlb 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Code to reproduce:
```python
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3.5"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="flash_attention_2",
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    torch_dtype=torch_dtype,
    device=device,
    return_timestamps=True
)

result = pipe(str(audio_path))
```

The audio file is long (~2 hours). Note that the model runs fine if flash attention is disabled. Have tested with various Whisper models and get the same error. 

Rolling back to 4.46.1 seems to run ok, but I'm not sure flash attn is actually working. Transcription on a 5090 for some models (e.g. openai/whisper-base) is significantly slower than a 3090, but this might be due to hardware optimisations on the flash attn side...

Error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[13], line 28
     15 processor = AutoProcessor.from_pretrained(model_id)
     17 pipe = pipeline(
     18     "automatic-speech-recognition",
     19     model=model,
   (...)
     25     return_timestamps=True
     26 )
---> 28 result = pipe(str(audio_path))

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/pipelines/automatic_speech_recognition.py:295, in AutomaticSpeechRecognitionPipeline.__call__(self, inputs, **kwargs)
    234 def __call__(
    235     self,
    236     inputs: Union[np.ndarray, bytes, str],
    237     **kwargs,
    238 ):
    239     """
    240     Transcribe the audio sequence(s) given as inputs to text. See the [`AutomaticSpeechRecognitionPipeline`]
    241     documentation for more information.
   (...)
    293                 `"".join(chunk["text"] for chunk in output["chunks"])`.
    294     """
--> 295     return super().__call__(inputs, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/pipelines/base.py:1423, in Pipeline.__call__(self, inputs, num_workers, batch_size, *args, **kwargs)
   1421     return self.iterate(inputs, preprocess_params, forward_params, postprocess_params)
   1422 elif self.framework == "pt" and isinstance(self, ChunkPipeline):
-> 1423     return next(
   1424         iter(
   1425             self.get_iterator(
   1426                 [inputs], num_workers, batch_size, preprocess_params, forward_params, postprocess_params
   1427             )
   1428         )
   1429     )
   1430 else:
   1431     return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/pipelines/pt_utils.py:124, in PipelineIterator.__next__(self)
    121     return self.loader_batch_item()
    123 # We're out of items within a batch
--> 124 item = next(self.iterator)
    125 processed = self.infer(item, **self.params)
    126 # We now have a batch of "inferred things".

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/pipelines/pt_utils.py:269, in PipelinePackIterator.__next__(self)
    266             return accumulator
    268 while not is_last:
--> 269     processed = self.infer(next(self.iterator), **self.params)
    270     if self.loader_batch_size is not None:
    271         if isinstance(processed, torch.Tensor):

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/pipelines/base.py:1338, in Pipeline.forward(self, model_inputs, **forward_params)
   1336     with inference_context():
   1337         model_inputs = self._ensure_tensor_on_device(model_inputs, device=self.device)
-> 1338         model_outputs = self._forward(model_inputs, **forward_params)
   1339         model_outputs = self._ensure_tensor_on_device(model_outputs, device=torch.device("cpu"))
   1340 else:

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/pipelines/automatic_speech_recognition.py:527, in AutomaticSpeechRecognitionPipeline._forward(self, model_inputs, return_timestamps, **generate_kwargs)
    524 if "generation_config" not in generate_kwargs:
    525     generate_kwargs["generation_config"] = self.generation_config
--> 527 tokens = self.model.generate(
    528     inputs=inputs,
    529     attention_mask=attention_mask,
    530     **generate_kwargs,
    531 )
    532 # whisper longform generation stores timestamps in "segments"
    533 if return_timestamps == "word" and self.type == "seq2seq_whisper":

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/generation_whisper.py:774, in WhisperGenerationMixin.generate(self, input_features, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, return_timestamps, task, language, is_multilingual, prompt_ids, prompt_condition_type, condition_on_prev_tokens, temperature, compression_ratio_threshold, logprob_threshold, no_speech_threshold, num_segment_frames, attention_mask, time_precision, time_precision_features, return_token_timestamps, return_segments, return_dict_in_generate, force_unique_generate_call, **kwargs)
    765             proc.set_begin_index(decoder_input_ids.shape[-1])
    767 # 6.6 Run generate with fallback
    768 (
    769     seek_sequences,
    770     seek_outputs,
    771     should_skip,
    772     do_condition_on_prev_tokens,
    773     model_output_type,
--> 774 ) = self.generate_with_fallback(
    775     segment_input=segment_input,
    776     decoder_input_ids=decoder_input_ids,
    777     cur_bsz=cur_bsz,
    778     batch_idx_map=batch_idx_map,
    779     seek=seek,
    780     num_segment_frames=num_segment_frames,
    781     max_frames=max_frames,
    782     temperatures=temperatures,
    783     generation_config=generation_config,
    784     logits_processor=logits_processor,
    785     stopping_criteria=stopping_criteria,
    786     prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
    787     synced_gpus=synced_gpus,
    788     return_token_timestamps=return_token_timestamps,
    789     do_condition_on_prev_tokens=do_condition_on_prev_tokens,
    790     is_shortform=is_shortform,
    791     batch_size=batch_size,
    792     attention_mask=attention_mask,
    793     kwargs=kwargs,
    794 )
    796 # 6.7 In every generated sequence, split by timestamp tokens and extract segments
    797 for i, seek_sequence in enumerate(seek_sequences):

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/generation_whisper.py:950, in WhisperGenerationMixin.generate_with_fallback(self, segment_input, decoder_input_ids, cur_bsz, batch_idx_map, seek, num_segment_frames, max_frames, temperatures, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, return_token_timestamps, do_condition_on_prev_tokens, is_shortform, batch_size, attention_mask, kwargs)
    945     if generate_kwargs.get("encoder_outputs") is not None:
    946         generate_kwargs["encoder_outputs"] = F.pad(
    947             generate_kwargs["encoder_outputs"], (0, 0, 0, 0, 0, batch_size - cur_bsz), value=0
    948         )
--> 950 seek_outputs = super().generate(
    951     segment_input,
    952     generation_config=generation_config,
    953     logits_processor=logits_processor,
    954     stopping_criteria=stopping_criteria,
    955     prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
    956     synced_gpus=synced_gpus,
    957     decoder_input_ids=decoder_input_ids,
    958     attention_mask=attention_mask,
    959     **generate_kwargs,
    960 )
    962 model_output_type = type(seek_outputs)
    964 # post-process sequence tokens and outputs to be in list form

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/generation/utils.py:2616, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, use_model_defaults, custom_generate, **kwargs)
   2609     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2610         input_ids=input_ids,
   2611         expand_size=generation_config.num_beams,
   2612         is_encoder_decoder=self.config.is_encoder_decoder,
   2613         **model_kwargs,
   2614     )
   2615     # 12. run beam sample
-> 2616     result = self._beam_search(
   2617         input_ids,
   2618         logits_processor=prepared_logits_processor,
   2619         stopping_criteria=prepared_stopping_criteria,
   2620         generation_config=generation_config,
   2621         synced_gpus=synced_gpus,
   2622         **model_kwargs,
   2623     )
   2625 elif generation_mode == GenerationMode.GROUP_BEAM_SEARCH:
   2626     # 11. prepare beam search scorer
   2627     beam_scorer = BeamSearchScorer(
   2628         batch_size=batch_size,
   2629         num_beams=generation_config.num_beams,
   (...)
   2635         max_length=generation_config.max_length,
   2636     )

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/generation/utils.py:4030, in GenerationMixin._beam_search(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, **model_kwargs)
   4027 model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
   4028 model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
-> 4030 model_outputs = self(**model_inputs, return_dict=True)
   4032 # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
   4033 model_kwargs = self._update_model_kwargs_for_generation(
   4034     model_outputs,
   4035     model_kwargs,
   4036     is_encoder_decoder=self.config.is_encoder_decoder,
   4037 )

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1751, in Module._wrapped_call_impl(self, *args, **kwargs)
   1749     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1750 else:
-> 1751     return self._call_impl(*args, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1762, in Module._call_impl(self, *args, **kwargs)
   1757 # If we don't have any hooks, we want to skip the rest of the logic in
   1758 # this function, and just call forward.
   1759 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1760         or _global_backward_pre_hooks or _global_backward_hooks
   1761         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1762     return forward_call(*args, **kwargs)
   1764 result = None
   1765 called_always_called_hooks = set()

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/modeling_whisper.py:1694, in WhisperForConditionalGeneration.forward(self, input_features, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, decoder_inputs_embeds, decoder_position_ids, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
   1689     if decoder_input_ids is None and decoder_inputs_embeds is None:
   1690         decoder_input_ids = shift_tokens_right(
   1691             labels, self.config.pad_token_id, self.config.decoder_start_token_id
   1692         )
-> 1694 outputs = self.model(
   1695     input_features,
   1696     attention_mask=attention_mask,
   1697     decoder_input_ids=decoder_input_ids,
   1698     encoder_outputs=encoder_outputs,
   1699     decoder_attention_mask=decoder_attention_mask,
   1700     head_mask=head_mask,
   1701     decoder_head_mask=decoder_head_mask,
   1702     cross_attn_head_mask=cross_attn_head_mask,
   1703     past_key_values=past_key_values,
   1704     decoder_inputs_embeds=decoder_inputs_embeds,
   1705     decoder_position_ids=decoder_position_ids,
   1706     use_cache=use_cache,
   1707     output_attentions=output_attentions,
   1708     output_hidden_states=output_hidden_states,
   1709     return_dict=return_dict,
   1710     cache_position=cache_position,
   1711 )
   1712 lm_logits = self.proj_out(outputs[0])
   1714 loss = None

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1751, in Module._wrapped_call_impl(self, *args, **kwargs)
   1749     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1750 else:
-> 1751     return self._call_impl(*args, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1762, in Module._call_impl(self, *args, **kwargs)
   1757 # If we don't have any hooks, we want to skip the rest of the logic in
   1758 # this function, and just call forward.
   1759 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1760         or _global_backward_pre_hooks or _global_backward_hooks
   1761         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1762     return forward_call(*args, **kwargs)
   1764 result = None
   1765 called_always_called_hooks = set()

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/modeling_whisper.py:1529, in WhisperModel.forward(self, input_features, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, decoder_inputs_embeds, decoder_position_ids, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
   1522     encoder_outputs = BaseModelOutput(
   1523         last_hidden_state=encoder_outputs[0],
   1524         hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
   1525         attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
   1526     )
   1528 # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
-> 1529 decoder_outputs = self.decoder(
   1530     input_ids=decoder_input_ids,
   1531     attention_mask=decoder_attention_mask,
   1532     encoder_hidden_states=encoder_outputs[0],
   1533     head_mask=decoder_head_mask,
   1534     cross_attn_head_mask=cross_attn_head_mask,
   1535     past_key_values=past_key_values,
   1536     inputs_embeds=decoder_inputs_embeds,
   1537     position_ids=decoder_position_ids,
   1538     use_cache=use_cache,
   1539     output_attentions=output_attentions,
   1540     output_hidden_states=output_hidden_states,
   1541     return_dict=return_dict,
   1542     cache_position=cache_position,
   1543 )
   1545 if not return_dict:
   1546     return decoder_outputs + encoder_outputs

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1751, in Module._wrapped_call_impl(self, *args, **kwargs)
   1749     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1750 else:
-> 1751     return self._call_impl(*args, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1762, in Module._call_impl(self, *args, **kwargs)
   1757 # If we don't have any hooks, we want to skip the rest of the logic in
   1758 # this function, and just call forward.
   1759 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1760         or _global_backward_pre_hooks or _global_backward_hooks
   1761         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1762     return forward_call(*args, **kwargs)
   1764 result = None
   1765 called_always_called_hooks = set()

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/modeling_whisper.py:1188, in WhisperDecoder.forward(self, input_ids, attention_mask, encoder_hidden_states, head_mask, cross_attn_head_mask, past_key_values, inputs_embeds, position_ids, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
   1174     layer_outputs = self._gradient_checkpointing_func(
   1175         decoder_layer.__call__,
   1176         hidden_states,
   (...)
   1185         cache_position,
   1186     )
   1187 else:
-> 1188     layer_outputs = decoder_layer(
   1189         hidden_states,
   1190         attention_mask=causal_mask,
   1191         encoder_hidden_states=encoder_hidden_states,
   1192         layer_head_mask=(head_mask[idx] if head_mask is not None else None),
   1193         cross_attn_layer_head_mask=(
   1194             cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
   1195         ),
   1196         past_key_value=past_key_values if use_cache else None,
   1197         output_attentions=output_attentions,
   1198         use_cache=use_cache,
   1199         cache_position=cache_position,
   1200     )
   1201 hidden_states = layer_outputs[0]
   1203 if output_attentions:

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1751, in Module._wrapped_call_impl(self, *args, **kwargs)
   1749     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1750 else:
-> 1751     return self._call_impl(*args, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1762, in Module._call_impl(self, *args, **kwargs)
   1757 # If we don't have any hooks, we want to skip the rest of the logic in
   1758 # this function, and just call forward.
   1759 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1760         or _global_backward_pre_hooks or _global_backward_hooks
   1761         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1762     return forward_call(*args, **kwargs)
   1764 result = None
   1765 called_always_called_hooks = set()

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/modeling_whisper.py:727, in WhisperDecoderLayer.forward(self, hidden_states, attention_mask, encoder_hidden_states, encoder_attention_mask, layer_head_mask, cross_attn_layer_head_mask, past_key_value, output_attentions, use_cache, cache_position)
    725 residual = hidden_states
    726 hidden_states = self.encoder_attn_layer_norm(hidden_states)
--> 727 hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
    728     hidden_states=hidden_states,
    729     key_value_states=encoder_hidden_states,
    730     attention_mask=encoder_attention_mask,
    731     layer_head_mask=cross_attn_layer_head_mask,
    732     past_key_value=past_key_value,
    733     output_attentions=output_attentions,
    734 )
    735 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    736 hidden_states = residual + hidden_states

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1751, in Module._wrapped_call_impl(self, *args, **kwargs)
   1749     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1750 else:
-> 1751     return self._call_impl(*args, **kwargs)

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1762, in Module._call_impl(self, *args, **kwargs)
   1757 # If we don't have any hooks, we want to skip the rest of the logic in
   1758 # this function, and just call forward.
   1759 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1760         or _global_backward_pre_hooks or _global_backward_hooks
   1761         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1762     return forward_call(*args, **kwargs)
   1764 result = None
   1765 called_always_called_hooks = set()

File ~/github/tapesearch/.venv/lib/python3.12/site-packages/transformers/models/whisper/modeling_whisper.py:401, in WhisperFlashAttention2.forward(self, hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, output_attentions, cache_position)
    399     value_states = past_key_value.value_cache[self.layer_idx]
    400 else:
--> 401     key_states = self.k_proj(current_states).view(bsz, tgt_len, self.num_heads, self.head_dim)
    402     value_states = self.v_proj(current_states).view(bsz, tgt_len, self.num_heads, self.head_dim)
    403     key_states = key_states.transpose(1, 2).contiguous()

RuntimeError: shape '[5, 3, 20, 64]' is invalid for input of size 9600000
```


### Expected behavior

Expect the pipeline to run in the same way as when flash attention is disabled (but faster!)

Comments:
Hi @eustlb, I’d like to grab this, please.



cc @vasqu since you've been working on Whisper-related attention issues

@Anjum48 Can you try the latest transformers version (`4.53.0`)? We have rewritten the attention module in #38235 to work with all attention types (fa2 was known to not work as expected).

Regarding speed: The transformers implementation with flash attention might be slow(er) on small inputs (<2k) because of how we have to handle attention masks. I would suggest SDPA in most cases. #36585 might be interesting to you (as there are perf comparisons on whisper parts)

Thanks @vasqu! Can confirm that FA2 is now working and on benchmarking, sdpa is slightly faster as discussed above

Hi @vasqu - reopening this as after some benchmarking between 4.46 and 4.53, it looks like the new attention module is slowing things down quite significantly (times are in M.SS):

```
transformers==4.46.1 sdpa

openai/whisper-base: 4.50
distil-whisper/distil-medium.en: 1.20
distil-whisper/distil-large-v3: 1.23
distil-whisper/distil-large-v3.5: 1.28
openai/whisper-large-v3-turbo: 3.36


transformers==4.46.1 FA2

openai/whisper-base: 4.57
distil-whisper/distil-medium.en: 1.18
distil-whisper/distil-large-v3:  1.21
distil-whisper/distil-large-v3.5: 1.30 
openai/whisper-large-v3-turbo:  4.00


transformers==4.53.1 sdpa

openai/whisper-base: 6.13
distil-whisper/distil-medium.en: 3.07
distil-whisper/distil-large-v3: 3.04
distil-whisper/distil-large-v3.5:  3.05
openai/whisper-large-v3-turbo: 5.12


transformers==4.53.1 FA2

openai/whisper-base: 6.42
distil-whisper/distil-medium.en: 3.10
distil-whisper/distil-large-v3: 3.04
distil-whisper/distil-large-v3.5: 3.08
openai/whisper-large-v3-turbo: 5.45
```

This is benchmarking with 125 mins of audio from [here](https://traffic.libsyn.com/secure/skepticsguide/skepticast2025-05-31.mp3)

Code used for tests:

```python
import torch
from whisper.utils import get_writer
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from pathlib import Path

audio_path = Path("/home/anjum/Downloads/skepticast2025-05-31.mp3")  # 124 mins
output_path = Path("/home/anjum/Documents")

device_id = 1
device = f"cuda:{device_id}" if torch.cuda.is_available() else "cpu"
torch_dtype = (
    torch.float16 if torch.cuda.is_available() else torch.float32
)

torch.set_float32_matmul_precision("high")
torch.cuda.empty_cache()

def format_hf_to_whisper(result):
    results_formatted = {"text": result["text"], "segments": [], "language": None}

    for chunk in result["chunks"]:
        if chunk["timestamp"][0] is None:
            start = 0
        else:
            start = chunk["timestamp"][0]

        if chunk["timestamp"][1] is None:
            end = start
        else:
            end = chunk["timestamp"][1]

        segment = {
            "start": start,
            "end": end,
            "text": chunk["text"],
        }
        results_formatted["segments"].append(segment)
    return results_formatted

def run_hf_model(model_id, params={}):
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        low_cpu_mem_usage=True,
        use_safetensors=True,
        # attn_implementation="flash_attention_2",
        attn_implementation="sdpa",
    )
    model.to(device)

    processor = AutoProcessor.from_pretrained(model_id)

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=torch_dtype,
        device=device,
        return_timestamps=True,
        **params,
    )

    generate_kwargs = {
        "max_new_tokens": 128,
        "return_timestamps": True,
    }

    result = pipe(str(audio_path), generate_kwargs=generate_kwargs)
    result = format_hf_to_whisper(result)

    writer = get_writer("all", str(output_path))
    writer_args = {
        "highlight_words": False,
        "max_line_count": None,
        "max_line_width": None,
    }
    writer(result, f"hf_{model_id.replace('/', '_')}", writer_args)
    torch.cuda.empty_cache()

batch_size = 16

run_hf_model("distil-whisper/distil-large-v3", {"chunk_length_s": 25, "batch_size": batch_size, "ignore_warning": True}) 
```

It looks like in these tests `distil-whisper/distil-large-v3` is about 2x slower :(

@Anjum48 This is indeed not so nice 👀 Since it's a few minor versions, it might not be related to the attention calculation tbh - there have been some changes to the mask creation and the generation code as well, for example. Not sure when I would get to dig deeper into this; if you'd like to, you could also investigate potential bad commits with `git bisect`.

Hi @vasqu, I wasn't sure how to tackle a `git bisect` since I'm not 100% sure which files to watch for changes.

However, I did a manual binary search :) and I located that the performance regression happened between 4.51.3 -> 4.52.1. My tests included 2 repeats and confirmed the 2x slower times I was seeing above.

The [release notes](https://github.com/huggingface/transformers/releases/tag/v4.52.1) show a number of PRs related to attention and a few other things that might be relevant. Probably worth you scanning through those PRs to see if anything jumps out to you as a potential culprit.

This is the only Whisper related PR in that release:
> Fix incorrect attention mask truncate in WhisperFlashAttention2 by @OliBomby in https://github.com/huggingface/transformers/pull/36477

Tagging @eustlb too for visibility.

This is obviously turning into a different issue than the one I originally opened. Happy to rename this issue or open a new one - whichever you all think is best.

Indeed, it would be nice to track this in a separate issue. 

Thanks for already looking into it @Anjum48 :D I'm a bit suprised that it's v4.52 because a lot of changes happened during v4.53 and it makes me believe that it's not attention-related but something else. For `git bisect`, [stackoverflow](https://stackoverflow.com/questions/4713088/how-do-i-use-git-bisect) has a nice answer here - you will need to work on the git repository instead of via pip and you could just test with your performance test each time to determine good vs bad commit. 

Opened a new issue so closing this one