Repository: huggingface/transformers
Pull request: [bugfix] fix ATTN_MASK_NPU device mismatch error on multi-device NPU …

# What does this PR do?

This PR fixes a device mismatch error that occurs when using NPU flash attention functions across multiple devices. Currently, running the following test script with more than one npus

```
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "/serving_model/qwen2.5-vl-7b-instruct", torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="flash_attention_2"
)
processor = AutoProcessor.from_pretrained("/serving_model/qwen2.5-vl-7b-instruct")
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
            },
            {"type": "text", "text": "Describe this image."},
        ],
    }
]
text = processor.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
    text=[text],
    images=image_inputs,
    videos=video_inputs,
    padding=True,
    return_tensors="pt",
)
inputs = inputs.to(model.device)

generated_ids = model.generate(**inputs, max_new_tokens=128)
generated_ids_trimmed = [
    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)
```
will result in
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, npu:1 and npu:0! (when checking argument for argument atten_mask in method wrapper__npu_fusion_attention)
```

Solution: replaced the global ATTN_MASK_NPU variable with a device-specific cache


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?




Changed files:
filename: src/transformers/integrations/npu_flash_attention.py
patch: @@ -38,7 +38,14 @@
         "or 3 (down-right aligned causal mask)."
     )
 
-ATTN_MASK_NPU = None
+ATTN_MASK_NPU_CACHE = {}
+
+
+def get_attn_mask_npu(device):
+    """Get or create attention mask for the specified device."""
+    if device not in ATTN_MASK_NPU_CACHE:
+        ATTN_MASK_NPU_CACHE[device] = torch.triu(torch.ones([2048, 2048], device=device), diagonal=1).bool()
+    return ATTN_MASK_NPU_CACHE[device]
 
 
 def is_npu_fa2_top_left_aligned_causal_mask():
@@ -174,9 +181,7 @@ def npu_flash_attn_func(
         head_num = q.shape[2]
         output = torch_npu.npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=keep_prob, scale=softmax_scale)[0]
     else:
-        global ATTN_MASK_NPU
-        if ATTN_MASK_NPU is None:
-            ATTN_MASK_NPU = torch.triu(torch.ones([2048, 2048], device=q.device), diagonal=1).bool()
+        attn_mask_npu = get_attn_mask_npu(q.device)
         head_num = q.shape[2]
         output = torch_npu.npu_fusion_attention(
             q,
@@ -186,7 +191,7 @@ def npu_flash_attn_func(
             "BSND",
             keep_prob=keep_prob,
             scale=softmax_scale,
-            atten_mask=ATTN_MASK_NPU,
+            atten_mask=attn_mask_npu,
             sparse_mode=SPARSE_MODE,
         )[0]
 
@@ -227,9 +232,7 @@ def npu_flash_attn_varlen_func(
             actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()),
         )[0]
     else:
-        global ATTN_MASK_NPU
-        if ATTN_MASK_NPU is None:
-            ATTN_MASK_NPU = torch.triu(torch.ones([2048, 2048], device=q.device), diagonal=1).bool()
+        attn_mask_npu = get_attn_mask_npu(q.device)
         head_num = q.shape[1]
         output = torch_npu.npu_fusion_attention(
             q,
@@ -238,7 +241,7 @@ def npu_flash_attn_varlen_func(
             head_num,
             pse=None,
             padding_mask=None,
-            atten_mask=ATTN_MASK_NPU,
+            atten_mask=attn_mask_npu,
             scale=softmax_scale,
             keep_prob=keep_prob,
             input_layout="TND",