[FA] Cleanup loading logic

[vasqu]

We currently ignore anything passed through the interface and forced correct loading at attn setting time. This changes it to be more flexible where power users may customly switch it up by just setting it in the config.

[vasqu]

This didn't have any effect currently and we only used the forced loading to get the correct implementation

[ArthurZucker]

no this is super important for paged_attention wrappers 😓

[vasqu]

Changed the logic so that CB is now able to use all fa versions with lazy loading, i.e. fa2, fa3, kernels fas

[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]

Lazy import only happens at init or explicit setting time

transformers/src/transformers/modeling_utils.py

Lines 2580 to 2582 in 766ac18

	# preload flash attention here to allow compile with fullgraph
	if "flash" in applicable_attn_implementation:
	lazy_import_flash_attention(applicable_attn_implementation)

transformers/src/transformers/generation/continuous_batching/continuous_api.py

Lines 614 to 618 in ced85ca

	# lazy loading flash attention including kernel variations
	if "flash" in attn_implementation:
	from ...modeling_flash_attention_utils import lazy_import_paged_flash_attention
	lazy_import_paged_flash_attention(attn_implementation)

Otherwise, it can happen when users only set the attention within the config, i.e. the interface detects a different implementation name than the currently loaded one

transformers/src/transformers/modeling_flash_attention_utils.py

Lines 143 to 144 in 766ac18

	if implementation is not None and _loaded_implementation != implementation:
	_loaded_implementation = implementation

[ArthurZucker]

make sure you test the continious_batching_simple script with fa2 !

[ArthurZucker]

no this is super important for paged_attention wrappers 😓

[vasqu]

@ArthurZucker updated the logic a bit to include CB now as well. The loaded attention is now properly dependent on the config like in the base fa version + we can now use fa3 as well for CB.

[ArthurZucker]

If we do:

- partial(attention_wrapper, implementation=kernel)

in that case we should remove this arg from all xxx_paged

[vasqu]

Seems like the implementation kwarg was only used in flash_paged, the other implementations don't have it:

• eager

  transformers/src/transformers/integrations/eager_paged.py

  Lines 21 to 29 in a99b1be

  	def eager_paged_attention_forward(
  	module: nn.Module,
  	query: torch.Tensor,
  	key: torch.Tensor,
  	value: torch.Tensor,
  	attention_mask: Optional[torch.Tensor], # shape [seqlen_q, seqlen_k]
  	scaling: float,
  	**kwargs,
  	):

• sdpa

  transformers/src/transformers/integrations/sdpa_paged.py

  Lines 20 to 29 in a99b1be

  	def sdpa_attention_paged_forward(
  	module: torch.nn.Module,
  	query: torch.Tensor,
  	key: torch.Tensor,
  	value: torch.Tensor,
  	attention_mask: Optional[torch.Tensor],
  	dropout: float = 0.0,
  	scaling: Optional[float] = None,
  	**kwargs,
  	) -> tuple[torch.Tensor, None]:

• flash removed in this PR

  transformers/src/transformers/integrations/flash_paged.py

  Lines 9 to 21 in 1ece193

  	def paged_attention_forward(
  	module: torch.nn.Module,
  	q: torch.Tensor,
  	k: torch.Tensor,
  	v: torch.Tensor,
  	attention_mask: Optional[torch.Tensor] = None,
  	cache: PagedAttentionCache = None,
  	cu_seq_lens_q=None,
  	cu_seq_lens_k=None,
  	max_seqlen_q=None,
  	max_seqlen_k=None,
  	**kwargs,
  	) -> torch.Tensor:

  
  So we should be good @ArthurZucker?