openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=40031&platform=gitcode&codeHostingPlatformFlag=gitcode

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

💡 Tip:

• Committer can comment /approve or /lgtm
• Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

ascend-ds-bot, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend docs pipeline is running...

✅ 跳过 docs ci 检查，没有需要检查的文档文件

• compile、compile_inductor、compile_torchair : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

变更摘要

此 PR 修复了 profiler FLOPs 计算中的一个 bug：在融合注意力（fusion attention）场景下，npu_fusion_attention_flops 之前错误地将 head_num（查询头数）直接作为 KV 头数使用，未考虑 GQA/MQA 等场景下 KV 头数与查询头数不一致的情况。同时新增了 tx_enabled 属性，使 _msprof_tx 和 _mstx 两种 tx 启用方式统一判断，避免遗漏。

主要改动

• 新增 _infer_kv_heads 函数：在 _flops_formulas.py 中新增该函数，根据布局类型（BNSD/BSND/BSH/SBH）从 K 张量形状中推断实际的 KV 头数，并通过 q_head_dim 反推 kv_heads，解决 GQA/MQA 场景下 KV 头数与 Q 头数不一致的 FLOPs 计算偏差。
• 修正 npu_fusion_attention_flops 的 KV 头数参数：将 _calculate_common_layout_flops 调用中的 kv_heads 参数从固定的 head_num 改为通过 _infer_kv_heads 推断的值，确保非 TND 布局下 FLOPs 计算使用正确的 KV 头数。
• 新增 tx_enabled 属性：在 _ExperimentalConfig 类中增加 tx_enabled 属性，统一检查 _msprof_tx 与 _mstx 两个标志位，任一为真即表示 tx 已启用。
• 统一 tx 启用判断逻辑：在 profiler_interface.py 的 start 和 stop 方法中，将原有的 self.experimental_config._msprof_tx 直接访问替换为 self.experimental_config.tx_enabled，确保 FlopsHookManager 的安装/卸载条件覆盖所有 tx 启用路径。

代码审查

经过全面审查所有三个变更文件，我对本次 diff 进行了深入分析：

文件 1: torch_npu/profiler/_flops_formulas.py

变更内容：

• 新增 _infer_kv_heads 函数（第 428-443 行），从 k_shape 推断 KV head 数量
• 在 npu_fusion_attention_flops 中调用该函数（第 36 行），并将结果传递给 _calculate_common_layout_flops（第 38 行）

审查结论： 逻辑正确。

_infer_kv_heads 正确覆盖了四种布局：

• BNSD: 直接从 k_shape[1] 提取 head 数量 ✓
• BSND: 直接从 k_shape[2] 提取 head 数量 ✓
• BSH/SBH: 通过计算 k_hidden / (q_hidden / q_heads) 推断 KV head 数量 ✓
• 未知布局: 回退到 q_heads（保持旧行为），后续 _parse_dims 会因未识别布局抛出 ValueError ✓

该修复正确处理了 GQA（分组查询注意力）场景，之前 head_num 被同时作为 Q 和 K/V 的 head 数量传入，当两者不一致时会导致 ValueError: Hidden size K must be divisible by the number of heads Q。

文件 2: torch_npu/profiler/experimental_config.py

变更内容： 新增 tx_enabled 属性（第 151-153 行），返回 self._msprof_tx or self._mstx。

审查结论： 逻辑正确。

• 该属性与 __call__ 方法中已存在的 msprof_tx=self._msprof_tx or self._mstx 逻辑一致 ✓
• _msprof_tx 和 _mstx 均在 __init__ 中初始化（默认 False），不会出现 AttributeError ✓

文件 3: torch_npu/profiler/profiler_interface.py

变更内容： 将 self.experimental_config._msprof_tx 替换为 self.experimental_config.tx_enabled（第 159 和 168 行）。

审查结论： 逻辑正确。

• 旧代码直接访问私有属性 _msprof_tx，当仅设置 mstx=True 时不会安装 FlopsHookManager
• 新代码通过 tx_enabled 同时检查两个标志，与 C++ 配置逻辑一致 ✓
• tx_enabled 作为 property 封装了内部实现，提高了可维护性 ✓

总结

整体风险评估：低风险。 变更精准地修复了 GQA 场景下 FLOPs 计算的 bug，同时修复了 mstx 参数未被 FlopsHookManager 识别的遗漏。所有逻辑经逐路径追踪验证正确，无新增缺陷。

已审查的文件：

• torch_npu/profiler/_flops_formulas.py — 无问题
• torch_npu/profiler/experimental_config.py — 无问题
• torch_npu/profiler/profiler_interface.py — 无问题

⚠️ 已识别出整体风险，但无法提取行内评论，请参考整体评估。

/approve

/approve

The following label is not ready.

ci-pipeline-passed: The ci-pipeline-passed label is expired. Please compile again.

compile

ascend docs pipeline is running...

✅ 跳过 docs ci 检查，没有需要检查的文档文件

compile

ascend docs pipeline is running...

✅ 跳过 docs ci 检查，没有需要检查的文档文件

compile

ascend docs pipeline is running...

✅ 跳过 docs ci 检查，没有需要检查的文档文件

• compile、compile_inductor、compile_torchair : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线