[performance_optim] define flash attention mask on NPU device directly

[FightingZhen]

What does this PR do?

When using Flash Attention2 on Ascend NPU, we have found that CPU memory keep increasing when calling func npu_flash_attn_varlen_func or npu_flash_attn_func.

The root cause is that the attention mask generated by func torch.ones() is initially defined on the CPU side, occupying CPU memory before being transferred to the NPU device. As the func npu_flash_attn_varlen_func or npu_flash_attn_func is called repeatedly, the CPU memory consumption continues to accumulate, which is not optimal solution. Below is one example:

transformers/src/transformers/integrations/npu_flash_attention.py

Line 225 in 12f65ee

attn_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(q.device)

Therefore, this PR is committed for solving this problem by defining attention mask tensor with torch.ones() on NPU device directy.

Fixes # (issue)
Not releated.

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?

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[FightingZhen]

@MekkCyber @SunMarc please help me review it, thanks : )

[MekkCyber]

Sounds good ! thanks for catching that 🤗

[SunMarc]

SGTM !

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

[FightingZhen]

@MekkCyber @SunMarc I have rebased the code to main branch and repushed it, I am not clear why build_pr_documentation CI pending, and the CI error seems not related with my PR, and can you please help me merge this PR? thanks : )

[FightingZhen]

@MekkCyber @SunMarc All CI passes and this PR seems to be ready for merge :)

[MekkCyber]

Merged 🎊 ! Thanks a lot