From 9d7049d229e7d04bf6ff3f94379ae062c24e1a02 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 06:49:15 +0000 Subject: [PATCH 1/4] Make PTA OOM injection controllable Co-authored-by: yjyang62 --- .../csrc/core/npu/NPUCachingAllocator.cpp | 85 +++++++++++++++++-- 1 file changed, 78 insertions(+), 7 deletions(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 2a4049a0dc..529b0b4194 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -1,5 +1,7 @@ #include +#include #include +#include #include #include #include @@ -23,6 +25,8 @@ #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/NPUGraphsUtils.h" +#include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "NPUBlockHandle.h" #include "torch_npu/csrc/core/npu/NpuVariables.h" #include "torch_npu/csrc/core/npu/GetCANNInfo.h" @@ -107,7 +111,10 @@ 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; +constexpr int64_t kDefaultPtaOomTriggerCount = 1000; +constexpr size_t kDefaultPtaOomMinAllocSize = kSmallSize; +static std::atomic g_pta_oom_candidate_count{0}; +static std::atomic g_pta_oom_triggered{false}; static char SHAREABLE_HANDLE_VERSION = 1; enum ShareableHandleType : char { SHAREABLE_NPU_MALLOC = 'c', @@ -116,6 +123,74 @@ enum ShareableHandleType : char { using StatTypes = std::array(StatType::NUM_TYPES)>; +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; +} + +size_t getPtaOomMinAllocSize() +{ + const static size_t min_alloc_size = []() -> size_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_MIN_ALLOC_BYTES"); + int64_t env_flag = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : + static_cast(kDefaultPtaOomMinAllocSize); + return env_flag > 0 ? static_cast(env_flag) : 0; + }(); + return min_alloc_size; +} + +int64_t getPtaOomTargetDevice() +{ + const static int64_t target_device = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DEVICE"); + return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : -1; + }(); + return target_device; +} + +bool isPtaOomEnabled() +{ + const static bool enabled = []() -> bool { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_ENABLE"); + return (env_val == nullptr) || (strtol(env_val, nullptr, 10) != 0); + }(); + return enabled; +} + +void maybeThrowPtaOom(int device, size_t size) +{ + if (!isPtaOomEnabled() || g_pta_oom_triggered.load()) { + return; + } + + const int64_t trigger_count = getPtaOomTriggerCount(); + if (trigger_count <= 0 || size == 0 || size < getPtaOomMinAllocSize()) { + return; + } + + const int64_t target_device = getPtaOomTargetDevice(); + if (target_device >= 0 && target_device != device) { + return; + } + + if (c10_npu::currentStreamCaptureStatus() != c10_npu::CaptureStatus::None) { + return; + } + + const int64_t current_count = ++g_pta_oom_candidate_count; + if (current_count > trigger_count && current_count < trigger_count + 2) { + g_pta_oom_triggered.store(true); + auto retmsg = std::string("NPU out of memory. Injected PTA OOM after ") + + std::to_string(current_count) + " eligible allocations. Tried to allocate " + + format_size(size) + " on NPU " + std::to_string(device) + "."; + TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); + } +} + void update_stat(Stat &stat, int64_t amount) { stat.current += amount; @@ -1152,12 +1227,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."); - if (g_malloc_call_count > 30004 && g_malloc_call_count < 30006) { - TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); - } + { 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... @@ -3518,6 +3588,7 @@ class NpuCachingAllocator : public NPUAllocator { int device = 0; NPU_CHECK_ERROR(c10_npu::GetDevice(&device)); + maybeThrowPtaOom(device, size); LazySetDevice(device); void *devPtr = nullptr; void (*deleteFunc)(void *) = &local_raw_delete; From aa70de7c3c2978ee84fbdd776a31b8697e51c0c7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 07:53:55 +0000 Subject: [PATCH 2/4] Narrow PTA OOM recovery trigger scope Co-authored-by: yjyang62 --- .../csrc/core/npu/NPUCachingAllocator.cpp | 42 ++++++++++++++++++- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 529b0b4194..55a0596f42 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -112,7 +112,7 @@ 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 constexpr int64_t kDefaultPtaOomTriggerCount = 1000; -constexpr size_t kDefaultPtaOomMinAllocSize = kSmallSize; +constexpr size_t kDefaultPtaOomMinAllocSize = 64 * 1024 * 1024; static std::atomic g_pta_oom_candidate_count{0}; static std::atomic g_pta_oom_triggered{false}; static char SHAREABLE_HANDLE_VERSION = 1; @@ -143,6 +143,16 @@ size_t getPtaOomMinAllocSize() return min_alloc_size; } +size_t getPtaOomMaxAllocSize() +{ + const static size_t max_alloc_size = []() -> size_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_MAX_ALLOC_BYTES"); + int64_t env_flag = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : -1; + return env_flag > 0 ? static_cast(env_flag) : 0; + }(); + return max_alloc_size; +} + int64_t getPtaOomTargetDevice() { const static int64_t target_device = []() -> int64_t { @@ -152,11 +162,29 @@ int64_t getPtaOomTargetDevice() return target_device; } +int64_t getPtaOomTargetRank() +{ + const static int64_t target_rank = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_RANK"); + if (env_val != nullptr) { + return strtol(env_val, nullptr, 10); + } + return std::getenv("RANK") != nullptr ? 0 : -1; + }(); + return target_rank; +} + +int64_t getCurrentRank() +{ + char *rank = std::getenv("RANK"); + return rank != nullptr ? strtol(rank, nullptr, 10) : -1; +} + bool isPtaOomEnabled() { const static bool enabled = []() -> bool { char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_ENABLE"); - return (env_val == nullptr) || (strtol(env_val, nullptr, 10) != 0); + return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0); }(); return enabled; } @@ -172,11 +200,21 @@ void maybeThrowPtaOom(int device, size_t size) return; } + const size_t max_alloc_size = getPtaOomMaxAllocSize(); + if (max_alloc_size != 0 && size > max_alloc_size) { + return; + } + const int64_t target_device = getPtaOomTargetDevice(); if (target_device >= 0 && target_device != device) { return; } + const int64_t target_rank = getPtaOomTargetRank(); + if (target_rank >= 0 && target_rank != getCurrentRank()) { + return; + } + if (c10_npu::currentStreamCaptureStatus() != c10_npu::CaptureStatus::None) { return; } From 3746efc76abd89655669ea4e93105504ae7ff519 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 07:55:54 +0000 Subject: [PATCH 3/4] Enable PTA OOM on all ranks by default Co-authored-by: yjyang62 --- torch_npu/csrc/core/npu/NPUCachingAllocator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 55a0596f42..f656b01519 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -169,7 +169,7 @@ int64_t getPtaOomTargetRank() if (env_val != nullptr) { return strtol(env_val, nullptr, 10); } - return std::getenv("RANK") != nullptr ? 0 : -1; + return -1; }(); return target_rank; } From bdf9ef4b755fa6e582597f07c38fbaac00f911c8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 09:02:57 +0000 Subject: [PATCH 4/4] Use normal tensor allocation for PTA OOM injection Co-authored-by: yjyang62 --- .../csrc/core/npu/NPUCachingAllocator.cpp | 81 +------------------ 1 file changed, 4 insertions(+), 77 deletions(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index f656b01519..fad36bae04 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -111,8 +111,7 @@ 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 -constexpr int64_t kDefaultPtaOomTriggerCount = 1000; -constexpr size_t kDefaultPtaOomMinAllocSize = 64 * 1024 * 1024; +constexpr int64_t kDefaultPtaOomTriggerCount = 30004; static std::atomic g_pta_oom_candidate_count{0}; static std::atomic g_pta_oom_triggered{false}; static char SHAREABLE_HANDLE_VERSION = 1; @@ -132,86 +131,14 @@ int64_t getPtaOomTriggerCount() return trigger_count; } -size_t getPtaOomMinAllocSize() -{ - const static size_t min_alloc_size = []() -> size_t { - char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_MIN_ALLOC_BYTES"); - int64_t env_flag = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : - static_cast(kDefaultPtaOomMinAllocSize); - return env_flag > 0 ? static_cast(env_flag) : 0; - }(); - return min_alloc_size; -} - -size_t getPtaOomMaxAllocSize() -{ - const static size_t max_alloc_size = []() -> size_t { - char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_MAX_ALLOC_BYTES"); - int64_t env_flag = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : -1; - return env_flag > 0 ? static_cast(env_flag) : 0; - }(); - return max_alloc_size; -} - -int64_t getPtaOomTargetDevice() -{ - const static int64_t target_device = []() -> int64_t { - char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DEVICE"); - return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : -1; - }(); - return target_device; -} - -int64_t getPtaOomTargetRank() -{ - const static int64_t target_rank = []() -> int64_t { - char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_RANK"); - if (env_val != nullptr) { - return strtol(env_val, nullptr, 10); - } - return -1; - }(); - return target_rank; -} - -int64_t getCurrentRank() -{ - char *rank = std::getenv("RANK"); - return rank != nullptr ? strtol(rank, nullptr, 10) : -1; -} - -bool isPtaOomEnabled() -{ - const static bool enabled = []() -> bool { - char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_ENABLE"); - return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0); - }(); - return enabled; -} - void maybeThrowPtaOom(int device, size_t size) { - if (!isPtaOomEnabled() || g_pta_oom_triggered.load()) { + if (g_pta_oom_triggered.load() || size == 0) { return; } const int64_t trigger_count = getPtaOomTriggerCount(); - if (trigger_count <= 0 || size == 0 || size < getPtaOomMinAllocSize()) { - return; - } - - const size_t max_alloc_size = getPtaOomMaxAllocSize(); - if (max_alloc_size != 0 && size > max_alloc_size) { - return; - } - - const int64_t target_device = getPtaOomTargetDevice(); - if (target_device >= 0 && target_device != device) { - return; - } - - const int64_t target_rank = getPtaOomTargetRank(); - if (target_rank >= 0 && target_rank != getCurrentRank()) { + if (trigger_count <= 0) { return; } @@ -223,7 +150,7 @@ void maybeThrowPtaOom(int device, size_t size) if (current_count > trigger_count && current_count < trigger_count + 2) { g_pta_oom_triggered.store(true); auto retmsg = std::string("NPU out of memory. Injected PTA OOM after ") + - std::to_string(current_count) + " eligible allocations. Tried to allocate " + + std::to_string(current_count) + " normal tensor allocations. Tried to allocate " + format_size(size) + " on NPU " + std::to_string(device) + "."; TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); }