[Bug]: ACLNN 执行回调误用 enqueue 状态导致切流后控核资源绑定错误-ops-transformer-AtomGit

已关闭

[Bug]: ACLNN 执行回调误用 enqueue 状态导致切流后控核资源绑定错误 #5452

PerrySkywalker创建于 23 天前关闭于 19 天前

23 天前创建

问题描述

torch_extension/cann_ops_transformer/common/aclnn_common.h 的 ACLNN_CMD 在 acl_call 执行回调内部调用 check_enqueue_need_use。外层入队准备与内层执行回调需要分别维护 enqueue/dequeue 的 stream 资源绑定状态；混用后，跨 stream 执行可能漏绑资源，导致算子使用上一个 stream 的核数限制。

复现与实际结果

系统 Python、Ascend950，TASK_QUEUE_ENABLE=1。用 C++ 最小回调通过 OpCommand::SetCustomHandler / Run 执行 A→B→A：

stream A 设置 Cube/Vector=8/16，回调使用现有 enqueue 判断。
stream B 设置 Cube/Vector=4/8，回调使用标准 dequeue 判断，模拟与原生 op-plugin 算子混用。
使用 aclrtGetResInCurrentThread 读取消费线程实际绑定资源。

首次 A 为 8/16，B 为 4/8；再次 A 时 enqueue 判断返回 false，实际仍为 4/8，期望为 8/16。外层与回调线程 ID 不同。

将执行回调统一改为 dequeue 判断后，A→B→A 的实际资源为 8/16→4/8→8/16。TASK_QUEUE_ENABLE=0 的同步对照也通过。

修复方案

仅将 acl_call 内部的 check_enqueue_need_use 改为 check_dequeue_need_use，保留外层 enqueue 判断以及现有 aclrtUseStreamResInCurrentThread 调用。同步修复 master 和 9.2.0。

回归验证

已修改实际安装包头文件，并在系统 Python 下重新编译 flash_mla_with_kvcache Torch 扩展：

8 Cube / 16 Vector，BF16，aclgraph + fullgraph=True，CPU 对照通过，最大绝对误差 0.00203973。
16 Cube / 32 Vector，FP16、causal、batch=2、Q长度=2、KV长度=257，aclgraph + fullgraph=True + static_kernel_compile=True，CPU 对照通过，最大绝对误差 0.000186145；确认新生成的 .run 包包含 FlashMlaWithKvcache .o。
两个 stream 按 8/16→4/8→8/16 切换，混合原生 torch.add 与 MLA，三次 CPU 对照均通过，最大绝对误差 0.00029558。

两个目标分支各提交同一处一行修复，分别通过 staged 和完整 PR 差异的 pre-commit（含 OAT）检查。上述上卡测试使用本机安装包重编译，不代表两个分支各自全量构建或远端 CI 已通过。