npu_flash_attn_utils.py — NPU FlashAttention 适配¶
文件路径: verl/utils/npu_flash_attn_utils.py
文件概述¶
为华为 Ascend NPU 提供 FlashAttention 相关工具函数的适配实现。在 NVIDIA GPU 上使用 flash_attn 库，NPU 上使用 transformers 库内置的 NPU attention 实现。
核心功能¶
提供与 flash_attn.bert_padding 相同接口的函数：
- unpad_input(): 去除 padding
- pad_input(): 恢复 padding
- index_first_axis(): 索引操作
- rearrange(): 张量重排
与其他模块的关系¶
- 被
attention_utils.py在检测到 NPU 环境时导入 - 提供与 CUDA 版本完全相同的接口
小结¶
NPU 设备上 FlashAttention 工具函数的替代实现，确保跨硬件兼容。