Skip to content

[TRTLLM-4406][feat] LLM sleep & wakeup Part 1: virtual device memory #5034

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

Jump to bottom
Merged
merged 16 commits into from
Aug 4, 2025
Merged

[TRTLLM-4406][feat] LLM sleep & wakeup Part 1: virtual device memory #5034

merged 16 commits into from
Aug 4, 2025

Conversation

tongyuantongyu
Copy link
Member

@tongyuantongyu tongyuantongyu commented Jun 9, 2025 •
edited by coderabbitai bot
Loading

Description

This PR implements virtual device memory that can be released and later reallocated without changing memory address. This paves the road to support LLM sleep and wakeup functionality.

Preliminary test shows we can reduce the GPU memory usage to less than 1GB during sleep.

Test Coverage

C++: cpp/tests/unit_tests/runtime/virtualMemoryTest.cpp

Python: tests/unittest/_torch/test_virtual_memory.py

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]

Launch build/test pipelines. All previously running jobs will be killed.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-[Post-Merge]-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

Summary by CodeRabbit

  • New Features

    • Introduced a CUDA virtual memory management system, enabling advanced allocation, release, and rematerialization of GPU memory.
    • Added Python and C++ APIs to control virtual memory allocation, including tag-based batch operations and restore modes (NONE, CPU, PINNED, MEMSET).
    • Provided PyTorch integration with context managers for scoped virtual memory allocation and utilities for releasing or materializing memory by tag.

  • Bug Fixes

    • Improved exception safety and error tracking for virtual memory operations.

  • Tests

    • Added comprehensive unit and integration tests for CUDA virtual memory management, covering allocation, release, restoration, and PyTorch workflows.

  • Chores

    • Updated build configurations to include new virtual memory sources and tests.

@tongyuantongyu tongyuantongyu requested a review from joyang-nv June 9, 2025 10:08
@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch from 74bbfe2 to 065f219 Compare June 9, 2025 10:14
@WilliamTambellini
Copy link
Contributor

Cool @tongyuantongyu
I ve looked at the change but did nt have time to dig the detail of UVMAllocator.
Does it use cudaMallocManaged() ?
eg.
https://developer.nvidia.com/blog/tag/unified-memory/
Best

@tongyuantongyu
Copy link
Member Author

I ve looked at the change but did nt have time to dig the detail of UVMAllocator.
Does it use cudaMallocManaged() ?

UVMAllocator uses cudaMallocManaged under the hood.

But this PR is not about UVMAllocator. It uses the low level CUDA VA API: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html.

hchings
hchings reviewed Jun 16, 2025
View reviewed changes
@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch 2 times, most recently from 3644b98 to 08e298b Compare June 30, 2025 07:37
@tongyuantongyu tongyuantongyu marked this pull request as ready for review June 30, 2025 08:56
@tongyuantongyu tongyuantongyu requested a review from a team as a code owner June 30, 2025 08:56
@tongyuantongyu
Copy link
Member Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10325 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10325 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #7637 completed with status: 'FAILURE'

@tongyuantongyu
Copy link
Member Author

/bot run

@tongyuantongyu
Copy link
Member Author

/bot kill

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10416 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10418 [ kill ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10416 [ run ] completed with state ABORTED

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10418 [ kill ] completed with state SUCCESS
Successfully killed previous jobs for commit 08e298b

@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch from cd982b5 to 9feef48 Compare July 1, 2025 07:14
@tongyuantongyu tongyuantongyu requested a review from a team as a code owner July 1, 2025 07:14
@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch from 9feef48 to 217b7c0 Compare July 1, 2025 07:14
@tongyuantongyu
Copy link
Member Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10455 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10455 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #7736 completed with status: 'FAILURE'

@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch 2 times, most recently from 379f7f4 to 6610455 Compare July 1, 2025 09:15
@tongyuantongyu
Copy link
Member Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10477 [ run ] triggered by Bot

@coderabbitai coderabbitai bot requested a review from pcastonguay July 31, 2025 04:00
@tongyuantongyu tongyuantongyu requested a review from ixlmar July 31, 2025 04:03
@tensorrt-cicd
Copy link
Collaborator

PR_Github #13603 [ run ] triggered by Bot

@tongyuantongyu
use shared_mutex
ee6676a 
Signed-off-by: Yuan Tong <[email protected]>
@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch from b738dc0 to ee6676a Compare July 31, 2025 05:42
@tongyuantongyu
Copy link
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13617 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13603 [ run ] completed with state ABORTED

@tongyuantongyu
Copy link
Member Author

/bot kill

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13629 [ kill ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13617 [ run ] completed with state ABORTED

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13629 [ kill ] completed with state SUCCESS
Successfully killed previous jobs for commit ee6676a

ixlmar
ixlmar approved these changes Jul 31, 2025
View reviewed changes
Copy link
Collaborator

@ixlmar ixlmar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Just one detail: PyTorch's use_mem_pool appears to be thread-local. If we used thread-local variables instead of globals for the allocator, our API would align better with theirs.

@ixlmar
Copy link
Collaborator

ixlmar commented Jul 31, 2025

LGTM

Just one detail: PyTorch's use_mem_pool appears to be thread-local. If we used thread-local variables instead of globals for the allocator, our API would align better with theirs.

But that might still cause issues with async, so I suggest we leave it as is and add contextvars once there is a need.

@tongyuantongyu
Copy link
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13658 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

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

@tongyuantongyu tongyuantongyu force-pushed the ytong/releasable_memory branch from 334daa5 to 9210aa1 Compare August 1, 2025 07:23
@tongyuantongyu
Copy link
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13784 [ run ] triggered by Bot

@tongyuantongyu tongyuantongyu changed the title feat: LLM sleep & wakeup Part 1: virtual device memory [TRTLLM-4406][feat] LLM sleep & wakeup Part 1: virtual device memory Aug 1, 2025
@tensorrt-cicd
Copy link
Collaborator

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

@tongyuantongyu
Copy link
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #13897 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

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

@tongyuantongyu tongyuantongyu merged commit a2f271c into NVIDIA:main Aug 4, 2025
4 of 6 checks passed
@tongyuantongyu tongyuantongyu deleted the ytong/releasable_memory branch August 4, 2025 05:52
lancelly pushed a commit to lancelly/TensorRT-LLM that referenced this pull request Aug 6, 2025
@tongyuantongyu @lancelly
[TRTLLM-4406][feat] LLM sleep & wakeup Part 1: virtual device memory (N…
32b1641 
…VIDIA#5034)

Signed-off-by: Yuan Tong <[email protected]>
Signed-off-by: Lanyu Liao <[email protected]>
jain-ria pushed a commit to jain-ria/TensorRT-LLM that referenced this pull request Aug 7, 2025
@tongyuantongyu
[TRTLLM-4406][feat] LLM sleep & wakeup Part 1: virtual device memory (N…
ae04001 
…VIDIA#5034)

Signed-off-by: Yuan Tong <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@hchings hchings hchings left review comments

@coderabbitai coderabbitai[bot] coderabbitai[bot] left review comments

@HuiGao-NV HuiGao-NV HuiGao-NV left review comments

@liji-nv liji-nv liji-nv approved these changes

@joyang-nv joyang-nv joyang-nv approved these changes

@lfr-0531 lfr-0531 Awaiting requested review from lfr-0531 lfr-0531 is a code owner automatically assigned from NVIDIA/trt-llm-torch-devs

@nv-yilinf nv-yilinf Awaiting requested review from nv-yilinf nv-yilinf is a code owner automatically assigned from NVIDIA/trt-llm-torch-devs

@DomBrown DomBrown Awaiting requested review from DomBrown

@nv-guomingz nv-guomingz Awaiting requested review from nv-guomingz

@chzblych chzblych Awaiting requested review from chzblych

@juney-nvidia juney-nvidia Awaiting requested review from juney-nvidia

@shaharmor98 shaharmor98 Awaiting requested review from shaharmor98

@pcastonguay pcastonguay Awaiting requested review from pcastonguay

+1 more reviewer

@ixlmar ixlmar ixlmar approved these changes

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
8 participants
@tongyuantongyu @WilliamTambellini @tensorrt-cicd @ixlmar @hchings @liji-nv @joyang-nv @HuiGao-NV