From 64b229b04bbd4833e952131c6bad057795432fd8 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 03:59:40 +0000
Subject: [PATCH 01/11] Add Python-level inject_oom for inference OOM testing

Provide a single top-level OutOfMemoryError injection point for workloads
such as Qwen 235B that bypass C++ allocator and HCCL hooks. Controlled via
NPU_INJECT_OOM_STEP and a call at the top of each inference iteration.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/npu/__init__.py  |  3 ++-
 torch_npu/npu/_recovery.py | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/torch_npu/npu/__init__.py b/torch_npu/npu/__init__.py
index c8acaf21f7..f4e590276c 100644
--- a/torch_npu/npu/__init__.py
+++ b/torch_npu/npu/__init__.py
@@ -100,6 +100,7 @@
     "current_blas_handle",
     "stop_device",
     "restart_device",
+    "inject_oom",
     "check_uce_in_memory",
     "config",
     "matmul",
@@ -159,7 +160,7 @@
                     get_sync_debug_mode, init_dump, current_blas_handle, is_bf16_supported,
                     finalize_dump, set_dump, get_npu_overflow_flag, clear_npu_overflow_flag,
                     check_uce_in_memory, stress_detect, _get_uce_addr, ipc_collect, set_op_timeout_ms)
-from ._recovery import restart_device, stop_device
+from ._recovery import restart_device, stop_device, inject_oom
 from .streams import Stream, Event, SyncLaunchStream, ExternalEvent
 from .mstx import mstx
 from .npu_config import *  # noqa: F403
diff --git a/torch_npu/npu/_recovery.py b/torch_npu/npu/_recovery.py
index 00d983d2cd..0422030a66 100644
--- a/torch_npu/npu/_recovery.py
+++ b/torch_npu/npu/_recovery.py
@@ -1,3 +1,4 @@
+import os
 import torch
 from torch.distributed.distributed_c10d import _pg_map
 
@@ -54,6 +55,20 @@ def _recovery_all_npu_stream(device: int) -> None:
     return torch_npu._C._recovery_all_npu_stream(device)
 
 
+def inject_oom(step: int = 1, trigger_step: int = None, device: int = None) -> None:
+    if trigger_step is None:
+        trigger_step = int(os.getenv("NPU_INJECT_OOM_STEP", "0"))
+    if trigger_step <= 0 or step != trigger_step:
+        return
+    if device is None:
+        device = torch_npu.npu.current_device()
+    raise torch.OutOfMemoryError(
+        f"NPU out of memory. Injected OOM on NPU {device}. "
+        f"Recover with: torch_npu.npu.restart_device({device}, rebuild_all_resources=True). "
+        + pta_error(ErrCode.MEMORY)
+    )
+
+
 def restart_device(device_id: int, rebuild_all_resources: int = False):
     torch_npu.npu._lazy_init()
     if rebuild_all_resources:

From f70d2505c384340f3b53a961b9dbfa0a559ca719 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 04:34:50 +0000
Subject: [PATCH 02/11] Add automatic PTA OOM injection at upper C++ layers

Remove Python inject_oom API. Introduce maybeThrowPtaOom() controlled by
PTA_OOM_TRIGGER_COUNT (default 6000), called from MakeSureQueueEmpty (async
queue drain, where OOM propagates upward) and malloc() entry. Skips graph
capture. Raises torch.OutOfMemoryError with PTA memory fault code.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 .../csrc/core/npu/NPUCachingAllocator.cpp     |  8 ++--
 torch_npu/csrc/core/npu/NPUException.cpp      | 38 +++++++++++++++++++
 torch_npu/csrc/core/npu/NPUException.h        |  2 +
 torch_npu/csrc/core/npu/NPUQueue.cpp          |  4 ++
 torch_npu/npu/__init__.py                     |  3 +-
 torch_npu/npu/_recovery.py                    | 15 --------
 6 files changed, 49 insertions(+), 21 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
index 783d27969a..edd25c9227 100644
--- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
+++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
@@ -107,7 +107,6 @@ const std::string kMinDriverVersion = "25.0.RC1";     // minimum driver version
 const std::string kCannModule = "CANN";               // cann module name
 constexpr int kPrecision = 4;                         // precision of the memory usage information
 constexpr size_t kLazyQuerySize = 512;                // lazy query event size
-static int64_t g_malloc_call_count = 0;
 static char SHAREABLE_HANDLE_VERSION = 1;
 enum ShareableHandleType : char {
     SHAREABLE_NPU_MALLOC = 'c',
@@ -1152,9 +1151,10 @@ class DeviceCachingAllocator {
     // Thus, do not call a public method from another public method.
 
     Block *malloc(int device, size_t orig_size, aclrtStream stream, uint8_t allocator_type = 0)
-    {           
-        g_malloc_call_count++;
-        auto retmsg = std::string("NPU out of memory. Tried to allocate more than 1EB memory.");
+    {
+        if (currentStreamCaptureStatus() == CaptureStatus::None) {
+            maybeThrowPtaOom("malloc");
+        }
 
         TORCH_NPU_MEMORY_LOGD("Allocating memory: size=%zu, device=%d", orig_size, device);
         // done outside the lock because we don't know what locks the recorder needs
diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index e9f8dd356c..5ac501a7be 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -1,3 +1,6 @@
+#include <atomic>
+#include <cstdlib>
+
 #include "torch_npu/csrc/core/npu/NPUException.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
 #include "torch_npu/csrc/core/npu/NPUStream.h"
@@ -433,4 +436,39 @@ bool isCannOOM(const std::string &errMsg)
     return false;
 }
 
+namespace {
+static std::atomic<int64_t> g_pta_oom_call_count{0};
+static constexpr int64_t kDefaultPtaOomTriggerCount = 6000;
+
+int64_t getPtaOomTriggerCount()
+{
+    const static int64_t trigger_count = []() -> int64_t {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_COUNT");
+        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultPtaOomTriggerCount;
+    }();
+    return trigger_count;
+}
+} // namespace
+
+void maybeThrowPtaOom(const char *context)
+{
+    const int64_t trigger_count = getPtaOomTriggerCount();
+    if (trigger_count <= 0) {
+        return;
+    }
+
+    const int64_t current_count = ++g_pta_oom_call_count;
+    if (current_count > trigger_count && current_count < trigger_count + 2) {
+        auto retmsg = std::string("NPU out of memory. Injected PTA OOM after ") +
+            std::to_string(current_count) + " PTA operations";
+        if (context != nullptr && context[0] != '\0') {
+            retmsg += ", context is ";
+            retmsg += context;
+        }
+        retmsg += ". ";
+        retmsg += PTA_ERROR(ErrCode::MEMORY);
+        TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
+    }
+}
+
 } // namespace c10_npu
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index 987d9fa6a1..066ddbe3bc 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -336,6 +336,8 @@ std::string handleSuspectRemoteError(int errorCode);
 
 bool isCannOOM(const std::string &errMsg);
 
+void maybeThrowPtaOom(const char *context = nullptr);
+
 bool ShouldAppendDeviceErrorVerbose();
 
 void clear_device_error_info();
diff --git a/torch_npu/csrc/core/npu/NPUQueue.cpp b/torch_npu/csrc/core/npu/NPUQueue.cpp
index 4d820bbf9b..dd9b46a956 100644
--- a/torch_npu/csrc/core/npu/NPUQueue.cpp
+++ b/torch_npu/csrc/core/npu/NPUQueue.cpp
@@ -8,6 +8,7 @@
 #include "torch_npu/csrc/framework/OpCommand.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/logging/LogContext.h"
 
 #ifndef BUILD_LIBTORCH
@@ -348,6 +349,9 @@ NPUStatus Repository::MakeSureQueueEmpty(bool check_error)
             throw std::runtime_error(runtime_error);
         }
     }
+    if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
+        maybeThrowPtaOom("MakeSureQueueEmpty");
+    }
     logger->debug("MakeSureQueueEmpty: clearing successful, device = %d, write_idx = %u, read_idx = %u, status = %d",
         device_idx, write_idx.idx, read_idx.idx, GetStatus());
 
diff --git a/torch_npu/npu/__init__.py b/torch_npu/npu/__init__.py
index f4e590276c..c8acaf21f7 100644
--- a/torch_npu/npu/__init__.py
+++ b/torch_npu/npu/__init__.py
@@ -100,7 +100,6 @@
     "current_blas_handle",
     "stop_device",
     "restart_device",
-    "inject_oom",
     "check_uce_in_memory",
     "config",
     "matmul",
@@ -160,7 +159,7 @@
                     get_sync_debug_mode, init_dump, current_blas_handle, is_bf16_supported,
                     finalize_dump, set_dump, get_npu_overflow_flag, clear_npu_overflow_flag,
                     check_uce_in_memory, stress_detect, _get_uce_addr, ipc_collect, set_op_timeout_ms)
-from ._recovery import restart_device, stop_device, inject_oom
+from ._recovery import restart_device, stop_device
 from .streams import Stream, Event, SyncLaunchStream, ExternalEvent
 from .mstx import mstx
 from .npu_config import *  # noqa: F403
diff --git a/torch_npu/npu/_recovery.py b/torch_npu/npu/_recovery.py
index 0422030a66..00d983d2cd 100644
--- a/torch_npu/npu/_recovery.py
+++ b/torch_npu/npu/_recovery.py
@@ -1,4 +1,3 @@
-import os
 import torch
 from torch.distributed.distributed_c10d import _pg_map
 
@@ -55,20 +54,6 @@ def _recovery_all_npu_stream(device: int) -> None:
     return torch_npu._C._recovery_all_npu_stream(device)
 
 
-def inject_oom(step: int = 1, trigger_step: int = None, device: int = None) -> None:
-    if trigger_step is None:
-        trigger_step = int(os.getenv("NPU_INJECT_OOM_STEP", "0"))
-    if trigger_step <= 0 or step != trigger_step:
-        return
-    if device is None:
-        device = torch_npu.npu.current_device()
-    raise torch.OutOfMemoryError(
-        f"NPU out of memory. Injected OOM on NPU {device}. "
-        f"Recover with: torch_npu.npu.restart_device({device}, rebuild_all_resources=True). "
-        + pta_error(ErrCode.MEMORY)
-    )
-
-
 def restart_device(device_id: int, rebuild_all_resources: int = False):
     torch_npu.npu._lazy_init()
     if rebuild_all_resources:

From 477da6ff6e3f2f8490b1947de9dec0c157244f20 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 04:37:26 +0000
Subject: [PATCH 03/11] Make PTA OOM injection full-card fault

On trigger, mark all allocator blocks on the device unsafe and enable
unsafe-data checking before throwing OutOfMemoryError. Reset inject state
on restart_device and clear unsafe check flag after recovery.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 .../csrc/core/npu/NPUCachingAllocator.cpp     |  2 +-
 torch_npu/csrc/core/npu/NPUException.cpp      | 47 ++++++++++++++-----
 torch_npu/csrc/core/npu/NPUException.h        |  4 +-
 torch_npu/csrc/core/npu/NPUQueue.cpp          |  2 +-
 torch_npu/csrc/npu/Module.cpp                 |  1 +
 torch_npu/npu/_recovery.py                    |  2 +
 6 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
index edd25c9227..d5143d5135 100644
--- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
+++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
@@ -1153,7 +1153,7 @@ class DeviceCachingAllocator {
     Block *malloc(int device, size_t orig_size, aclrtStream stream, uint8_t allocator_type = 0)
     {
         if (currentStreamCaptureStatus() == CaptureStatus::None) {
-            maybeThrowPtaOom("malloc");
+            maybeThrowPtaOom("malloc", device);
         }
 
         TORCH_NPU_MEMORY_LOGD("Allocating memory: size=%zu, device=%d", orig_size, device);
diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 5ac501a7be..8bd3d78b3e 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -2,7 +2,9 @@
 #include <cstdlib>
 
 #include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
+#include "torch_npu/csrc/core/npu/NPURecovery.h"
 #include "torch_npu/csrc/core/npu/NPUStream.h"
 #include "torch_npu/csrc/core/npu/NpuVariables.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
@@ -438,6 +440,7 @@ bool isCannOOM(const std::string &errMsg)
 
 namespace {
 static std::atomic<int64_t> g_pta_oom_call_count{0};
+static std::atomic<bool> g_pta_oom_injected{false};
 static constexpr int64_t kDefaultPtaOomTriggerCount = 6000;
 
 int64_t getPtaOomTriggerCount()
@@ -450,7 +453,7 @@ int64_t getPtaOomTriggerCount()
 }
 } // namespace
 
-void maybeThrowPtaOom(const char *context)
+void maybeThrowPtaOom(const char *context, int device)
 {
     const int64_t trigger_count = getPtaOomTriggerCount();
     if (trigger_count <= 0) {
@@ -458,17 +461,39 @@ void maybeThrowPtaOom(const char *context)
     }
 
     const int64_t current_count = ++g_pta_oom_call_count;
-    if (current_count > trigger_count && current_count < trigger_count + 2) {
-        auto retmsg = std::string("NPU out of memory. Injected PTA OOM after ") +
-            std::to_string(current_count) + " PTA operations";
-        if (context != nullptr && context[0] != '\0') {
-            retmsg += ", context is ";
-            retmsg += context;
-        }
-        retmsg += ". ";
-        retmsg += PTA_ERROR(ErrCode::MEMORY);
-        TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
+    if (current_count <= trigger_count || current_count >= trigger_count + 2) {
+        return;
+    }
+
+    if (g_pta_oom_injected.exchange(true)) {
+        return;
+    }
+
+    if (device < 0) {
+        NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
+    }
+
+    NPUCachingAllocator::markAllBlockUnsafe(device);
+    c10_npu::set_npu_data_unsafe_flag(true);
+
+    auto retmsg = std::string("NPU out of memory. Injected full-card PTA OOM on NPU ") +
+        std::to_string(device) +
+        ". All existing tensors on this device are marked unsafe. "
+        "Triggered after " + std::to_string(current_count) + " PTA operations";
+    if (context != nullptr && context[0] != '\0') {
+        retmsg += ", context is ";
+        retmsg += context;
     }
+    retmsg += ". Recover with: torch_npu.npu.restart_device(" +
+        std::to_string(device) + ", rebuild_all_resources=True). ";
+    retmsg += PTA_ERROR(ErrCode::MEMORY);
+    TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
+}
+
+void resetPtaOomInjectState()
+{
+    g_pta_oom_injected.store(false);
+    g_pta_oom_call_count.store(0);
 }
 
 } // namespace c10_npu
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index 066ddbe3bc..f459b58148 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -336,7 +336,9 @@ std::string handleSuspectRemoteError(int errorCode);
 
 bool isCannOOM(const std::string &errMsg);
 
-void maybeThrowPtaOom(const char *context = nullptr);
+void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
+
+void resetPtaOomInjectState();
 
 bool ShouldAppendDeviceErrorVerbose();
 
diff --git a/torch_npu/csrc/core/npu/NPUQueue.cpp b/torch_npu/csrc/core/npu/NPUQueue.cpp
index dd9b46a956..2143c45254 100644
--- a/torch_npu/csrc/core/npu/NPUQueue.cpp
+++ b/torch_npu/csrc/core/npu/NPUQueue.cpp
@@ -350,7 +350,7 @@ NPUStatus Repository::MakeSureQueueEmpty(bool check_error)
         }
     }
     if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
-        maybeThrowPtaOom("MakeSureQueueEmpty");
+        maybeThrowPtaOom("MakeSureQueueEmpty", device_idx);
     }
     logger->debug("MakeSureQueueEmpty: clearing successful, device = %d, write_idx = %u, read_idx = %u, status = %d",
         device_idx, write_idx.idx, read_idx.idx, GetStatus());
diff --git a/torch_npu/csrc/npu/Module.cpp b/torch_npu/csrc/npu/Module.cpp
index 812f1f0c43..9110cb50d7 100644
--- a/torch_npu/csrc/npu/Module.cpp
+++ b/torch_npu/csrc/npu/Module.cpp
@@ -864,6 +864,7 @@ PyObject* THNPModule_restart_device_wrap(PyObject* self, PyObject* arg)
     }
     setDefaultStreamsStatus(device, c10_npu::RepoStatus::INIT);
     c10_npu::NPUCachingAllocator::cleanEvent();
+    c10_npu::resetPtaOomInjectState();
     loggerRecovery->info("NPU restart device end, device is %d.", device);
 
     Py_RETURN_NONE;
diff --git a/torch_npu/npu/_recovery.py b/torch_npu/npu/_recovery.py
index 00d983d2cd..c3c039423d 100644
--- a/torch_npu/npu/_recovery.py
+++ b/torch_npu/npu/_recovery.py
@@ -62,6 +62,8 @@ def restart_device(device_id: int, rebuild_all_resources: int = False):
         _recovery_all_npu_stream(device_id)
     torch_npu._C._npu_restart_device(device_id)
     _except_handler.set_force_stop_exception(False)
+    if rebuild_all_resources:
+        set_npu_tensor_unsafe_check_flag(False)
     # pg recovery
     npu_device = torch.device('npu')
     for pg in _pg_map:

From 94d2eb96387ccb414c33a1591bffa2d87add87e3 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 06:09:58 +0000
Subject: [PATCH 04/11] Remove PTA OOM inject recovery hooks for MindIE fast
 recovery

Drop resetPtaOomInjectState, restart_device reset call, and unsafe-flag
clearing. Keep full-card OOM fault injection only; recovery is handled
externally by MindIE.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/csrc/core/npu/NPUException.cpp | 9 +--------
 torch_npu/csrc/core/npu/NPUException.h   | 2 --
 torch_npu/csrc/npu/Module.cpp            | 1 -
 torch_npu/npu/_recovery.py               | 2 --
 4 files changed, 1 insertion(+), 13 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 8bd3d78b3e..3d31f3e5da 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -484,16 +484,9 @@ void maybeThrowPtaOom(const char *context, int device)
         retmsg += ", context is ";
         retmsg += context;
     }
-    retmsg += ". Recover with: torch_npu.npu.restart_device(" +
-        std::to_string(device) + ", rebuild_all_resources=True). ";
+    retmsg += ". ";
     retmsg += PTA_ERROR(ErrCode::MEMORY);
     TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
 }
 
-void resetPtaOomInjectState()
-{
-    g_pta_oom_injected.store(false);
-    g_pta_oom_call_count.store(0);
-}
-
 } // namespace c10_npu
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index f459b58148..c44a26c793 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -338,8 +338,6 @@ bool isCannOOM(const std::string &errMsg);
 
 void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
 
-void resetPtaOomInjectState();
-
 bool ShouldAppendDeviceErrorVerbose();
 
 void clear_device_error_info();
diff --git a/torch_npu/csrc/npu/Module.cpp b/torch_npu/csrc/npu/Module.cpp
index 9110cb50d7..812f1f0c43 100644
--- a/torch_npu/csrc/npu/Module.cpp
+++ b/torch_npu/csrc/npu/Module.cpp
@@ -864,7 +864,6 @@ PyObject* THNPModule_restart_device_wrap(PyObject* self, PyObject* arg)
     }
     setDefaultStreamsStatus(device, c10_npu::RepoStatus::INIT);
     c10_npu::NPUCachingAllocator::cleanEvent();
-    c10_npu::resetPtaOomInjectState();
     loggerRecovery->info("NPU restart device end, device is %d.", device);
 
     Py_RETURN_NONE;
diff --git a/torch_npu/npu/_recovery.py b/torch_npu/npu/_recovery.py
index c3c039423d..00d983d2cd 100644
--- a/torch_npu/npu/_recovery.py
+++ b/torch_npu/npu/_recovery.py
@@ -62,8 +62,6 @@ def restart_device(device_id: int, rebuild_all_resources: int = False):
         _recovery_all_npu_stream(device_id)
     torch_npu._C._npu_restart_device(device_id)
     _except_handler.set_force_stop_exception(False)
-    if rebuild_all_resources:
-        set_npu_tensor_unsafe_check_flag(False)
     # pg recovery
     npu_device = torch.device('npu')
     for pg in _pg_map:

From 4ab4cee7b45503c1d3f66f3c777e5bdb31d5beba Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 07:18:01 +0000
Subject: [PATCH 05/11] Raise default PTA_OOM_TRIGGER_COUNT to 400000 to skip
 warmup

Avoid injecting full-card OOM during model warmup; aligns with prior
malloc counter tuning for large models like Qwen 235B.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/csrc/core/npu/NPUException.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 3d31f3e5da..31e095395b 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -441,7 +441,7 @@ bool isCannOOM(const std::string &errMsg)
 namespace {
 static std::atomic<int64_t> g_pta_oom_call_count{0};
 static std::atomic<bool> g_pta_oom_injected{false};
-static constexpr int64_t kDefaultPtaOomTriggerCount = 6000;
+static constexpr int64_t kDefaultPtaOomTriggerCount = 400000;
 
 int64_t getPtaOomTriggerCount()
 {

From f01ec55ba6bb84eafc25bf42ddd6d9606c05ac6b Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 08:54:44 +0000
Subject: [PATCH 06/11] Add PTA OOM counter on P static graph execution paths

Record progress in InnerRunOpApi, ExecFunc, ExecFuncOpApi, and
AclmdlRIExecuteAsync so prefill static-graph inference increments the
counter. Defer throw to malloc/MakeSureQueueEmpty via pending flag.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/csrc/core/npu/NPUException.cpp      | 51 ++++++++++++++-----
 torch_npu/csrc/core/npu/NPUException.h        |  2 +
 .../csrc/core/npu/interface/AclInterface.cpp  |  4 ++
 torch_npu/csrc/framework/OpParamMaker.cpp     | 12 +++++
 4 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 31e095395b..2425586ce0 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -441,6 +441,8 @@ bool isCannOOM(const std::string &errMsg)
 namespace {
 static std::atomic<int64_t> g_pta_oom_call_count{0};
 static std::atomic<bool> g_pta_oom_injected{false};
+static std::atomic<bool> g_pta_oom_pending{false};
+static std::atomic<int64_t> g_pta_oom_pending_count{0};
 static constexpr int64_t kDefaultPtaOomTriggerCount = 400000;
 
 int64_t getPtaOomTriggerCount()
@@ -451,20 +453,9 @@ int64_t getPtaOomTriggerCount()
     }();
     return trigger_count;
 }
-} // namespace
 
-void maybeThrowPtaOom(const char *context, int device)
+void throwFullCardPtaOom(const char *context, int device)
 {
-    const int64_t trigger_count = getPtaOomTriggerCount();
-    if (trigger_count <= 0) {
-        return;
-    }
-
-    const int64_t current_count = ++g_pta_oom_call_count;
-    if (current_count <= trigger_count || current_count >= trigger_count + 2) {
-        return;
-    }
-
     if (g_pta_oom_injected.exchange(true)) {
         return;
     }
@@ -476,6 +467,7 @@ void maybeThrowPtaOom(const char *context, int device)
     NPUCachingAllocator::markAllBlockUnsafe(device);
     c10_npu::set_npu_data_unsafe_flag(true);
 
+    const int64_t current_count = g_pta_oom_pending_count.load();
     auto retmsg = std::string("NPU out of memory. Injected full-card PTA OOM on NPU ") +
         std::to_string(device) +
         ". All existing tensors on this device are marked unsafe. "
@@ -488,5 +480,40 @@ void maybeThrowPtaOom(const char *context, int device)
     retmsg += PTA_ERROR(ErrCode::MEMORY);
     TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
 }
+} // namespace
+
+void recordPtaOomProgress()
+{
+    const int64_t trigger_count = getPtaOomTriggerCount();
+    if (trigger_count <= 0 || g_pta_oom_injected.load()) {
+        return;
+    }
+
+    const int64_t current_count = ++g_pta_oom_call_count;
+    if (current_count > trigger_count && current_count < trigger_count + 2) {
+        g_pta_oom_pending_count.store(current_count);
+        g_pta_oom_pending.store(true);
+    }
+}
+
+void maybeThrowPtaOom(const char *context, int device)
+{
+    const int64_t trigger_count = getPtaOomTriggerCount();
+    if (trigger_count <= 0) {
+        return;
+    }
+
+    if (g_pta_oom_pending.load() && !g_pta_oom_injected.load()) {
+        g_pta_oom_pending.store(false);
+        throwFullCardPtaOom(context, device);
+        return;
+    }
+
+    recordPtaOomProgress();
+    if (g_pta_oom_pending.load() && !g_pta_oom_injected.load()) {
+        g_pta_oom_pending.store(false);
+        throwFullCardPtaOom(context, device);
+    }
+}
 
 } // namespace c10_npu
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index c44a26c793..4774ad916d 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -336,6 +336,8 @@ std::string handleSuspectRemoteError(int errorCode);
 
 bool isCannOOM(const std::string &errMsg);
 
+void recordPtaOomProgress();
+
 void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
 
 bool ShouldAppendDeviceErrorVerbose();
diff --git a/torch_npu/csrc/core/npu/interface/AclInterface.cpp b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
index 1c9f1793ef..560c18e672 100644
--- a/torch_npu/csrc/core/npu/interface/AclInterface.cpp
+++ b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
@@ -8,6 +8,7 @@
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/NPUException.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/GetCANNInfo.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -1076,6 +1077,9 @@ aclError AclmdlRIDebugPrint(aclmdlRI modelRI)
 aclError AclmdlRIExecuteAsync(aclmdlRI modelRI, aclrtStream stream)
 {
     ACL_CALL_LOG("aclmdlRIExecuteAsync", "modelRI=" << modelRI << ", stream=" << stream);
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::recordPtaOomProgress();
+    }
     typedef aclError (*AclmdlRIExecuteAsync)(aclmdlRI, aclrtStream);
     static AclmdlRIExecuteAsync func = nullptr;
     if (func == nullptr) {
diff --git a/torch_npu/csrc/framework/OpParamMaker.cpp b/torch_npu/csrc/framework/OpParamMaker.cpp
index 6fd381f8e8..846bf9a1ac 100644
--- a/torch_npu/csrc/framework/OpParamMaker.cpp
+++ b/torch_npu/csrc/framework/OpParamMaker.cpp
@@ -3,6 +3,8 @@
 
 #include "torch_npu/csrc/core/npu/CachingHostAllocator.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
+#include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/NPUQueue.h"
 #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
 #include "torch_npu/csrc/distributed/HCCLUtils.hpp"
@@ -273,6 +275,10 @@ aclError OpCommandImpl::InnerRun(
 
 aclError OpCommandImpl::InnerRunOpApi(const string &op_name, PROC_FUNC func)
 {
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::recordPtaOomProgress();
+    }
+
     aclError ret;
     auto stream = c10_npu::getCurrentNPUStream();
     if (stream.getRepoStopFlag()) {
@@ -332,6 +338,9 @@ bool ContainsAny(const std::string& str, std::initializer_list<std::string> patt
 int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<ExecuteParas *>(in->paramVal);
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::recordPtaOomProgress();
+    }
     ASCEND_LOGD("Op %s Run.", cur_paras->opType);
     logger->debug("ExecFunc: Op %s Run.", cur_paras->opType);
     aclError ret;
@@ -426,6 +435,9 @@ int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<ExecuteParasOpApi *>(in->paramVal);
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::recordPtaOomProgress();
+    }
     ASCEND_LOGD("Op %s Run.", cur_paras->opType);
     logger->debug("ExecFuncOpApi: Op %s Run.", cur_paras->opType);
     aclError ret;

From d3cf4601c15bf4fcd0ea54e889fec2fd6aa1bfea Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 09:15:05 +0000
Subject: [PATCH 07/11] Count PTA OOM at stream/device sync for Qwen P GE
 prefill

TorchAir GE static graph bypasses OpCommand during prefill body; hook
AclrtSynchronizeStreamWithTimeout and npuSynchronizeDevice so each
prefill forward completion increments and can throw OOM on MindIE path.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/csrc/core/npu/NPUStream.cpp              | 4 ++++
 torch_npu/csrc/core/npu/interface/AclInterface.cpp | 9 +++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUStream.cpp b/torch_npu/csrc/core/npu/NPUStream.cpp
index ad554f5694..a29e590889 100644
--- a/torch_npu/csrc/core/npu/NPUStream.cpp
+++ b/torch_npu/csrc/core/npu/NPUStream.cpp
@@ -14,6 +14,7 @@
 #include "torch_npu/csrc/core/npu/NPUGuard.h"
 #include "torch_npu/csrc/core/npu/NPUQueue.h"
 #include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
 #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
@@ -559,6 +560,9 @@ bool npuSynchronizeDevice(bool check_error)
             ASCEND_LOGE("MakeSureQueueEmpty fail, ret: %s", ret.c_str());
         }
     }
+    if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
+        maybeThrowPtaOom("npuSynchronizeDevice");
+    }
     auto acl_ret = c10_npu::acl::AclrtSynchronizeDeviceWithTimeout();
     if (acl_ret != ACL_ERROR_NONE) {
         CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(acl_ret);
diff --git a/torch_npu/csrc/core/npu/interface/AclInterface.cpp b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
index 560c18e672..0d155f837e 100644
--- a/torch_npu/csrc/core/npu/interface/AclInterface.cpp
+++ b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
@@ -595,8 +595,9 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) {
     typedef aclError (*AclrtSynchronizeStreamWithTimeout)(aclrtStream, int32_t);
     static AclrtSynchronizeStreamWithTimeout func = (AclrtSynchronizeStreamWithTimeout)GET_FUNC(aclrtSynchronizeStreamWithTimeout);
     int32_t timeout = c10_npu::option::OptionsManager::GetACLExecTimeout();
+    aclError ret = ACL_ERROR_NONE;
     if (func != nullptr) {
-        return func(stream, timeout);
+        ret = func(stream, timeout);
     } else {
         TORCH_NPU_WARN_ONCE(func, "Failed to find function", "aclrtSynchronizeStreamWithTimeout");
         typedef aclError (*AclrtSynchronizeStream)(aclrtStream);
@@ -605,8 +606,12 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) {
             func_backup = (AclrtSynchronizeStream)GET_FUNC(aclrtSynchronizeStream);
         }
         TORCH_CHECK(func_backup, "Failed to find function", "aclrtSynchronizeStreamWithTimeout and aclrtSynchronizeStream", PROF_ERROR(ErrCode::NOT_FOUND));
-        return func_backup(stream);
+        ret = func_backup(stream);
     }
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::maybeThrowPtaOom("AclrtSynchronizeStreamWithTimeout");
+    }
+    return ret;
 }
 
 aclError AclrtDestroyStreamForce(aclrtStream stream) {

From dfdccaa5fed9cf4dcebe27387bb1971088de5515 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 09:56:20 +0000
Subject: [PATCH 08/11] Guarantee PTA OOM on P prefill via forward boundary
 hooks

Add PTA_OOM_INJECT=1 for immediate full-card OOM on first forward
boundary (stream/device/event sync, queue drain, memcpy).

Add PTA_OOM_FORWARD_SYNC_TRIGGER_COUNT (default 10) separate from the
400k op counter so Qwen GE prefill can trigger without per-op PTA paths.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/csrc/core/npu/NPUEvent.cpp          |  4 ++
 torch_npu/csrc/core/npu/NPUException.cpp      | 52 +++++++++++++++++++
 torch_npu/csrc/core/npu/NPUException.h        |  2 +
 torch_npu/csrc/core/npu/NPUStream.cpp         |  6 ++-
 torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp |  5 ++
 .../csrc/core/npu/interface/AclInterface.cpp  | 11 ++--
 torch_npu/csrc/framework/OpParamMaker.cpp     |  3 ++
 7 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUEvent.cpp b/torch_npu/csrc/core/npu/NPUEvent.cpp
index 3a20235889..d25d282243 100644
--- a/torch_npu/csrc/core/npu/NPUEvent.cpp
+++ b/torch_npu/csrc/core/npu/NPUEvent.cpp
@@ -2,6 +2,7 @@
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
 #include "torch_npu/csrc/core/npu/NPUGuard.h"
 #include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
 #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
 #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h"
@@ -202,6 +203,9 @@ void NPUEvent::synchronize() const
         }
         NPU_CHECK_ERROR(aclrtSynchronizeEvent(event_));
         ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event_);
+        if (currentStreamCaptureStatus() == CaptureStatus::None) {
+            maybeThrowPtaOomOnForwardBoundary("NPUEvent::synchronize");
+        }
 #ifndef BUILD_LIBTORCH
         const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
         if (C10_UNLIKELY(trigger)) {
diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 2425586ce0..655f0b715d 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -440,10 +440,12 @@ bool isCannOOM(const std::string &errMsg)
 
 namespace {
 static std::atomic<int64_t> g_pta_oom_call_count{0};
+static std::atomic<int64_t> g_pta_forward_boundary_count{0};
 static std::atomic<bool> g_pta_oom_injected{false};
 static std::atomic<bool> g_pta_oom_pending{false};
 static std::atomic<int64_t> g_pta_oom_pending_count{0};
 static constexpr int64_t kDefaultPtaOomTriggerCount = 400000;
+static constexpr int64_t kDefaultForwardSyncTriggerCount = 10;
 
 int64_t getPtaOomTriggerCount()
 {
@@ -454,6 +456,24 @@ int64_t getPtaOomTriggerCount()
     return trigger_count;
 }
 
+int64_t getForwardSyncTriggerCount()
+{
+    const static int64_t trigger_count = []() -> int64_t {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_FORWARD_SYNC_TRIGGER_COUNT");
+        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultForwardSyncTriggerCount;
+    }();
+    return trigger_count;
+}
+
+bool isPtaOomInjectEnabled()
+{
+    const static bool enabled = []() -> bool {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_INJECT");
+        return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0);
+    }();
+    return enabled;
+}
+
 void throwFullCardPtaOom(const char *context, int device)
 {
     if (g_pta_oom_injected.exchange(true)) {
@@ -516,4 +536,36 @@ void maybeThrowPtaOom(const char *context, int device)
     }
 }
 
+void maybeThrowPtaOomOnForwardBoundary(const char *context, int device)
+{
+    if (g_pta_oom_injected.load()) {
+        return;
+    }
+
+    if (isPtaOomInjectEnabled()) {
+        throwFullCardPtaOom(context, device);
+        return;
+    }
+
+    if (currentStreamCaptureStatus() != CaptureStatus::None) {
+        return;
+    }
+
+    const int64_t forward_trigger = getForwardSyncTriggerCount();
+    if (forward_trigger > 0) {
+        const int64_t forward_count = ++g_pta_forward_boundary_count;
+        if (forward_count > forward_trigger && forward_count < forward_trigger + 2) {
+            g_pta_oom_pending_count.store(forward_count);
+            throwFullCardPtaOom(context, device);
+            return;
+        }
+    }
+
+    recordPtaOomProgress();
+    if (g_pta_oom_pending.load() && !g_pta_oom_injected.load()) {
+        g_pta_oom_pending.store(false);
+        throwFullCardPtaOom(context, device);
+    }
+}
+
 } // namespace c10_npu
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index 4774ad916d..2ca4f4c9a8 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -340,6 +340,8 @@ void recordPtaOomProgress();
 
 void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
 
+void maybeThrowPtaOomOnForwardBoundary(const char *context = nullptr, int device = -1);
+
 bool ShouldAppendDeviceErrorVerbose();
 
 void clear_device_error_info();
diff --git a/torch_npu/csrc/core/npu/NPUStream.cpp b/torch_npu/csrc/core/npu/NPUStream.cpp
index a29e590889..29245390d0 100644
--- a/torch_npu/csrc/core/npu/NPUStream.cpp
+++ b/torch_npu/csrc/core/npu/NPUStream.cpp
@@ -509,6 +509,10 @@ NPUStatus emptyAllNPUStream(bool check_error)
         }
     }
 
+    if (check_error) {
+        maybeThrowPtaOomOnForwardBoundary("emptyAllNPUStream");
+    }
+
     return NPU_STATUS_SUCCESS;
 }
 
@@ -561,7 +565,7 @@ bool npuSynchronizeDevice(bool check_error)
         }
     }
     if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
-        maybeThrowPtaOom("npuSynchronizeDevice");
+        maybeThrowPtaOomOnForwardBoundary("npuSynchronizeDevice");
     }
     auto acl_ret = c10_npu::acl::AclrtSynchronizeDeviceWithTimeout();
     if (acl_ret != ACL_ERROR_NONE) {
diff --git a/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp b/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp
index 52411a045d..19b30c4527 100644
--- a/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp
+++ b/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp
@@ -12,6 +12,8 @@
 #include "torch_npu/csrc/core/NPUSerialization.h"
 #include "torch_npu/csrc/core/npu/NPUHooksInterface.h"
 #include "torch_npu/csrc/core/npu/NPUEventManager.h"
+#include "torch_npu/csrc/core/npu/NPUException.h"
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
 
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -221,6 +223,9 @@ void NPUGuardImpl::synchronizeEvent(void* event) const
 
     NPU_CHECK_ERROR_WITHOUT_UCE(aclrtSynchronizeEvent(npu_event));
     ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", npu_event);
+    if (currentStreamCaptureStatus() == CaptureStatus::None) {
+        maybeThrowPtaOomOnForwardBoundary("NPUGuardImpl::synchronizeEvent");
+    }
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
     if (C10_UNLIKELY(trigger)) {
diff --git a/torch_npu/csrc/core/npu/interface/AclInterface.cpp b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
index 0d155f837e..7950787d8e 100644
--- a/torch_npu/csrc/core/npu/interface/AclInterface.cpp
+++ b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
@@ -609,7 +609,7 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) {
         ret = func_backup(stream);
     }
     if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::maybeThrowPtaOom("AclrtSynchronizeStreamWithTimeout");
+        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclrtSynchronizeStreamWithTimeout");
     }
     return ret;
 }
@@ -998,8 +998,9 @@ aclError AclrtSynchronizeDeviceWithTimeout(void)
     typedef aclError (*AclrtSynchronizeDeviceWithTimeout)(int32_t);
     static AclrtSynchronizeDeviceWithTimeout func = (AclrtSynchronizeDeviceWithTimeout)GET_FUNC(aclrtSynchronizeDeviceWithTimeout);
     int32_t timeout = c10_npu::option::OptionsManager::GetACLDeviceSyncTimeout();
+    aclError ret = ACL_ERROR_NONE;
     if (func != nullptr) {
-        return func(timeout);
+        ret = func(timeout);
     } else {
         if (timeout > 0) {
             TORCH_NPU_WARN_ONCE("The ACL_DEVICE_SYNC_TIMEOUT does not take effect. If you want to enable this env, please upgrade CANN to the matching version.");
@@ -1010,8 +1011,12 @@ aclError AclrtSynchronizeDeviceWithTimeout(void)
             func_backup = (AclrtSynchronizeDevice)GET_FUNC(aclrtSynchronizeDevice);
         }
         TORCH_CHECK(func_backup, "Failed to find function ", "aclrtSynchronizeDeviceWithTimeout and aclrtSynchronizeDevice", PTA_ERROR(ErrCode::NOT_FOUND));
-        return func_backup();
+        ret = func_backup();
+    }
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclrtSynchronizeDeviceWithTimeout");
     }
+    return ret;
 }
 
 aclError AclrtEventGetTimestamp(aclrtEvent event, uint64_t *timestamp)
diff --git a/torch_npu/csrc/framework/OpParamMaker.cpp b/torch_npu/csrc/framework/OpParamMaker.cpp
index 846bf9a1ac..db108f99d1 100644
--- a/torch_npu/csrc/framework/OpParamMaker.cpp
+++ b/torch_npu/csrc/framework/OpParamMaker.cpp
@@ -480,6 +480,9 @@ int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream)
 int MemcopyAsyncFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<c10_npu::queue::CopyParas *>(in->paramVal);
+    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::maybeThrowPtaOomOnForwardBoundary("MemcopyAsyncFunc");
+    }
     logger->debug("MemcopyAsyncFunc Run.");
     aclError ret;
     bool flag;

From f0aba97c68b9d42c0c33265ff13e9c3b32747953 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 15:35:06 +0000
Subject: [PATCH 09/11] Fix P prefill OOM: forward boundary mode when main
 counter disabled

When PTA_OOM_TRIGGER_COUNT=0 (P instance config), route malloc/queue/op
hooks to the forward-boundary counter instead of no-op.

Add PTA_OOM_SKIP_WARMUP_COUNT, PTA_OOM_DEBUG, route stream sync through
hooked AclrtSynchronizeStreamWithTimeout, and hook D2H memcpy plus
allocator event sync as additional P prefill boundaries.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 .../csrc/aten/ops/op_api/CopyKernelOpApi.cpp  |  3 +-
 .../csrc/core/npu/NPUCachingAllocator.cpp     |  3 +
 torch_npu/csrc/core/npu/NPUException.cpp      | 85 ++++++++++++++-----
 torch_npu/csrc/core/npu/NPUException.h        |  2 +-
 torch_npu/csrc/core/npu/NPUFunctions.cpp      |  3 +-
 .../csrc/core/npu/interface/AclInterface.cpp  |  6 +-
 torch_npu/csrc/framework/OpParamMaker.cpp     |  6 +-
 7 files changed, 82 insertions(+), 26 deletions(-)

diff --git a/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp b/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp
index 560b3437b3..5d600efb44 100644
--- a/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp
+++ b/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp
@@ -23,6 +23,7 @@
 #include "torch_npu/csrc/core/npu/CachingHostAllocator.h"
 #include "torch_npu/csrc/aten/NPUOpApiNativeFunctions.h"
 #include "torch_npu/csrc/aten/NPUNativeFunctions.h"
+#include "torch_npu/csrc/core/npu/interface/AclInterface.h"
 #include "third_party/op-plugin/op_plugin/utils/op_api_common.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -49,7 +50,7 @@ void copy_between_host_and_device_opapi(at::Tensor& dst, const at::Tensor& src,
         void* currentPtr = torch_npu::utils::is_npu(dst) ? src.data_ptr() : dst.data_ptr();
         process_non_blocking_copy(storage, currentPtr, stream, kind);
     } else {
-        aclError error = aclrtSynchronizeStream(stream);
+        aclError error = c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream);
         auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch(
             std::make_pair(dst.storage().unsafeGetStorageImpl(), dst.storage_offset() * dst.itemsize()), nbytes,
             std::make_pair(src.storage().unsafeGetStorageImpl(), src.storage_offset() * src.itemsize()), nbytes, kind);
diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
index d5143d5135..a8265ab1b9 100644
--- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
+++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
@@ -2937,6 +2937,9 @@ class DeviceCachingAllocator {
                 } else {
                     TORCH_NPU_MEMORY_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event.get());
                 }
+                if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
+                    maybeThrowPtaOomOnForwardBoundary("NPUCachingAllocator::synchronize_and_free_events");
+                }
 #ifndef BUILD_LIBTORCH
                 const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
                 if (C10_UNLIKELY(trigger)) {
diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 655f0b715d..e5a2f2e9bf 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -447,6 +447,24 @@ static std::atomic<int64_t> g_pta_oom_pending_count{0};
 static constexpr int64_t kDefaultPtaOomTriggerCount = 400000;
 static constexpr int64_t kDefaultForwardSyncTriggerCount = 10;
 
+bool isPtaOomDebugEnabled()
+{
+    const static bool enabled = []() -> bool {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DEBUG");
+        return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0);
+    }();
+    return enabled;
+}
+
+int64_t getWarmupSkipCount()
+{
+    const static int64_t skip_count = []() -> int64_t {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_SKIP_WARMUP_COUNT");
+        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : 0;
+    }();
+    return skip_count;
+}
+
 int64_t getPtaOomTriggerCount()
 {
     const static int64_t trigger_count = []() -> int64_t {
@@ -474,6 +492,11 @@ bool isPtaOomInjectEnabled()
     return enabled;
 }
 
+bool isForwardBoundaryModeEnabled()
+{
+    return isPtaOomInjectEnabled() || getForwardSyncTriggerCount() > 0;
+}
+
 void throwFullCardPtaOom(const char *context, int device)
 {
     if (g_pta_oom_injected.exchange(true)) {
@@ -502,10 +525,17 @@ void throwFullCardPtaOom(const char *context, int device)
 }
 } // namespace
 
-void recordPtaOomProgress()
+void recordPtaOomProgress(const char *context)
 {
+    if (g_pta_oom_injected.load()) {
+        return;
+    }
+
     const int64_t trigger_count = getPtaOomTriggerCount();
-    if (trigger_count <= 0 || g_pta_oom_injected.load()) {
+    if (trigger_count <= 0) {
+        if (isForwardBoundaryModeEnabled()) {
+            maybeThrowPtaOomOnForwardBoundary(context);
+        }
         return;
     }
 
@@ -518,19 +548,26 @@ void recordPtaOomProgress()
 
 void maybeThrowPtaOom(const char *context, int device)
 {
+    if (g_pta_oom_injected.load()) {
+        return;
+    }
+
     const int64_t trigger_count = getPtaOomTriggerCount();
     if (trigger_count <= 0) {
+        if (isForwardBoundaryModeEnabled()) {
+            maybeThrowPtaOomOnForwardBoundary(context, device);
+        }
         return;
     }
 
-    if (g_pta_oom_pending.load() && !g_pta_oom_injected.load()) {
+    if (g_pta_oom_pending.load()) {
         g_pta_oom_pending.store(false);
         throwFullCardPtaOom(context, device);
         return;
     }
 
-    recordPtaOomProgress();
-    if (g_pta_oom_pending.load() && !g_pta_oom_injected.load()) {
+    recordPtaOomProgress(context);
+    if (g_pta_oom_pending.load()) {
         g_pta_oom_pending.store(false);
         throwFullCardPtaOom(context, device);
     }
@@ -542,28 +579,38 @@ void maybeThrowPtaOomOnForwardBoundary(const char *context, int device)
         return;
     }
 
-    if (isPtaOomInjectEnabled()) {
-        throwFullCardPtaOom(context, device);
+    if (currentStreamCaptureStatus() != CaptureStatus::None) {
         return;
     }
 
-    if (currentStreamCaptureStatus() != CaptureStatus::None) {
+    const int64_t forward_count = ++g_pta_forward_boundary_count;
+    const int64_t warmup_skip = getWarmupSkipCount();
+    if (forward_count <= warmup_skip) {
+        if (isPtaOomDebugEnabled() && (forward_count % 1000 == 0 || forward_count == warmup_skip)) {
+            ASCEND_LOGI("PTA OOM forward boundary warmup: count=%lld skip=%lld context=%s",
+                static_cast<long long>(forward_count),
+                static_cast<long long>(warmup_skip),
+                context != nullptr ? context : "");
+        }
         return;
     }
+    const int64_t effective_count = forward_count - warmup_skip;
 
-    const int64_t forward_trigger = getForwardSyncTriggerCount();
-    if (forward_trigger > 0) {
-        const int64_t forward_count = ++g_pta_forward_boundary_count;
-        if (forward_count > forward_trigger && forward_count < forward_trigger + 2) {
-            g_pta_oom_pending_count.store(forward_count);
-            throwFullCardPtaOom(context, device);
-            return;
-        }
+    if (isPtaOomDebugEnabled()) {
+        ASCEND_LOGI("PTA OOM forward boundary: effective=%lld context=%s",
+            static_cast<long long>(effective_count),
+            context != nullptr ? context : "");
     }
 
-    recordPtaOomProgress();
-    if (g_pta_oom_pending.load() && !g_pta_oom_injected.load()) {
-        g_pta_oom_pending.store(false);
+    if (isPtaOomInjectEnabled()) {
+        throwFullCardPtaOom(context, device);
+        return;
+    }
+
+    const int64_t forward_trigger = getForwardSyncTriggerCount();
+    if (forward_trigger > 0 &&
+        effective_count > forward_trigger && effective_count < forward_trigger + 2) {
+        g_pta_oom_pending_count.store(effective_count);
         throwFullCardPtaOom(context, device);
     }
 }
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index 2ca4f4c9a8..4bef90c88a 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -336,7 +336,7 @@ std::string handleSuspectRemoteError(int errorCode);
 
 bool isCannOOM(const std::string &errMsg);
 
-void recordPtaOomProgress();
+void recordPtaOomProgress(const char *context = "pta_op");
 
 void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
 
diff --git a/torch_npu/csrc/core/npu/NPUFunctions.cpp b/torch_npu/csrc/core/npu/NPUFunctions.cpp
index 02cbc299c0..ac38d5c9be 100644
--- a/torch_npu/csrc/core/npu/NPUFunctions.cpp
+++ b/torch_npu/csrc/core/npu/NPUFunctions.cpp
@@ -6,6 +6,7 @@
 #include "torch_npu/csrc/core/npu/NPUAffinityController.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 #include "torch_npu/csrc/core/npu/GetCANNInfo.h"
+#include "torch_npu/csrc/core/npu/interface/AclInterface.h"
 #include "third_party/acl/inc/acl/acl_rt.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -363,7 +364,7 @@ void stream_synchronize(aclrtStream stream)
         trigger->traceNpuStreamSynchronization(reinterpret_cast<uintptr_t>(stream));
     }
 #endif
-    NPU_CHECK_ERROR(aclrtSynchronizeStream(stream));
+    NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream));
 }
 
 aclError SetDeviceResLimit(int32_t device, int32_t type, uint32_t value)
diff --git a/torch_npu/csrc/core/npu/interface/AclInterface.cpp b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
index 7950787d8e..bed4aa4b82 100644
--- a/torch_npu/csrc/core/npu/interface/AclInterface.cpp
+++ b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
@@ -1088,7 +1088,7 @@ aclError AclmdlRIExecuteAsync(aclmdlRI modelRI, aclrtStream stream)
 {
     ACL_CALL_LOG("aclmdlRIExecuteAsync", "modelRI=" << modelRI << ", stream=" << stream);
     if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress();
+        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclmdlRIExecuteAsync");
     }
     typedef aclError (*AclmdlRIExecuteAsync)(aclmdlRI, aclrtStream);
     static AclmdlRIExecuteAsync func = nullptr;
@@ -1588,6 +1588,10 @@ aclError AclrtMemcpyAsyncWithCondition(void *dst, size_t destMax, const void *sr
         func = (AclrtMemcpyAsyncWithConditionFunc)GET_FUNC(aclrtMemcpyAsyncWithCondition);
     }
     TORCH_CHECK(func, "Failed to find function ", "aclrtMemcpyAsyncWithCondition", PROF_ERROR(ErrCode::NOT_FOUND));
+    if (kind == aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST &&
+        c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
+        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclrtMemcpyAsyncWithCondition");
+    }
     return func(dst, destMax, src, count, kind, stream);
 }
 
diff --git a/torch_npu/csrc/framework/OpParamMaker.cpp b/torch_npu/csrc/framework/OpParamMaker.cpp
index db108f99d1..998ddd1033 100644
--- a/torch_npu/csrc/framework/OpParamMaker.cpp
+++ b/torch_npu/csrc/framework/OpParamMaker.cpp
@@ -276,7 +276,7 @@ aclError OpCommandImpl::InnerRun(
 aclError OpCommandImpl::InnerRunOpApi(const string &op_name, PROC_FUNC func)
 {
     if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress();
+        c10_npu::recordPtaOomProgress("InnerRunOpApi");
     }
 
     aclError ret;
@@ -339,7 +339,7 @@ int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<ExecuteParas *>(in->paramVal);
     if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress();
+        c10_npu::recordPtaOomProgress("ExecFunc");
     }
     ASCEND_LOGD("Op %s Run.", cur_paras->opType);
     logger->debug("ExecFunc: Op %s Run.", cur_paras->opType);
@@ -436,7 +436,7 @@ int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<ExecuteParasOpApi *>(in->paramVal);
     if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress();
+        c10_npu::recordPtaOomProgress("ExecFuncOpApi");
     }
     ASCEND_LOGD("Op %s Run.", cur_paras->opType);
     logger->debug("ExecFuncOpApi: Op %s Run.", cur_paras->opType);

From 94abaac0e63c0361a4678c2453e28347d0d62a39 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 15:51:11 +0000
Subject: [PATCH 10/11] Add timer-based PTA OOM trigger after configurable
 delay

Support PTA_OOM_TRIGGER_AFTER_SECONDS, PTA_OOM_TRIGGER_AFTER_MINUTES,
and PTA_OOM_TIMER=1 (default 6 minutes). Timer starts at PTA load and
throws full-card OOM on the next hook after the deadline.

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 torch_npu/csrc/core/npu/NPUException.cpp | 115 ++++++++++++++++++++++-
 1 file changed, 112 insertions(+), 3 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index e5a2f2e9bf..6a6a713532 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -1,5 +1,7 @@
 #include <atomic>
+#include <chrono>
 #include <cstdlib>
+#include <thread>
 
 #include "torch_npu/csrc/core/npu/NPUException.h"
 #include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
@@ -444,8 +446,12 @@ static std::atomic<int64_t> g_pta_forward_boundary_count{0};
 static std::atomic<bool> g_pta_oom_injected{false};
 static std::atomic<bool> g_pta_oom_pending{false};
 static std::atomic<int64_t> g_pta_oom_pending_count{0};
+static std::atomic<bool> g_pta_oom_timer_started{false};
+static std::atomic<bool> g_pta_oom_timer_expired{false};
+static std::chrono::steady_clock::time_point g_pta_oom_start_time;
 static constexpr int64_t kDefaultPtaOomTriggerCount = 400000;
 static constexpr int64_t kDefaultForwardSyncTriggerCount = 10;
+static constexpr int64_t kDefaultTriggerAfterSeconds = 360;
 
 bool isPtaOomDebugEnabled()
 {
@@ -492,9 +498,101 @@ bool isPtaOomInjectEnabled()
     return enabled;
 }
 
+int64_t getTriggerAfterSeconds()
+{
+    const static int64_t trigger_seconds = []() -> int64_t {
+        char *seconds_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_AFTER_SECONDS");
+        if (seconds_val != nullptr) {
+            return strtol(seconds_val, nullptr, 10);
+        }
+        char *minutes_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_AFTER_MINUTES");
+        if (minutes_val != nullptr) {
+            return strtol(minutes_val, nullptr, 10) * 60;
+        }
+        char *timer_mode = c10_npu::option::get_and_log_env("PTA_OOM_TIMER");
+        if (timer_mode != nullptr && strtol(timer_mode, nullptr, 10) != 0) {
+            return kDefaultTriggerAfterSeconds;
+        }
+        return 0;
+    }();
+    return trigger_seconds;
+}
+
+bool isTimerTriggerEnabled()
+{
+    return getTriggerAfterSeconds() > 0;
+}
+
+void throwFullCardPtaOom(const char *context, int device);
+
+void ensurePtaOomTimerStarted()
+{
+    if (g_pta_oom_timer_started.load()) {
+        return;
+    }
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
+    if (trigger_seconds <= 0) {
+        return;
+    }
+    bool expected = false;
+    if (!g_pta_oom_timer_started.compare_exchange_strong(expected, true)) {
+        return;
+    }
+    g_pta_oom_start_time = std::chrono::steady_clock::now();
+    ASCEND_LOGI("PTA OOM timer started, will trigger after %lld seconds",
+        static_cast<long long>(trigger_seconds));
+    std::thread([trigger_seconds]() {
+        std::this_thread::sleep_for(std::chrono::seconds(trigger_seconds));
+        if (!g_pta_oom_injected.load()) {
+            g_pta_oom_timer_expired.store(true);
+            g_pta_oom_pending_count.store(trigger_seconds);
+            ASCEND_LOGI("PTA OOM timer expired after %lld seconds, pending throw on next hook",
+                static_cast<long long>(trigger_seconds));
+        }
+    }).detach();
+}
+
+int64_t getElapsedSecondsSincePtaOomStart()
+{
+    const auto now = std::chrono::steady_clock::now();
+    return std::chrono::duration_cast<std::chrono::seconds>(now - g_pta_oom_start_time).count();
+}
+
+void maybeThrowPtaOomByTimer(const char *context, int device = -1)
+{
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
+    if (trigger_seconds <= 0 || g_pta_oom_injected.load()) {
+        return;
+    }
+    ensurePtaOomTimerStarted();
+    const bool expired = g_pta_oom_timer_expired.load() ||
+        (g_pta_oom_timer_started.load() && getElapsedSecondsSincePtaOomStart() >= trigger_seconds);
+    if (!expired) {
+        if (isPtaOomDebugEnabled()) {
+            const int64_t elapsed = getElapsedSecondsSincePtaOomStart();
+            if (elapsed > 0 && elapsed % 60 == 0) {
+                ASCEND_LOGI("PTA OOM timer: elapsed=%lld/%lld seconds context=%s",
+                    static_cast<long long>(elapsed),
+                    static_cast<long long>(trigger_seconds),
+                    context != nullptr ? context : "");
+            }
+        }
+        return;
+    }
+    g_pta_oom_pending_count.store(trigger_seconds);
+    throwFullCardPtaOom(context, device);
+}
+
+const bool kPtaOomTimerArmed = []() {
+    if (getTriggerAfterSeconds() > 0) {
+        ensurePtaOomTimerStarted();
+    }
+    return true;
+}();
+
 bool isForwardBoundaryModeEnabled()
 {
-    return isPtaOomInjectEnabled() || getForwardSyncTriggerCount() > 0;
+    return isPtaOomInjectEnabled() || getForwardSyncTriggerCount() > 0 || isTimerTriggerEnabled();
 }
 
 void throwFullCardPtaOom(const char *context, int device)
@@ -511,10 +609,15 @@ void throwFullCardPtaOom(const char *context, int device)
     c10_npu::set_npu_data_unsafe_flag(true);
 
     const int64_t current_count = g_pta_oom_pending_count.load();
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
     auto retmsg = std::string("NPU out of memory. Injected full-card PTA OOM on NPU ") +
         std::to_string(device) +
-        ". All existing tensors on this device are marked unsafe. "
-        "Triggered after " + std::to_string(current_count) + " PTA operations";
+        ". All existing tensors on this device are marked unsafe. ";
+    if (trigger_seconds > 0 && (g_pta_oom_timer_expired.load() || current_count == trigger_seconds)) {
+        retmsg += "Triggered after " + std::to_string(trigger_seconds) + " seconds";
+    } else {
+        retmsg += "Triggered after " + std::to_string(current_count) + " PTA operations";
+    }
     if (context != nullptr && context[0] != '\0') {
         retmsg += ", context is ";
         retmsg += context;
@@ -531,6 +634,8 @@ void recordPtaOomProgress(const char *context)
         return;
     }
 
+    maybeThrowPtaOomByTimer(context);
+
     const int64_t trigger_count = getPtaOomTriggerCount();
     if (trigger_count <= 0) {
         if (isForwardBoundaryModeEnabled()) {
@@ -552,6 +657,8 @@ void maybeThrowPtaOom(const char *context, int device)
         return;
     }
 
+    maybeThrowPtaOomByTimer(context, device);
+
     const int64_t trigger_count = getPtaOomTriggerCount();
     if (trigger_count <= 0) {
         if (isForwardBoundaryModeEnabled()) {
@@ -579,6 +686,8 @@ void maybeThrowPtaOomOnForwardBoundary(const char *context, int device)
         return;
     }
 
+    maybeThrowPtaOomByTimer(context, device);
+
     if (currentStreamCaptureStatus() != CaptureStatus::None) {
         return;
     }

From d6b0ffa2d042e6a3b72e6890f4499257bf14b379 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 24 Jun 2026 16:53:15 +0000
Subject: [PATCH 11/11] Remove count-based PTA OOM injection, keep timer only

Delete PTA_OOM_TRIGGER_COUNT, forward sync counter, warmup skip, and
PTA_OOM_INJECT. All hooks now call maybeThrowPtaOom which only checks
the timer (PTA_OOM_TRIGGER_AFTER_SECONDS/MINUTES or PTA_OOM_TIMER=1).

Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
---
 .../csrc/core/npu/NPUCachingAllocator.cpp     |   8 +-
 torch_npu/csrc/core/npu/NPUEvent.cpp          |   4 +-
 torch_npu/csrc/core/npu/NPUException.cpp      | 190 ++----------------
 torch_npu/csrc/core/npu/NPUException.h        |   4 -
 torch_npu/csrc/core/npu/NPUQueue.cpp          |   2 +-
 torch_npu/csrc/core/npu/NPUStream.cpp         |   6 +-
 torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp |   4 +-
 .../csrc/core/npu/interface/AclInterface.cpp  |  17 +-
 torch_npu/csrc/framework/OpParamMaker.cpp     |  16 +-
 9 files changed, 35 insertions(+), 216 deletions(-)

diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
index a8265ab1b9..5abd2b9354 100644
--- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
+++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
@@ -1152,9 +1152,7 @@ class DeviceCachingAllocator {
 
     Block *malloc(int device, size_t orig_size, aclrtStream stream, uint8_t allocator_type = 0)
     {
-        if (currentStreamCaptureStatus() == CaptureStatus::None) {
-            maybeThrowPtaOom("malloc", device);
-        }
+        maybeThrowPtaOom("malloc", device);
 
         TORCH_NPU_MEMORY_LOGD("Allocating memory: size=%zu, device=%d", orig_size, device);
         // done outside the lock because we don't know what locks the recorder needs
@@ -2937,8 +2935,8 @@ class DeviceCachingAllocator {
                 } else {
                     TORCH_NPU_MEMORY_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event.get());
                 }
-                if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
-                    maybeThrowPtaOomOnForwardBoundary("NPUCachingAllocator::synchronize_and_free_events");
+                if (check_error) {
+                    maybeThrowPtaOom("NPUCachingAllocator::synchronize_and_free_events");
                 }
 #ifndef BUILD_LIBTORCH
                 const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
diff --git a/torch_npu/csrc/core/npu/NPUEvent.cpp b/torch_npu/csrc/core/npu/NPUEvent.cpp
index d25d282243..8082fece88 100644
--- a/torch_npu/csrc/core/npu/NPUEvent.cpp
+++ b/torch_npu/csrc/core/npu/NPUEvent.cpp
@@ -203,9 +203,7 @@ void NPUEvent::synchronize() const
         }
         NPU_CHECK_ERROR(aclrtSynchronizeEvent(event_));
         ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event_);
-        if (currentStreamCaptureStatus() == CaptureStatus::None) {
-            maybeThrowPtaOomOnForwardBoundary("NPUEvent::synchronize");
-        }
+        maybeThrowPtaOom("NPUEvent::synchronize");
 #ifndef BUILD_LIBTORCH
         const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
         if (C10_UNLIKELY(trigger)) {
diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp
index 6a6a713532..4a8f41a0ab 100644
--- a/torch_npu/csrc/core/npu/NPUException.cpp
+++ b/torch_npu/csrc/core/npu/NPUException.cpp
@@ -7,7 +7,6 @@
 #include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
 #include "torch_npu/csrc/core/npu/NPURecovery.h"
-#include "torch_npu/csrc/core/npu/NPUStream.h"
 #include "torch_npu/csrc/core/npu/NpuVariables.h"
 #include "torch_npu/csrc/core/npu/register/OptionsManager.h"
 
@@ -441,16 +440,10 @@ bool isCannOOM(const std::string &errMsg)
 }
 
 namespace {
-static std::atomic<int64_t> g_pta_oom_call_count{0};
-static std::atomic<int64_t> g_pta_forward_boundary_count{0};
 static std::atomic<bool> g_pta_oom_injected{false};
-static std::atomic<bool> g_pta_oom_pending{false};
-static std::atomic<int64_t> g_pta_oom_pending_count{0};
 static std::atomic<bool> g_pta_oom_timer_started{false};
 static std::atomic<bool> g_pta_oom_timer_expired{false};
 static std::chrono::steady_clock::time_point g_pta_oom_start_time;
-static constexpr int64_t kDefaultPtaOomTriggerCount = 400000;
-static constexpr int64_t kDefaultForwardSyncTriggerCount = 10;
 static constexpr int64_t kDefaultTriggerAfterSeconds = 360;
 
 bool isPtaOomDebugEnabled()
@@ -462,42 +455,6 @@ bool isPtaOomDebugEnabled()
     return enabled;
 }
 
-int64_t getWarmupSkipCount()
-{
-    const static int64_t skip_count = []() -> int64_t {
-        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_SKIP_WARMUP_COUNT");
-        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : 0;
-    }();
-    return skip_count;
-}
-
-int64_t getPtaOomTriggerCount()
-{
-    const static int64_t trigger_count = []() -> int64_t {
-        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_COUNT");
-        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultPtaOomTriggerCount;
-    }();
-    return trigger_count;
-}
-
-int64_t getForwardSyncTriggerCount()
-{
-    const static int64_t trigger_count = []() -> int64_t {
-        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_FORWARD_SYNC_TRIGGER_COUNT");
-        return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultForwardSyncTriggerCount;
-    }();
-    return trigger_count;
-}
-
-bool isPtaOomInjectEnabled()
-{
-    const static bool enabled = []() -> bool {
-        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_INJECT");
-        return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0);
-    }();
-    return enabled;
-}
-
 int64_t getTriggerAfterSeconds()
 {
     const static int64_t trigger_seconds = []() -> int64_t {
@@ -518,11 +475,6 @@ int64_t getTriggerAfterSeconds()
     return trigger_seconds;
 }
 
-bool isTimerTriggerEnabled()
-{
-    return getTriggerAfterSeconds() > 0;
-}
-
 void throwFullCardPtaOom(const char *context, int device);
 
 void ensurePtaOomTimerStarted()
@@ -545,7 +497,6 @@ void ensurePtaOomTimerStarted()
         std::this_thread::sleep_for(std::chrono::seconds(trigger_seconds));
         if (!g_pta_oom_injected.load()) {
             g_pta_oom_timer_expired.store(true);
-            g_pta_oom_pending_count.store(trigger_seconds);
             ASCEND_LOGI("PTA OOM timer expired after %lld seconds, pending throw on next hook",
                 static_cast<long long>(trigger_seconds));
         }
@@ -558,31 +509,6 @@ int64_t getElapsedSecondsSincePtaOomStart()
     return std::chrono::duration_cast<std::chrono::seconds>(now - g_pta_oom_start_time).count();
 }
 
-void maybeThrowPtaOomByTimer(const char *context, int device = -1)
-{
-    const int64_t trigger_seconds = getTriggerAfterSeconds();
-    if (trigger_seconds <= 0 || g_pta_oom_injected.load()) {
-        return;
-    }
-    ensurePtaOomTimerStarted();
-    const bool expired = g_pta_oom_timer_expired.load() ||
-        (g_pta_oom_timer_started.load() && getElapsedSecondsSincePtaOomStart() >= trigger_seconds);
-    if (!expired) {
-        if (isPtaOomDebugEnabled()) {
-            const int64_t elapsed = getElapsedSecondsSincePtaOomStart();
-            if (elapsed > 0 && elapsed % 60 == 0) {
-                ASCEND_LOGI("PTA OOM timer: elapsed=%lld/%lld seconds context=%s",
-                    static_cast<long long>(elapsed),
-                    static_cast<long long>(trigger_seconds),
-                    context != nullptr ? context : "");
-            }
-        }
-        return;
-    }
-    g_pta_oom_pending_count.store(trigger_seconds);
-    throwFullCardPtaOom(context, device);
-}
-
 const bool kPtaOomTimerArmed = []() {
     if (getTriggerAfterSeconds() > 0) {
         ensurePtaOomTimerStarted();
@@ -590,11 +516,6 @@ const bool kPtaOomTimerArmed = []() {
     return true;
 }();
 
-bool isForwardBoundaryModeEnabled()
-{
-    return isPtaOomInjectEnabled() || getForwardSyncTriggerCount() > 0 || isTimerTriggerEnabled();
-}
-
 void throwFullCardPtaOom(const char *context, int device)
 {
     if (g_pta_oom_injected.exchange(true)) {
@@ -608,16 +529,11 @@ void throwFullCardPtaOom(const char *context, int device)
     NPUCachingAllocator::markAllBlockUnsafe(device);
     c10_npu::set_npu_data_unsafe_flag(true);
 
-    const int64_t current_count = g_pta_oom_pending_count.load();
     const int64_t trigger_seconds = getTriggerAfterSeconds();
     auto retmsg = std::string("NPU out of memory. Injected full-card PTA OOM on NPU ") +
         std::to_string(device) +
-        ". All existing tensors on this device are marked unsafe. ";
-    if (trigger_seconds > 0 && (g_pta_oom_timer_expired.load() || current_count == trigger_seconds)) {
-        retmsg += "Triggered after " + std::to_string(trigger_seconds) + " seconds";
-    } else {
-        retmsg += "Triggered after " + std::to_string(current_count) + " PTA operations";
-    }
+        ". All existing tensors on this device are marked unsafe. "
+        "Triggered after " + std::to_string(trigger_seconds) + " seconds";
     if (context != nullptr && context[0] != '\0') {
         retmsg += ", context is ";
         retmsg += context;
@@ -628,100 +544,28 @@ void throwFullCardPtaOom(const char *context, int device)
 }
 } // namespace
 
-void recordPtaOomProgress(const char *context)
-{
-    if (g_pta_oom_injected.load()) {
-        return;
-    }
-
-    maybeThrowPtaOomByTimer(context);
-
-    const int64_t trigger_count = getPtaOomTriggerCount();
-    if (trigger_count <= 0) {
-        if (isForwardBoundaryModeEnabled()) {
-            maybeThrowPtaOomOnForwardBoundary(context);
-        }
-        return;
-    }
-
-    const int64_t current_count = ++g_pta_oom_call_count;
-    if (current_count > trigger_count && current_count < trigger_count + 2) {
-        g_pta_oom_pending_count.store(current_count);
-        g_pta_oom_pending.store(true);
-    }
-}
-
 void maybeThrowPtaOom(const char *context, int device)
 {
-    if (g_pta_oom_injected.load()) {
-        return;
-    }
-
-    maybeThrowPtaOomByTimer(context, device);
-
-    const int64_t trigger_count = getPtaOomTriggerCount();
-    if (trigger_count <= 0) {
-        if (isForwardBoundaryModeEnabled()) {
-            maybeThrowPtaOomOnForwardBoundary(context, device);
-        }
-        return;
-    }
-
-    if (g_pta_oom_pending.load()) {
-        g_pta_oom_pending.store(false);
-        throwFullCardPtaOom(context, device);
-        return;
-    }
-
-    recordPtaOomProgress(context);
-    if (g_pta_oom_pending.load()) {
-        g_pta_oom_pending.store(false);
-        throwFullCardPtaOom(context, device);
-    }
-}
-
-void maybeThrowPtaOomOnForwardBoundary(const char *context, int device)
-{
-    if (g_pta_oom_injected.load()) {
-        return;
-    }
-
-    maybeThrowPtaOomByTimer(context, device);
-
-    if (currentStreamCaptureStatus() != CaptureStatus::None) {
+    const int64_t trigger_seconds = getTriggerAfterSeconds();
+    if (trigger_seconds <= 0 || g_pta_oom_injected.load()) {
         return;
     }
-
-    const int64_t forward_count = ++g_pta_forward_boundary_count;
-    const int64_t warmup_skip = getWarmupSkipCount();
-    if (forward_count <= warmup_skip) {
-        if (isPtaOomDebugEnabled() && (forward_count % 1000 == 0 || forward_count == warmup_skip)) {
-            ASCEND_LOGI("PTA OOM forward boundary warmup: count=%lld skip=%lld context=%s",
-                static_cast<long long>(forward_count),
-                static_cast<long long>(warmup_skip),
-                context != nullptr ? context : "");
+    ensurePtaOomTimerStarted();
+    const bool expired = g_pta_oom_timer_expired.load() ||
+        (g_pta_oom_timer_started.load() && getElapsedSecondsSincePtaOomStart() >= trigger_seconds);
+    if (!expired) {
+        if (isPtaOomDebugEnabled()) {
+            const int64_t elapsed = getElapsedSecondsSincePtaOomStart();
+            if (elapsed > 0 && elapsed % 60 == 0) {
+                ASCEND_LOGI("PTA OOM timer: elapsed=%lld/%lld seconds context=%s",
+                    static_cast<long long>(elapsed),
+                    static_cast<long long>(trigger_seconds),
+                    context != nullptr ? context : "");
+            }
         }
         return;
     }
-    const int64_t effective_count = forward_count - warmup_skip;
-
-    if (isPtaOomDebugEnabled()) {
-        ASCEND_LOGI("PTA OOM forward boundary: effective=%lld context=%s",
-            static_cast<long long>(effective_count),
-            context != nullptr ? context : "");
-    }
-
-    if (isPtaOomInjectEnabled()) {
-        throwFullCardPtaOom(context, device);
-        return;
-    }
-
-    const int64_t forward_trigger = getForwardSyncTriggerCount();
-    if (forward_trigger > 0 &&
-        effective_count > forward_trigger && effective_count < forward_trigger + 2) {
-        g_pta_oom_pending_count.store(effective_count);
-        throwFullCardPtaOom(context, device);
-    }
+    throwFullCardPtaOom(context, device);
 }
 
 } // namespace c10_npu
diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h
index 4bef90c88a..c44a26c793 100644
--- a/torch_npu/csrc/core/npu/NPUException.h
+++ b/torch_npu/csrc/core/npu/NPUException.h
@@ -336,12 +336,8 @@ std::string handleSuspectRemoteError(int errorCode);
 
 bool isCannOOM(const std::string &errMsg);
 
-void recordPtaOomProgress(const char *context = "pta_op");
-
 void maybeThrowPtaOom(const char *context = nullptr, int device = -1);
 
-void maybeThrowPtaOomOnForwardBoundary(const char *context = nullptr, int device = -1);
-
 bool ShouldAppendDeviceErrorVerbose();
 
 void clear_device_error_info();
diff --git a/torch_npu/csrc/core/npu/NPUQueue.cpp b/torch_npu/csrc/core/npu/NPUQueue.cpp
index 2143c45254..e775386c25 100644
--- a/torch_npu/csrc/core/npu/NPUQueue.cpp
+++ b/torch_npu/csrc/core/npu/NPUQueue.cpp
@@ -349,7 +349,7 @@ NPUStatus Repository::MakeSureQueueEmpty(bool check_error)
             throw std::runtime_error(runtime_error);
         }
     }
-    if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
+    if (check_error) {
         maybeThrowPtaOom("MakeSureQueueEmpty", device_idx);
     }
     logger->debug("MakeSureQueueEmpty: clearing successful, device = %d, write_idx = %u, read_idx = %u, status = %d",
diff --git a/torch_npu/csrc/core/npu/NPUStream.cpp b/torch_npu/csrc/core/npu/NPUStream.cpp
index 29245390d0..1ff6f1bd5f 100644
--- a/torch_npu/csrc/core/npu/NPUStream.cpp
+++ b/torch_npu/csrc/core/npu/NPUStream.cpp
@@ -510,7 +510,7 @@ NPUStatus emptyAllNPUStream(bool check_error)
     }
 
     if (check_error) {
-        maybeThrowPtaOomOnForwardBoundary("emptyAllNPUStream");
+        maybeThrowPtaOom("emptyAllNPUStream");
     }
 
     return NPU_STATUS_SUCCESS;
@@ -564,8 +564,8 @@ bool npuSynchronizeDevice(bool check_error)
             ASCEND_LOGE("MakeSureQueueEmpty fail, ret: %s", ret.c_str());
         }
     }
-    if (check_error && currentStreamCaptureStatus() == CaptureStatus::None) {
-        maybeThrowPtaOomOnForwardBoundary("npuSynchronizeDevice");
+    if (check_error) {
+        maybeThrowPtaOom("npuSynchronizeDevice");
     }
     auto acl_ret = c10_npu::acl::AclrtSynchronizeDeviceWithTimeout();
     if (acl_ret != ACL_ERROR_NONE) {
diff --git a/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp b/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp
index 19b30c4527..77cad1980d 100644
--- a/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp
+++ b/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp
@@ -223,9 +223,7 @@ void NPUGuardImpl::synchronizeEvent(void* event) const
 
     NPU_CHECK_ERROR_WITHOUT_UCE(aclrtSynchronizeEvent(npu_event));
     ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", npu_event);
-    if (currentStreamCaptureStatus() == CaptureStatus::None) {
-        maybeThrowPtaOomOnForwardBoundary("NPUGuardImpl::synchronizeEvent");
-    }
+    maybeThrowPtaOom("NPUGuardImpl::synchronizeEvent");
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
     if (C10_UNLIKELY(trigger)) {
diff --git a/torch_npu/csrc/core/npu/interface/AclInterface.cpp b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
index bed4aa4b82..75c5c611d4 100644
--- a/torch_npu/csrc/core/npu/interface/AclInterface.cpp
+++ b/torch_npu/csrc/core/npu/interface/AclInterface.cpp
@@ -608,9 +608,7 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) {
         TORCH_CHECK(func_backup, "Failed to find function", "aclrtSynchronizeStreamWithTimeout and aclrtSynchronizeStream", PROF_ERROR(ErrCode::NOT_FOUND));
         ret = func_backup(stream);
     }
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclrtSynchronizeStreamWithTimeout");
-    }
+    c10_npu::maybeThrowPtaOom("AclrtSynchronizeStreamWithTimeout");
     return ret;
 }
 
@@ -1013,9 +1011,7 @@ aclError AclrtSynchronizeDeviceWithTimeout(void)
         TORCH_CHECK(func_backup, "Failed to find function ", "aclrtSynchronizeDeviceWithTimeout and aclrtSynchronizeDevice", PTA_ERROR(ErrCode::NOT_FOUND));
         ret = func_backup();
     }
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclrtSynchronizeDeviceWithTimeout");
-    }
+    c10_npu::maybeThrowPtaOom("AclrtSynchronizeDeviceWithTimeout");
     return ret;
 }
 
@@ -1087,9 +1083,7 @@ aclError AclmdlRIDebugPrint(aclmdlRI modelRI)
 aclError AclmdlRIExecuteAsync(aclmdlRI modelRI, aclrtStream stream)
 {
     ACL_CALL_LOG("aclmdlRIExecuteAsync", "modelRI=" << modelRI << ", stream=" << stream);
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclmdlRIExecuteAsync");
-    }
+    c10_npu::maybeThrowPtaOom("AclmdlRIExecuteAsync");
     typedef aclError (*AclmdlRIExecuteAsync)(aclmdlRI, aclrtStream);
     static AclmdlRIExecuteAsync func = nullptr;
     if (func == nullptr) {
@@ -1588,9 +1582,8 @@ aclError AclrtMemcpyAsyncWithCondition(void *dst, size_t destMax, const void *sr
         func = (AclrtMemcpyAsyncWithConditionFunc)GET_FUNC(aclrtMemcpyAsyncWithCondition);
     }
     TORCH_CHECK(func, "Failed to find function ", "aclrtMemcpyAsyncWithCondition", PROF_ERROR(ErrCode::NOT_FOUND));
-    if (kind == aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST &&
-        c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::maybeThrowPtaOomOnForwardBoundary("AclrtMemcpyAsyncWithCondition");
+    if (kind == aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST) {
+        c10_npu::maybeThrowPtaOom("AclrtMemcpyAsyncWithCondition");
     }
     return func(dst, destMax, src, count, kind, stream);
 }
diff --git a/torch_npu/csrc/framework/OpParamMaker.cpp b/torch_npu/csrc/framework/OpParamMaker.cpp
index 998ddd1033..963e7bcdb6 100644
--- a/torch_npu/csrc/framework/OpParamMaker.cpp
+++ b/torch_npu/csrc/framework/OpParamMaker.cpp
@@ -275,9 +275,7 @@ aclError OpCommandImpl::InnerRun(
 
 aclError OpCommandImpl::InnerRunOpApi(const string &op_name, PROC_FUNC func)
 {
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress("InnerRunOpApi");
-    }
+    c10_npu::maybeThrowPtaOom("InnerRunOpApi");
 
     aclError ret;
     auto stream = c10_npu::getCurrentNPUStream();
@@ -338,9 +336,7 @@ bool ContainsAny(const std::string& str, std::initializer_list<std::string> patt
 int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<ExecuteParas *>(in->paramVal);
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress("ExecFunc");
-    }
+    c10_npu::maybeThrowPtaOom("ExecFunc");
     ASCEND_LOGD("Op %s Run.", cur_paras->opType);
     logger->debug("ExecFunc: Op %s Run.", cur_paras->opType);
     aclError ret;
@@ -435,9 +431,7 @@ int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<ExecuteParasOpApi *>(in->paramVal);
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::recordPtaOomProgress("ExecFuncOpApi");
-    }
+    c10_npu::maybeThrowPtaOom("ExecFuncOpApi");
     ASCEND_LOGD("Op %s Run.", cur_paras->opType);
     logger->debug("ExecFuncOpApi: Op %s Run.", cur_paras->opType);
     aclError ret;
@@ -480,9 +474,7 @@ int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream)
 int MemcopyAsyncFunc(c10_npu::queue::QueueParas *in, aclrtStream stream)
 {
     auto cur_paras = static_cast<c10_npu::queue::CopyParas *>(in->paramVal);
-    if (c10_npu::currentStreamCaptureStatus() == c10_npu::CaptureStatus::None) {
-        c10_npu::maybeThrowPtaOomOnForwardBoundary("MemcopyAsyncFunc");
-    }
+    c10_npu::maybeThrowPtaOom("MemcopyAsyncFunc");
     logger->debug("MemcopyAsyncFunc Run.");
     aclError ret;
     bool flag;