[Feature]: Improve vLLM CUDA Memory Utilization and Estimation

[yewentao256]

🚀 The feature, motivation and pitch

Thanks @houseroad for the context!

It’s a general problem, 1) we do CUDA graph with memory allocation for KV cache, 2) we run the torch.compile. Since we may do memory profiling in torch.compile which may consume a lot of memory, either we may need to reserve less memory for KV cache (with some HBM wasted at the end), or we run into OOM.

To solve this problem, the idea is described in #19480 (comment), @yinghai , which looks resonable to me.

1. Reserve a big range of memory address but allocate a small amount of KV memory using subset of those memory address so that we can do the profile and cudagraph capturing.
2. Deallocate the KV memory but hold the KV memory address. Estimate how much KV we can build now that we have information of cudagraph memory usage
3. Allocate the KV memory and map to the same reserved memory address for KV.

The CUDA memory allocation, deallocation, pause and resume can be found in https://github.com/fzyzcjy/torch_memory_saver/blob/0fa8e036b9f5f21d63bb1a06f70102b11f0b6743/csrc/core.cpp

I think this is an important problem to solve since a lot of Cuda memory issues raised during initialization and graph capturing like #25290 or #25597.

@yinghai Would you like to work on this?

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[yewentao256]

vllm serve deepseek-ai/DeepSeek-R1 -dp 8 --enable-expert-parallel --port 9256 Now could directly OOM while startup in main with 8 * B200