Add counter-based PTA OOM injection for P/D inference

torch_npu/csrc/core/npu/NPUCachingAllocator.cpp

@@ -23,6 +23,7 @@
 #include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"
 #include "torch_npu/csrc/core/npu/NPURecovery.h"
 #include "torch_npu/csrc/core/npu/NPUGuard.h"
+#include "torch_npu/csrc/core/npu/PtaOomInjector.h"
 #include "NPUBlockHandle.h"
 #include "torch_npu/csrc/core/npu/NpuVariables.h"
 #include "torch_npu/csrc/core/npu/GetCANNInfo.h"
@@ -107,7 +108,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,10 +1152,7 @@ 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.");
-
+    {
         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
         // to have...
@@ -3516,6 +3513,7 @@ class NpuCachingAllocator : public NPUAllocator {
 
         int device = 0;
         NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
+        c10_npu::pta_oom::maybeThrowAllocOom(device, size);
         LazySetDevice(device);
         void *devPtr = nullptr;
         void (*deleteFunc)(void *) = &local_raw_delete;
@@ -3546,6 +3544,7 @@ class NpuCachingAllocator : public NPUAllocator {
         }
         int device = 0;
         NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
+        c10_npu::pta_oom::maybeThrowAllocOom(device, size);
         void *realPtr = nullptr;
         void (*deleteFunc)(void *) = &local_raw_delete;
 

torch_npu/csrc/core/npu/PtaOomInjector.cpp

@@ -0,0 +1,125 @@
+#include "torch_npu/csrc/core/npu/PtaOomInjector.h"
+
+#include <atomic>
+#include <cstdlib>
+#include <string>
+
+#include <c10/util/Exception.h>
+
+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
+#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
+
+namespace c10_npu {
+namespace pta_oom {
+namespace {
+
+constexpr int64_t kDefaultTriggerCount = 30004;
+
+std::atomic<int64_t> g_candidate_count{0};
+std::atomic<bool> g_oom_triggered{false};
+
+std::string formatAllocSize(uint64_t size)
+{
+    if (size <= 1024) {
+        return std::to_string(size) + " bytes";
+    }
+    if (size <= 1048576) {
+        return std::to_string(size / 1024.0) + " KiB";
+    }
+    if (size <= 1073741824ULL) {
+        return std::to_string(size / 1048576.0) + " MiB";
+    }
+    return std::to_string(size / 1073741824.0) + " GiB";
+}
+
+int64_t getTriggerCount()
+{
+    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) : kDefaultTriggerCount;
+    }();
+    return trigger_count;
+}
+
+int64_t getOpTriggerCount()
+{
+    const static int64_t trigger_count = []() -> int64_t {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_OP_TRIGGER_COUNT");
+        if (env_val != nullptr) {
+            return strtol(env_val, nullptr, 10);
+        }
+        return getTriggerCount();
+    }();
+    return trigger_count;
+}
+
+bool shouldIgnoreCapture()
+{
+    const static bool ignore_capture = []() -> bool {
+        char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DURING_CAPTURE");
+        return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0);
+    }();
+    return ignore_capture;
+}
+
+bool isCaptureActive()
+{
+    if (shouldIgnoreCapture()) {
+        return false;
+    }
+    return c10_npu::currentStreamCaptureStatus() != c10_npu::CaptureStatus::None;
+}
+
+bool shouldTrigger(int64_t trigger_count, int64_t current_count, const std::string &message)
+{
+    if (current_count > trigger_count && current_count < trigger_count + 2) {
+        g_oom_triggered.store(true);
+        TORCH_CHECK_WITH(OutOfMemoryError, false, message.c_str());
+        return true;
+    }
+    return false;
+}
+
+} // namespace
+
+void maybeThrowAllocOom(int device, size_t size)
+{
+    if (size == 0 || g_oom_triggered.load() || isCaptureActive()) {
+        return;
+    }
+
+    const int64_t trigger_count = getTriggerCount();
+    if (trigger_count <= 0) {
+        return;
+    }
+
+    const int64_t current_count = ++g_candidate_count;
+    shouldTrigger(
+        trigger_count,
+        current_count,
+        std::string("NPU out of memory. Injected PTA alloc OOM after ") +
+        std::to_string(current_count) + " PTA events. Tried to allocate " +
+        formatAllocSize(size) + " on NPU " + std::to_string(device) + ".");
+}
+
+void maybeThrowOpOom(const std::string &op_name)
+{
+    if (g_oom_triggered.load() || isCaptureActive()) {
+        return;
+    }
+
+    const int64_t trigger_count = getOpTriggerCount();
+    if (trigger_count <= 0) {
+        return;
+    }
+
+    const int64_t current_count = ++g_candidate_count;
+    shouldTrigger(
+        trigger_count,
+        current_count,
+        std::string("NPU out of memory. Injected PTA op OOM after ") +
+        std::to_string(current_count) + " PTA events, last op is " + op_name + ".");
+}
+
+} // namespace pta_oom
+} // namespace c10_npu

torch_npu/csrc/core/npu/PtaOomInjector.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <cstddef>
+#include <string>
+
+namespace c10_npu {
+namespace pta_oom {
+
+// Trigger injected PTA OOM after PTA_OOM_TRIGGER_COUNT tensor allocations.
+// Useful for prefill (P) and other allocation-heavy phases.
+void maybeThrowAllocOom(int device, size_t size);
+
+// Trigger injected PTA OOM after PTA_OOM_OP_TRIGGER_COUNT NPU op executions.
+// Useful for decode (D) and graph-replay inference where allocations are reused.
+void maybeThrowOpOom(const std::string &op_name);
+
+} // namespace pta_oom
+} // namespace c10_npu

torch_npu/csrc/framework/OpCommand.cpp

@@ -14,6 +14,7 @@
 #include "torch_npu/csrc/aten/CustomFunctions.h"
 #include "torch_npu/csrc/core/npu/NPUFunctions.h"
 #include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
+#include "torch_npu/csrc/core/npu/PtaOomInjector.h"
 #include "torch_npu/csrc/logging/LogContext.h"
 #ifndef BUILD_LIBTORCH
 #include "torch_npu/csrc/sanitizer/NPUTrace.h"
@@ -130,6 +131,8 @@ OpCommand& OpCommand::Output(at::Tensor &output, const string &descName,
 
 void OpCommand::Run()
 {
+    const string &op_name = aclCmd->GetName();
+    c10_npu::pta_oom::maybeThrowOpOom(op_name);
     // Check for npu graph
     if (aclCmd->CheckCustomHandlerNull()) {
         g_used_aclop = true;
@@ -142,7 +145,6 @@ void OpCommand::Run()
     }
 
     aclCmd->SetEnginePriority();
-    const string &op_name = aclCmd->GetName();
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
 #endif
@@ -184,6 +186,7 @@ void OpCommand::Run()
 
 void OpCommand::RunOpApi(const string &op_name, PROC_FUNC func, bool sync)
 {
+    c10_npu::pta_oom::maybeThrowOpOom(op_name);
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
 #endif
@@ -231,6 +234,7 @@ void OpCommand::RunOpApi(const string &op_name, PROC_FUNC func, bool sync)
 
 void OpCommand::RunOpApiV2(const string &op_name, const PROC_FUNC &func, bool sync)
 {
+    c10_npu::pta_oom::maybeThrowOpOom(op_name);
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
 #endif
@@ -284,6 +288,7 @@ void OpCommand::RunOpApiV2(const string &op_name, const PROC_FUNC &func, bool sy
 
 void OpCommand::RunOpApiV3(const string &op_name, const PROC_FUNC &func, bool sync, c10_npu::NPUStream *task_stream)
 {
+    c10_npu::pta_oom::maybeThrowOpOom(op_name);
 #ifndef BUILD_LIBTORCH
     const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
 #endif