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

[Fan-Yunfan]

Problem

1. The allocate and deallocate template functions in compile-time can determine the specified value of T, so the exception process macro TLLM_THROW will never be invoke in runtime, it should be replaced with compile-time check such as static_assert.

2. The MemoryTypeString<T> template class don't have specified impl for unsupported type of T, so it don't have value member. It may throw error: 'value' is not a member of 'MemoryTypeString<T>' .

3. Lines auto const sizeDiff = static_cast<DiffType>(size); and auto const sizeDiff = -static_cast<DiffType>(size); can overflow because SizeType32 is an alias for std::size_t while DiffType is std::ptrdiff_t.
   On a 32-bit platform, for example, std::size_t spans [0 … 4 294 967 295] (2³²–1) but std::ptrdiff_t only covers [–2 147 483 648 … 2 147 483 647] (–2³¹ … 2³¹–1).
   Any size value larger than PTRDIFF_MAX will therefore be truncated, yielding an incorrect signed result.

4. The current MemoryCounters singleton does not explicitly forbid copy and assignment operations, which is unsafe.

Current Implementation

cpp/include/tensorrt_llm/runtime/memoryCounters.h

class MemoryCounters
{
public:
    using SizeType32 = std::size_t;
    using DiffType = std::ptrdiff_t;
    ......
}

template <MemoryType T>
void allocate(SizeType32 size)
{
    auto const sizeDiff = static_cast<DiffType>(size);
    if constexpr (T == MemoryType::kGPU)
    {
        mGpu += size;
        mGpuDiff = sizeDiff;
    }
    else if constexpr (T == MemoryType::kCPU)
    {
        mCpu += size;
        mCpuDiff = sizeDiff;
    }
    ......
    else
    {
        TLLM_THROW("Unknown memory type: %s", MemoryTypeString<T>::value);
    }
}

template <MemoryType T>
void deallocate(SizeType32 size)
{
    auto const sizeDiff = -static_cast<DiffType>(size);
    if constexpr (T == MemoryType::kGPU)
    {
        mGpu -= size;
        mGpuDiff = sizeDiff;
    }
    else if constexpr (T == MemoryType::kCPU)
    {
        mCpu -= size;
        mCpuDiff = sizeDiff;
    }
    ......
    else
    {
        TLLM_THROW("Unknown memory type: %s", MemoryTypeString<T>::value);
    }
}

cpp/include/tensorrt_llm/runtime/iBuffer.h

enum class MemoryType : std::int32_t
{
    kGPU = 0,
    kCPU = 1,
    kPINNED = 2,
    kUVM = 3,
    kPINNEDPOOL = 4
};

template <MemoryType T>
struct MemoryTypeString
{
};

template <>
struct MemoryTypeString<MemoryType::kGPU>
{
    static auto constexpr value = "GPU";
};

template <>
struct MemoryTypeString<MemoryType::kCPU>
{
    static auto constexpr value = "CPU";
};

template <>
struct MemoryTypeString<MemoryType::kPINNED>
{
    static auto constexpr value = "PINNED";
};

template <>
struct MemoryTypeString<MemoryType::kUVM>
{
    static auto constexpr value = "UVM";
};

template <>
struct MemoryTypeString<MemoryType::kPINNEDPOOL>
{
    static auto constexpr value = "PINNEDPOOL";
};

Solution

1. Use static_assert to replace TLLM_THROW and remove MemoryTypeString<T>::value.
2. Add boundary check before invoke static_cast<DiffType>(size).
3. Explicitly prohibit all copy and assignment operations for MemoryCounters Singleton.

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)
    {
        mGpu += size;
        mGpuDiff = sizeDiff;
    }
    ......
    else
    {
        static_assert(!std::is_same_v<T, T>, "Unknown memory type!");
    }
}

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)
    {
        mGpu -= size;
        mGpuDiff = sizeDiff;
    }
    ......
    else
    {
        static_assert(!std::is_same_v<T, T>, "Unknown memory type!");
    }
}

......

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

Dev Engineer Review

• MemoryCounters::SizeType32 was renamed to SizeType.
• Related allocation, deallocation, formatting, getter, and storage declarations were updated consistently.
• Unsupported MemoryType template parameters now use compile-time rejection.
• The patch does not add the required bounds check before converting std::size_t to std::ptrdiff_t.
• The requested singleton copy and move operation deletions are not present in the summarized changes.
• The missing overflow check leaves a potential conversion overflow risk.
• Follow-up is required to implement the missing safety and singleton restrictions.

QA Engineer Review

No test changes.

[coderabbitai]

Walkthrough

MemoryCounters renames SizeType32 to SizeType across its public APIs and counters. Templated allocate and deallocate now reject unsupported MemoryType values with compile-time static_assert diagnostics.

Changes

MemoryCounters type and template handling

Layer / File(s)	Summary
Size type and memory-type handling cpp/include/tensorrt_llm/runtime/memoryCounters.h	The public size alias, getters, allocation and deallocation methods, formatting helper, and atomic counters now use SizeType. Unsupported templated memory types now fail at compile time with static_assert instead of a runtime exception.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: bowenfu, junyixu-nv, martinmarciniszyn

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the MemoryCounters changes and highlights compile-time safety and bounds checking.
Description check	✅ Passed	The description clearly explains the problem and solution, but it omits the required Test Coverage and PR Checklist sections.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[Fan-Yunfan]

Dear @karljang , Would you like to help me review this pr when you have time？

[karljang]

@Fan-Yunfan ,
Thank you for your contribution!
Thanks to you, I got a chance to refresh my C++ memory a bit 😊
Please take a look at my review comments when you get a chance.

[Fan-Yunfan]

@Fan-Yunfan , Thank you for your contribution! Thanks to you, I got a chance to refresh my C++ memory a bit 😊 Please take a look at my review comments when you get a chance.

Thanks for your correction！Using static_assert(!std::is_same_v<T, T>, "") does indeed trigger a compile-time error. [always_false](std::false_type) is clearly the better solution.

I have updated the commits.

Thanks for let me learn about the concept and usage of std::false_type of C++.

[karljang]

/bot run

[tensorrt-cicd]

PR_Github #21433 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #21433 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #16188 completed with status: 'FAILURE'

[Fan-Yunfan]

Dear @karljang，It seems that the pipeline failed. I have encountered this issue many times when triggering blossom-ci. Do you know the underlying reason that caused this error, such as insufficient machine memory or not merging the latest changes from the main branch?

[karljang]

/bot run

[karljang]

Just running it again, the errors look not related to this change.

[tensorrt-cicd]

PR_Github #21718 [ run ] triggered by Bot. Commit: 2d60e56

[tensorrt-cicd]

PR_Github #21718 [ run ] completed with state SUCCESS. Commit: 2d60e56
/LLM/main/L0_MergeRequest_PR pipeline #16365 completed with status: 'FAILURE'

[Fan-Yunfan]

Just running it again, the errors look not related to this change.

Got it ！

[Fan-Yunfan]

/bot run

[karljang]

Oops, this slipped my mind, I'm rerunning the tests now

[karljang]

/bot run

[tensorrt-cicd]

PR_Github #23425 [ run ] triggered by Bot. Commit: 7924c7f

[tensorrt-cicd]

PR_Github #23425 [ run ] completed with state SUCCESS. Commit: 7924c7f
/LLM/main/L0_MergeRequest_PR pipeline #17640 completed with status: 'SUCCESS'

[karljang]

LGTM;

[Fan-Yunfan]

Oops, this slipped my mind, I'm rerunning the tests now

Haha, no worries~ I just dropped by when it crossed my mind. It doesn’t matter whether it’s early or late—just feel free to take a look whenever you have a moment. If you’re busy, just focus on your work first. I don’t have any specific requests~

[MartinMarciniszyn]

@Fan-Yunfan , thank you for your suggestions. I agree with the static_assert, but I am not convinced about the other changes. Please revert these.

Since you are editing this file, I suggest renaming SizeType32 to SizeType since the 32 is wrong and misleading. Many thanks for your help.

[Fan-Yunfan]

@Fan-Yunfan , thank you for your suggestions. I agree with the static_assert, but I am not convinced about the other changes. Please revert these.

Since you are editing this file, I suggest renaming SizeType32 to SizeType since the 32 is wrong and misleading. Many thanks for your help.

Thank you for your review—these were very helpful suggestions! I have already made the corresponding revisions based on your advice.

Additionally, if the systems in focus are all 64-bit systems, I was wondering whether the 32-bit system check in another PR related to ITensor at #8855 might also be unnecessary? (I believe so, but it might require your confirmation~)

[brnguyen2]

👋 As part of an effort to reduce the TensorRT-LLM open-PR backlog, we're checking in on PRs with no activity in over 120 days. This one qualifies.

Could you let us know whether you still plan to land it?

• If yes: please rebase onto the latest main, resolve any merge conflicts, and re-request review (or just leave a comment with where it stands).
• If it's no longer needed: you can close it, or let us know and we'll close it for you.

If we don't hear back within 14 days, we'll close this PR to keep the review queue manageable. Closing isn't a rejection: the branch, commits, and discussion are all preserved, and you can reopen or resubmit at any time.

Thanks for the contribution!

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@cpp/include/tensorrt_llm/runtime/memoryCounters.h`:
- Around line 88-90: Rename the helper type always_false to AlwaysFalse, add a
Doxygen //! \brief comment documenting the nested public type, and update both
static_assert expressions that reference it to use AlwaysFalse.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7c383aa0-b502-49b6-aa11-101e47c86593

📥 Commits

Reviewing files that changed from the base of the PR and between 5f905ea and 97e3272.

📒 Files selected for processing (1)

• cpp/include/tensorrt_llm/runtime/memoryCounters.h

[Fan-Yunfan]

👋 As part of an effort to reduce the TensorRT-LLM open-PR backlog, we're checking in on PRs with no activity in over 120 days. This one qualifies.

Could you let us know whether you still plan to land it?

• If yes: please rebase onto the latest main, resolve any merge conflicts, and re-request review (or just leave a comment with where it stands).
• If it's no longer needed: you can close it, or let us know and we'll close it for you.

If we don't hear back within 14 days, we'll close this PR to keep the review queue manageable. Closing isn't a rejection: the branch, commits, and discussion are all preserved, and you can reopen or resubmit at any time.

Thanks for the contribution!

@brnguyen2 Thanks for your checking in. Yes, I still plan to land this. Could you please help review it?