Performance regression with FA2 and long-context generation after 5.2.0
Author: andreasgoulas
Pull request: no
Labels: name: bug

### System Info

transformers: 5.2.0 vs 5.4.0 (I could not test 5.3.0 due to a bug)
PyTorch: 2.12
flash-attn: 2.8.3
CUDA: 13.2
Python: 3.12
cudnn: 9.20.0.48

### Who can help?

@zucchini-nlp 
@Cyrilvallez 

### 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

1. Install the latest version of transformers
2. Install FlashAttention2
3. Run the provided code
4. Downgrade to transformers version 5.2.0
5. Run the code again

```python
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
import torch

model = Qwen3VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen3-VL-2B-Thinking",
    dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
).to("cuda")

processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-2B-Thinking")

messages = [
    [
        {
            "role": "user",
            "content": [
                {"type": "video", "video": torch.zeros(768, 3, 512, 512)},
                {"type": "text", "text": "placeholder"},
            ],
        }
    ],
    [
        {
            "role": "user",
            "content": [
                {"type": "video", "video": torch.zeros(768, 3, 512, 512)},
                {"type": "text", "text": "different placeholder"},
            ],
        }
    ],
]

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    do_sample_frames=False,
    return_dict=True,
    return_tensors="pt",
    padding=True,
    padding_side="left",
).to("cuda")

MAX_LEN = 16 * 1024

while True:
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()
    out_ids = model.generate(
        **inputs, do_sample=False, min_length=MAX_LEN, max_length=MAX_LEN, use_cache=True
    )
    end.record()

    torch.cuda.synchronize()
    assert out_ids.size(0) == 2 and out_ids.size(1) == MAX_LEN

    cur_time = start.elapsed_time(end)
    print(cur_time)
```

### Expected behavior

Tested on an NVIDIA RTX 6000 PRO:
- v5.2.0 FA2 = 92 sec
- v5.4.0 FA2 = 108 sec (+17% overhead)
- v5.2.0 SDPA = 136 sec
- v5.4.0 SDPA = 136 sec (no change)

The performance regression seems to occur when the KV cache is sufficiently large, at which point the GPU memory bandwidth has been saturated. Furthermore, it seems to be related with the changes to the FA2 integration layer. It does not seem to be related with the nvidia-cudnn conv kernel performance regression which was fixed in a recent PyTorch version.

Comments:
will take a look tomorrow, can confirm there is a small diff (I got ~200sec and ~190sec on avg per release tag)

#46696 is the wrong fix but it still might have found the underlying perf regression(?), just as a heads up

We might need to cache the compiled state 

Looks like it was caused by https://github.com/huggingface/transformers/pull/42435. @vasqu could you take a look, checking out to the commit before that PR brings back the performance

Yea lemme check what happened. Feel like the is_tracing call is too expensive it seems

@vasqu It seems that patching `is_tracing` in import_utils.py to always return False does not significantly impact the generation latency (tested on 5.4.0). So, this is probably not the root cause.

According to my measurements, `flash_varlen_fn` is significantly more expensive on 5.4.0 vs 5.2.0 (240us vs 140us - excluding the kwargs preparation code). This is quite strange since the inputs and the flash-attn implementation should be identical across the two versions. @vasqu 

This is indeed super weird, the tracing call was just my intuition tbh. I haven't had time to properly investigate yet sadly, if you revert the flash attention stuff in the forward, do we still get the same perf then?

So revert the PR only for checking kwargs and the associated main forward

@vasqu After reverting the PR line-by-line, I think I found the root cause: removing the `.item()` call from max_seqlen_in_batch in `_get_unpad_data` somehow causes the next `_index_first_axis` calls to trigger this performance regression. It is not clear to me how this is possible - it could be the case that there is a memory leak or a hidden memcpy. A straightforward fix would be to re-introduce the `item()` call there, but I think that this issue requires a more careful investigation.

Ok wow that is really unexpected 😮 if you add a .item call before calling the index fn, everything is fine? Is there any difference between native FA2 and kernels FA2? `uv pip install kernels`

Because the native FA2 fn and the kernels ones should be slightly different 

Yes, adding `max_seqlen_in_batch_k = max_seqlen_in_batch_k.item()` before the 2 index calls resolves the issue, at least on an NVIDIA RTX 6000 PRO. I also tested an RTX 4090 but the performance gap was not that significant there. ~Installing kernels==0.14.1 did not significantly affect the measured latency.~ @vasqu 

Using the FA2 implementation from kernels, I can observe a small performance gain but the performance regression is still there.

That's super weird. So it seems to be GPU specific? I'm not sure whether it makes sense as general fix then, i.e. always call .item as it does have its own repercussions 😓 We can at least open a draft PR if you are open to it @andreasgoulas

@zucchini-nlp you also measured the latency, which GPU did you use? 

I check on the dgx with kernels FA, @vasqu

Yes, I will submit a draft PR reverting this specific change. I wouldn't say it is GPU-specific, but depending on the GPU the effect may be  more or less significant.

@vasqu I have submitted a PR that fixes this particular issue. Applying this patch on 5.4.0 fixes the regression completely. However, the timings on the latest version differ and, therefore, I marked it as draft while I investigate.

Awesome, thanks a lot for all your detailed work on this, really appreciate it!

@andreasgoulas @vasqu @zucchini-nlp – first of all, huge thanks for the deep debugging work here. The bisection down to the `.item()` removal in `_get_unpad_data` is a fantastic catch, and the draft PR (#47134) already proves the primary fix works on RTX 6000 Pro.

However, as @andreasgoulas  pointed out, timings on `main` **still** differ from the 5.2.0 baseline even with that patch applied – meaning there’s at least one secondary contributor stacked on top.

I have a clear roadmap to get this issue **fully resolved and safely merged**, and I’d like to take ownership of it:

### My immediate action plan if assigned:

1. **Validate the primary fix thoroughly**  
   - Run the `.item()` re‑addition across multiple GPUs (RTX 6000 Pro, RTX 4090, A100) to confirm parity with 5.2.0.  
   - Measure the exact overhead introduced by the device‑tensor path vs. Python scalar.

2. **Deep‑profile the root mechanism** (don’t just revert blindly)  
   - Use `torch.profiler` / `nsys` to compare both paths.  
   - Pinpoint whether the regression comes from:
     - a different kernel dispatch path inside `_index_first_axis` / `index_select`,  
     - repeated implicit synchronisations, or  
     - async kernel queue build‑up.  
   - Document the findings in the PR so future refactors don’t reintroduce the same trap.

3. **Hunt down the residual perf delta** (the unfinished business)  
   - With the `.item()` fix applied, bisect the remaining commits between 5.2.0 and current `main` that touch `modeling_flash_attention_utils.py` and the cache classes.  
   - Prime suspects: other removed `.item()`/`.cpu()` calls in the same refactor (#42435), or subtle changes to the `kwargs`‑building path for the FA2 call – the reporter already ruled out `is_tracing`, but I’ll systematically verify.

4. **Make the fix safe for `torch.compile` / graph capture**  
   - The original removal in #42435 was likely motivated by avoiding host‑device syncs under compilation.  
   - I’ll guard the `.item()` call with `torch.compiler.is_compiling()` so we get the performance fix in eager mode without breaking compiled graphs.

5. **Add a regression benchmark**  
   - Add a long‑context FA2 latency smoke test (or at least a documented perf script) to prevent this from regressing again – ideally gated in CI or added to the manual perf‑check docs.

I’ve already reviewed #47134 and #46696 (closed) to avoid duplication. I’m confident I can move #47134 out of draft, fix the remaining regression, and land this properly.

**Could one of you please assign this issue to me?** I’ll start by pushing an updated version of #47134 with the `compile`‑safe guard and the profiling results attached.

Looking forward to closing this one out 

Opened a clean PR scoped to just the fix: #47247 (closed #47215 which had unrelated commits mixed in).
@andreasgoulas — since #47134 already has your validated fix and profiling notes, I don't want to duplicate effort or create merge conflicts. Happy to close #47247 in favor of your PR, or rebase my compile-guard + benchmark additions on top of your branch — whatever's easiest for you. Let me know which you'd prefer.

Closing this out — #47134 was merged and resolves the regression. Thanks @andreasgoulas for the root-cause work and @vasqu for the review.