[Core] Support fully transparent sleep mode

[youkaichao]

there is a strong need to put vllm in sleep mode (offload weight, discard kv cache), from rlhf community. see #10714 and #11638

I have implemented the core functionality in https://github.com/vllm-project/vllm_allocator_adaptor , and add the integration in this PR.

Currently, we support:

• offload model weights
• discard kv cache

NOTE: we do nothing for cudagraph memory right now.

Because the underlying cumem API is very low-level, this PR is also compatible with cudagraph.

With this PR, when sleeping mode is enabled:

the current vLLM instance can use total_gpu_memory (79.22GiB) x gpu_memory_utilization (0.90) = 71.29GiB
model weights take 2.32GiB; non_torch_memory takes 0.67GiB; PyTorch activation peak memory takes 1.20GiB; the rest of the memory reserved for KV Cache is 67.11GiB.

Free memory before sleep: 8.34 GiB
Free memory after sleep: 78.29 GiB
(70GiB memory released, which is the sum of model weights and KV Cache size)

Why free memory after sleep is not 80 GiB? Because cudagraph memory pool is not released, and there are some other cost such as cuda context.

TODO:

• move https://github.com/vllm-project/vllm_allocator_adaptor into vllm repo (need someone familiar with CMake. Note that I don't use any pytorch ops in that file, so binding should be easy)
• expose the interface in API server, similar to profiling endpoints (can be a follow-up PR, and we need to discuss how to expose these endpoints. It should live under some dev-mode, with explicit opt-in)
• drop existing block manager's content when prefix-caching is enabled (after a second thought, I think this is out of the scope of the current PR. but we do need a way to expose this API interface)
• add the feature in V1
• check if any features are incompatible with this sleeping mode. I guess it should just work for cuda platforms, but need to check. If that is true, we can enable it by default. (NOTE: I find some bugs about PyTorch's pluggable allocator, see empty_cache does not work for CUDAPluggableAllocator + MemPool pytorch/pytorch#145168 . Therefore we should not enable it by default.)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

thanks to @cennn who helped me a lot along the way.

[youkaichao]

this pr benefits a lot from pytorch/pytorch#131152 and pytorch/pytorch#124807

[youkaichao]

TODO: in distributed inference, there's also NCCL memory to be considered. need to check how much memory that takes, and if we need to do anything to release that part (it might be quite difficult, as NCCL is quite black box)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @youkaichao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[comaniac]

Did a quick pass and overall LGTM. A high level question: would there be a correctness issue with prefix caching?

If prefix caching is enabled, and we sleep with level 2 (discard the weights), there will be correctness issue. That's why I'm asking for help to reset prefix caching state.

Make sense. Will provide the reset capability to prefix caching.

btw meanwhile, please raise an error if sleep mode level 2 is used and prefix caching is enabled in this PR, so that it can be unblocked.

[youkaichao]

btw meanwhile, please raise an error if sleep mode level 2 is used and prefix caching is enabled in this PR, so that it can be unblocked.

@comaniac fixed in b371bf3

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @youkaichao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[comaniac]

Otherwise LGTM. Approve to unblock first but other's comments are welcome.

[WoosukKwon]

LGTM when I skimmed through the PR. Since this PR is quite isolated, I don't have a concern in merging it.

Will take a closer look a few days later once I have more bandwidth.

[youkaichao]

Otherwise LGTM. Approve to unblock first but other's comments are welcome.

@comaniac thanks for the very detailed comments! Looking forward to #12284 being merged after this PR.