From e2d30d0a37a784aa019711364d0ec52c12b01d21 Mon Sep 17 00:00:00 2001
From: chen_liqing <jiangyiwen5@huawei.com>
Date: Mon, 18 Aug 2025 09:58:01 +0800
Subject: [PATCH] Consider tq=2 in reselect_static_kernel

---
 torch_npu/csrc/npu/Module.cpp | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/torch_npu/csrc/npu/Module.cpp b/torch_npu/csrc/npu/Module.cpp
index 21151d7e34..f61dc2e2b3 100644
--- a/torch_npu/csrc/npu/Module.cpp
+++ b/torch_npu/csrc/npu/Module.cpp
@@ -1613,7 +1613,20 @@ PyObject* THNPModule_aclnn_reselect_static_kernel(PyObject* self, PyObject* noar
     HANDLE_TH_ERRORS
     NPUStatus ret = c10_npu::emptyAllNPUStream();
     TORCH_CHECK(ret == SUCCESS, "Failed to empty NPU task queue, ret:", ret, PTA_ERROR(ErrCode::INTERNAL));
-    c10_npu::opapi::ReselectStaticKernel();
+
+    static const auto task_queue_enable = c10_npu::option::OptionsManager::GetTaskQueueEnable();
+    if (task_queue_enable == 2) {
+        auto acl_call = []()->int {
+            c10_npu::opapi::ReselectStaticKernel();
+            return 0;
+        };
+        at_npu::native::OpCommand::RunOpApiV2("reselect_static_kernel", acl_call);
+        NPUStatus ret = c10_npu::emptyAllNPUStream();
+        TORCH_CHECK(ret == SUCCESS, "Failed to empty NPU task queue, ret:", ret, PTA_ERROR(ErrCode::INTERNAL));
+    } else {
+        c10_npu::opapi::ReselectStaticKernel();
+    }
+
     Py_RETURN_NONE;
     END_HANDLE_TH_ERRORS
 }
-- 
Gitee