Fix flash_attention.py: wrong argument passing for attn_implementation

[TKONIY]

What does this PR do?

The name of the attn type argument for _flash_attention_forward() should be implementation, instead of attn_implementation which currently uses in the function call. This would result in wrong type specification.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

CC

@vasqu @ArthurZucker @Cyrilvallez

[Cyrilvallez]

Oh good catch there is indeed the wrong kwarg being passed on! I would prefer if you can change the name of the kwarg in _flash_attention_forward though, it would be more explicit!

[TKONIY]

Oh good catch there is indeed the wrong kwarg being passed on! I would prefer if you can change the name of the kwarg in _flash_attention_forward though, it would be more explicit!

Thank you. I have changed that. Please check.

[vasqu]

Can you also fix the docs naming then at

transformers/src/transformers/modeling_flash_attention_utils.py

Lines 567 to 568 in 0452f28

	implementation (`str`, *optional*):
	The attention implementation to use. If None, will default to the one based on the environment.

LGTM otherwise 🤗

[TKONIY]

Can you also fix the docs naming then at

transformers/src/transformers/modeling_flash_attention_utils.py

Lines 567 to 568 in 0452f28

	implementation (`str`, *optional*):
	The attention implementation to use. If None, will default to the one based on the environment.

LGTM otherwise 🤗

Thank you! Done!

[Cyrilvallez]

Thanks!! 🤗

[ArthurZucker]

https://github.com/huggingface/transformers/blob/update-from-pretrained/src/transformers/integrations/hub_kernels.py#L214-L214 is where we use implementation. If you do this it won't fallback to kernels, we need to make sure we use the one passed in load and register

[TKONIY]

https://github.com/huggingface/transformers/blob/update-from-pretrained/src/transformers/integrations/hub_kernels.py#L214-L214 is where we use implementation. If you do this it won't fallback to kernels, we need to make sure we use the one passed in load and register

So would it be better if I rollback to the fix that simply change attn_implementation= to implementation=？

[vasqu]

The link is broken, I think Arthur meant

transformers/src/transformers/integrations/hub_kernels.py

Lines 214 to 215 in caa14e7

	kernel_function = partial(attention_wrapper, implementation=kernel)
	lazy_import_flash_attention(kernel, force_import=True)

Yes, we should also change the kwarg there, totally forgot there, e.g. kernel_function = partial(attention_wrapper, attn_implementation=kernel) . But it's a bit messier tbh and I don't think we actually use the kwarg much at all except on first call (which would happen if someone does something custom with fa interface we have) --> the forced lazy import should load the correct kernel (checking in a second) and as we already loaded it, we never change it again there.

Edit: Still loads the correct kernel implementation, checked with kernels-community/flash-attn3