Skip to content

[None][fix] Fix and enhance MemoryCounters Singleton with compile-time safety and bounds checking #8140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open

[None][fix] Fix and enhance MemoryCounters Singleton with compile-time safety and bounds checking #8140

Changes from 9 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
a0d0511
[FIX] fix bugs caused by None attention_bias during qwen3 model conve…
Jul 25, 2025
47c8528
Merge branch 'main' into main
Fan-Yunfan Jul 28, 2025
8c017fd
Merge branch 'NVIDIA:main' into main
Fan-Yunfan Aug 6, 2025
4bd57ff
Merge branch 'NVIDIA:main' into main
Fan-Yunfan Aug 21, 2025
a35c38f
Merge branch 'NVIDIA:main' into main
Fan-Yunfan Oct 3, 2025
bbb2198
[None][fix] Enhance memory counters with compile-time safety and boun…
Fan-Yunfan Oct 4, 2025
4511c59
[None][fix] Explicitly prohibit all copy and assignment operations fo…
Fan-Yunfan Oct 4, 2025
a676106
Add header file for std::numeric_limits
Fan-Yunfan Oct 31, 2025
7924c7f
Merge branch 'main' into fyf_enhance_memory_counters
karljang Nov 3, 2025
8c43e55
[None][fix] Fix compile-time check and remove unnecessary checks and …
Fan-Yunfan Nov 6, 2025
1e9b5e2
Merge branch 'main' into fyf_enhance_memory_counters
Fan-Yunfan Nov 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions cpp/include/tensorrt_llm/runtime/memoryCounters.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

#include <atomic>
#include <cstddef>
#include <limits>
#include <string>

namespace tensorrt_llm::runtime
Expand Down Expand Up @@ -84,9 +85,18 @@ class MemoryCounters
return mPinnedPoolDiff;
}

template <MemoryType T>
struct always_false : std::false_type
{
};

template <MemoryType T>
void allocate(SizeType32 size)
{
if (size > static_cast<SizeType32>(std::numeric_limits<DiffType>::max()))
{
TLLM_THROW("Memory size too large for diff type: %zu", size);
}
auto const sizeDiff = static_cast<DiffType>(size);
if constexpr (T == MemoryType::kGPU)
{
Expand Down Expand Up @@ -115,7 +125,7 @@ class MemoryCounters
}
else
{
TLLM_THROW("Unknown memory type: %s", MemoryTypeString<T>::value);
static_assert(always_false<T>::value, "Unknown memory type!");
}
}

Expand All @@ -124,6 +134,10 @@ class MemoryCounters
template <MemoryType T>
void deallocate(SizeType32 size)
{
if (size > static_cast<SizeType32>(std::numeric_limits<DiffType>::max()))
{
TLLM_THROW("Memory size too large for diff type: %zu", size);
}
auto const sizeDiff = -static_cast<DiffType>(size);
if constexpr (T == MemoryType::kGPU)
{
Expand Down Expand Up @@ -152,14 +166,19 @@ class MemoryCounters
}
else
{
TLLM_THROW("Unknown memory type: %s", MemoryTypeString<T>::value);
static_assert(always_false<T>::value, "Unknown memory type!");
}
}

void deallocate(MemoryType memoryType, SizeType32 size);

static MemoryCounters& getInstance();

MemoryCounters(MemoryCounters const&) = delete;
MemoryCounters& operator=(MemoryCounters const&) = delete;
MemoryCounters(MemoryCounters&&) = delete;
MemoryCounters& operator=(MemoryCounters&&) = delete;

static std::string bytesToString(SizeType32 bytes, int precision = 2);

static std::string bytesToString(DiffType bytes, int precision = 2);
Expand Down