[remi-or]

This PR refactors two functions present in modeling_flash_attention_utils.py, not to change their behaviour but to make them more readable. It also fixes one bug.

_lazy_imports

• introduce fa_fallback_version to make the fallback system more readable: we fall back to the highest version of FA avaiable
• collapse the if/elif/else chain in that function: it used to be that in the else, there was another if/elif/else branch. Now everything is in the same branch
• moved the NPU-specific block before the FA block, where we try FA2, 3 and 4. First, it seems more logical with the block comment. Second, unless mutilples device types are available, this is does not change the behavior because those branches are exclusive: if is_npu_available is True, is_flash_attention_x_available if False
• added a tiny tuple[Callable, ...] return type hint

_flash_attention_forward

This is the change that matters the most IMO. This function is mainly about preparing the sequence-defining attributes (cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k) before calling the right flavor of flash. Currently, there are 4 cases:

1. there is an attention_mask, which we use to unpad the sequences and compute the sequence-defining attributes
2. there is no attention mask but the sequence-defining attributes are already passed
3. there is no attention mask, no sequence-defining attributes passed, and the position_ids show that the sequence is packed, therefore we compute the sequence-defining attributes from positions_ids
4. there is no attention mask, no sequence-defining attributes passed, and the sequence are not packed. In this case and this case only we run flash and not flash_varlen

Cases 1, 2 and 3 only differ in the way they pre-process and post-process the input / output of flash_varlen. Yet the currently, the main body of those cases is not shared: we duplicate some mps check, the call to the function, and checking for tuple.
Furthermore, the check for case 3. is always run. If positions_ids is passed, the check is expensive: GPU ops + D2H sync (incompatible with cuda graph), while totally avoidable if we are in case 1 or 2.
Finally, there is a "hole" in the checks: though we check if all sequence_defining attributes are passed, we only compute them from position_ids if the cu_seqlens are missing: this means that if a user does not pass max_length, the flash_varlen call will crash.

To fix all this:

• we separate the pre-processing, call to flash and post-processing into 3 different blocks. The only special case is 4.: since there is not pre or post processing to do, we return early in the pre-processing block
• we stop pre-computing is_fa_with_position_ids and only compute it case 1. and 2. fail.
• we stop repeating the mps check, call flash_varlen and tuple stripping block
• there is a tiny rename inside the function scope from flash_kwargs to flash_kwargs_fn because it's not a dict and it makes everything fit to one line when we call it

Side note: the diff on GH is quite ugly, I would recommend checking it out with Gitlens or smt. We can see there is much less change there.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vasqu]

Not completely opposed just wondering in combination with my other FA PR and how that fits in

[vasqu]

Thinking of #45153 and how this would play into your refactor here because we would need to support base FA via varlen (it's less accurate using varlen which is why I dont want to drop the base FA completely)

Also opinionated but how is this more readable? I know before we repeated quite a bit but now we have quite a few branches we need to remember like the last pad on the attention mask 😅 maybe we can adopt one function per scenario as in the other PR so we see what does what?

[remi-or]

we would need to support base FA via varlen (it's less accurate using varlen which is why I dont want to drop the base FA completely)

We still do? For that path the behavior has not changed, we still run base FA and exit if this is not varlen

Also opinionated but how is this more readable?

Well we just have a block to process the sequence defining variables, then the call, then the post processing. On the opposite I think it would make the code less readable to have 3 fused paths (4 with paged attention) and it will lead to repetitions inside the code. Plus we will still pass a lot of duplicate kwargs around. I think an if / else on the out is ok for users!

[vasqu]

We still do? For that path the behavior has not changed, we still run base FA and exit if this is not varlen

No I think you misunderstood, torch varlen only has FA varlen so there is no base FA and we need a workaround for that

[remi-or]

Ha ok! I don't understand how it relates to this PR though, sorry

[vasqu]

Hmm I was trying to be preemptive here: So we would introduce another prepare fa varlen path right? But yea no worries

[remi-or]

So we would introduce another prepare fa varlen path right

Depends on what arguments are given! I guess the torch varlen will have similar arguments, no?

[vasqu]

Yep we basically dummy create them because you can just infer them based on the bsz / seq len but should be easy to integrate 🫡

[vasqu]

Should be fine, I dont want to fight against it and I feel like im just becoming an opinionated slug :D

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 36857242507:1
Result: success | Jobs: 16 | Tests: 178,870 | Failures: 0 | Duration: 13h 45m