From 105de9d514d4d2565ddf2c5841f8f29bcca876d5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 24 Jun 2026 12:34:57 +0000 Subject: [PATCH] Add counter-based PTA OOM injection for P/D inference Implement PtaOomInjector with a shared counter across tensor allocations and NPU op executions so OOM faults can be triggered during both prefill (P) and decode (D) phases, including graph-replay inference on large models like Qwen-235B. - PTA_OOM_TRIGGER_COUNT: allocation event threshold (default 30004) - PTA_OOM_OP_TRIGGER_COUNT: optional separate op threshold - PTA_OOM_DURING_CAPTURE=1: allow injection during graph capture - HCCL OOM (HCCL_OOM_TRIGGER_COUNT) remains for P2P communication faults Co-authored-by: yjyang62 --- .../csrc/core/npu/NPUCachingAllocator.cpp | 9 +- torch_npu/csrc/core/npu/PtaOomInjector.cpp | 125 ++++++++++++++++++ torch_npu/csrc/core/npu/PtaOomInjector.h | 18 +++ torch_npu/csrc/framework/OpCommand.cpp | 7 +- 4 files changed, 153 insertions(+), 6 deletions(-) create mode 100644 torch_npu/csrc/core/npu/PtaOomInjector.cpp create mode 100644 torch_npu/csrc/core/npu/PtaOomInjector.h diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 783d27969a..c71c752d48 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/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; diff --git a/torch_npu/csrc/core/npu/PtaOomInjector.cpp b/torch_npu/csrc/core/npu/PtaOomInjector.cpp new file mode 100644 index 0000000000..40ed2ced53 --- /dev/null +++ b/torch_npu/csrc/core/npu/PtaOomInjector.cpp @@ -0,0 +1,125 @@ +#include "torch_npu/csrc/core/npu/PtaOomInjector.h" + +#include +#include +#include + +#include + +#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 g_candidate_count{0}; +std::atomic 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 diff --git a/torch_npu/csrc/core/npu/PtaOomInjector.h b/torch_npu/csrc/core/npu/PtaOomInjector.h new file mode 100644 index 0000000000..7e3fca6600 --- /dev/null +++ b/torch_npu/csrc/core/npu/PtaOomInjector.h @@ -0,0 +1,18 @@ +#pragma once + +#include +#include + +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 diff --git a/torch_npu/csrc/framework/OpCommand.cpp b/torch_npu/csrc/framework/OpCommand.cpp index 21bea98be0..1ef67a2705 100644 --- a/torch_npu/csrc/framework/OpCommand.cpp +++ b/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