Repository: huggingface/transformers
Pull request: Fixed paged|FA2 kernel loading logic and UT.

# What does this PR do?

**1. [General]** When running the test case `tests/generation/test_continuous_batching.py::ContinuousBatchingTest::test_continuous_batching_parity_qwen_flash`, it reports error as follows: 
```
                flash_attn_func = getattr(kernel, "flash_attn_func", None)
                flash_attn_varlen_func = getattr(kernel, "flash_attn_varlen_func", None)
                if flash_attn_varlen_func is None:
>                   raise ValueError(
                        f"Could not find the currently requested flash attention implementation at `{implementation}`."
                        "Make sure that you request a valid kernel from the hub, e.g. `kernels-community/flash-attn2`."
                    )
E                   ValueError: Could not find the currently requested flash attention implementation at `flash_attention_2`.Make sure that you request a valid kernel from the hub, e.g. `kernels-community/flash-attn2`.

src/transformers/modeling_flash_attention_utils.py:115: ValueError
```
The root cause is a kernel loading failure of `paged|flash_attention_2`. This PR fixes the issue.

**2. [XPU]** The above test case reports error `E   AssertionError: Test request_id = 'req_1' failed, no expected output was provided.....` after successfully running on both XPU and CUDA. However, on the XPU side, the test case can stably pass by using `require_deterministic_for_xpu`.

For enabling the related tests on XPU, please refer to PR https://github.com/huggingface/transformers/pull/42536.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -763,15 +763,9 @@ def __init__(
             num_kv_padding_intervals: (optional) Number of intervals used to pad the keys/values dimension
             allow_prefix_sharing: (optional) Whether to allow prefix sharing if the model has only full attention layers
         """
+        # Reloade paged version if necessary
         if "paged|" not in model.config._attn_implementation:
-            attn_implementation = f"paged|{model.config._attn_implementation}"
-            model.config._attn_implementation = attn_implementation
-
-            # 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)
+            model.set_attn_implementation(f"paged|{model.config._attn_implementation}")
 
         self.model = model.eval()
         generation_config = model.generation_config if generation_config is None else generation_config

filename: src/transformers/modeling_utils.py
patch: @@ -85,7 +85,7 @@
     verify_tp_plan,
 )
 from .loss.loss_utils import LOSS_MAPPING
-from .modeling_flash_attention_utils import lazy_import_flash_attention
+from .modeling_flash_attention_utils import lazy_import_flash_attention, lazy_import_paged_flash_attention
 from .pytorch_utils import id_tensor_storage
 from .quantizers import HfQuantizer
 from .quantizers.auto import get_hf_quantizer
@@ -1763,9 +1763,12 @@ def _check_and_adjust_attn_implementation(
         """
         applicable_attn_implementation = attn_implementation
 
+        is_paged = attn_implementation is not None and attn_implementation.startswith("paged|")
+
         # If FA not installed, do not fail but use kernels instead
         requested_original_flash_attn = attn_implementation is not None and (
-            attn_implementation == "flash_attention_2" or attn_implementation == "flash_attention_3"
+            attn_implementation.removeprefix("paged|") == "flash_attention_2"
+            or attn_implementation.removeprefix("paged|") == "flash_attention_3"
         )
         if (
             requested_original_flash_attn
@@ -1783,10 +1786,16 @@ def _check_and_adjust_attn_implementation(
             else:
                 applicable_attn_implementation = "kernels-community/vllm-flash-attn3"
 
+            if is_paged:
+                applicable_attn_implementation = f"paged|{applicable_attn_implementation}"
+
         if is_kernel(applicable_attn_implementation):
             try:
                 # preload flash attention here to allow compile with fullgraph
-                lazy_import_flash_attention(applicable_attn_implementation)
+                if is_paged:
+                    lazy_import_paged_flash_attention(applicable_attn_implementation)
+                else:
+                    lazy_import_flash_attention(applicable_attn_implementation)
 
                 # log that we used kernel fallback if successful
                 if requested_original_flash_attn:

filename: tests/generation/test_continuous_batching.py
patch: @@ -22,6 +22,7 @@
 from transformers.generation.continuous_batching.continuous_api import build_attention_mask
 from transformers.testing_utils import (
     Expectations,
+    require_deterministic_for_xpu,
     require_kernels,
     require_read_token,
     require_torch_accelerator,
@@ -137,6 +138,7 @@ def test_attention_mask(
                 f"Actual mask:\n{str_mask}"
             )
 
+    @require_deterministic_for_xpu
     def _continuous_batching_parity(
         self, model_id: str, attn_implementation: str, expected_outputs: dict[str, str]
     ) -> None: