chore: Unify Python NVTX call

[kaiyux]

• Move and unify NVTX implementation to tensorrt_llm/_utils.py
• Replace the original LLM API NVTX with nvtx_range_debug call

[Copilot]

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

[kaiyux]

/bot run

[tensorrt-cicd]

PR_Github #1775 [ run ] triggered by Bot

[tensorrt-cicd]

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

[Superjomn]

We mentioned having nvtx levels in an earlier discussion for tradeoff between details and performance with nsys, I noticed this MR introduces two levels, nvtx_range and nvtx_range_debug, are these two levels sufficient? Or will we introduce more or per-module switching in the future? @kaiyux @FrankD412

And I noticed this PR move nvtx_range to ._utils.py, will you unify the torch.cuda.nvtx_range usages in other places like pyexecutor? @kaiyux

[kaiyux]

@Superjomn Thanks for the comment

are these two levels sufficient? Or will we introduce more or per-module switching in the future?

I can't think of a scenario that we will need more levels - we already have layerwise nvtx markers as a feature.

And I noticed this PR move nvtx_range to ._utils.py, will you unify the torch.cuda.nvtx_range usages in other places like pyexecutor?

It should already be unified in this PR.

[Superjomn]

@Superjomn Thanks for the comment

are these two levels sufficient? Or will we introduce more or per-module switching in the future?

I can't think of a scenario that we will need more levels - we already have layerwise nvtx markers as a feature.

And I noticed this PR move nvtx_range to ._utils.py, will you unify the torch.cuda.nvtx_range usages in other places like pyexecutor?

It should already be unified in this PR.

Got it; that is clear. Thanks for your reply.

[kaiyux]

/bot run

[kaiyux]

/bot run

[tensorrt-cicd]

PR_Github #1851 [ run ] triggered by Bot

[FrankD412]

We mentioned having nvtx levels in an earlier discussion for tradeoff between details and performance with nsys, I noticed this MR introduces two levels, nvtx_range and nvtx_range_debug, are these two levels sufficient? Or will we introduce more or per-module switching in the future? @kaiyux @FrankD412

And I noticed this PR move nvtx_range to ._utils.py, will you unify the torch.cuda.nvtx_range usages in other places like pyexecutor? @kaiyux

I think two levels is fine -- when we profile we want as much information as possible so as long as it's documented how to do that I think we're happy on our end.

[tensorrt-cicd]

PR_Github #1851 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1370 completed with status: 'SUCCESS'

[Superjomn]

LGTM

[kaiyux]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #2347 [ reuse-pipeline ] triggered by Bot

[tensorrt-cicd]

PR_Github #2347 [ reuse-pipeline ] completed with state SUCCESS
Reusing PR_Github #1851 for commit 6379c3e