已开启
[Feature]: add NPU AMP autocast dispatch and HOP integration for flex attention #3824
liuyutong创建于 8月7日
8月7日 关联了看板：FrameworkPTAdapter 版本issue看板
8月7日 添加了label：feature
liuyutong8月7日 关联了pull request：feat: add NPU AMP autocast dispatch and HOP integration for flex attention
8月7日 关联了pull request：feat: add NPU AMP autocast dispatch and HOP integration for flex attention
8月7日 添加了label：bot-triaged
TorchNPU-Bot
8月7日 评论：
8月7日 评论：
检测到当前 issue 已关联 PR，自动添加标签：bot-triaged
liuyutong29 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
29 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
27 天前 修改了issue 的描述
liuyutong26 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
26 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
liuyutong26 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
26 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
liuyutong26 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
26 天前 关联了pull request：feat: add autocast registration for flex_attention both eager and inductor
liuyutong6 天前 关联了pull request：feat(flex_attention): autocast registration + device validation via PatchManager
6 天前 关联了pull request：feat(flex_attention): autocast registration + device validation via PatchManager
liuyutong6 天前 关联了pull request：feat(flex_attention): autocast registration + device validation via PatchManager
6 天前 关联了pull request：feat(flex_attention): autocast registration + device validation via PatchManager
提交提案之前，请先检索仓库内是否已有相同的提案，如已有请在同一提案中进行讨论。
💻 需求背景、当前现状、期望实现的功能内容、具体的设计方案、以及测试方案
FlexAttention 在 NPU 上无法使用 torch.autocast 混合精度
现象
torch.nn.attention.flex_attention在 NPU 设备上，当torch.autocast(device_type="npu", dtype=torch.bfloat16)激活时，抛出：根因
AutocastCUDA/AutocastCPU注册 flex_attention HOP 的 autocast kernel，AutocastPrivateUse1（NPU）未注册_validate_device为空实现，未对 NPU 设备做校验修复
AutocastPrivateUse1HOP kernel影响
NPU 上大模型可在eager模式和GPU行为保持一致，然后eager和inductor模式都能正常使用混合精度训练/推理。
替代方案
补充说明
欢迎加入社区，感谢您对社区的贡献 🎉!