[V1] Reuse V0's memory_profiling util for gpu worker memory profiling

[yeqcharlotte]

Purpose

Update: Rebase on top of #18974. Clean up duplicated calls, improve code readability and logging.

Follow-up PRs to reduce OOMs:

• Account for CUDA graph memory
• torch memory snapshot capture for profiling
• max-model-len auto for long ctx model

---

old context before rebase:
V1's memory profiling incorrectly counts memory used by other processes as non_torch_allocations. This means, if you try to start 2 vllm servers -- one uses 70% hbm, the other uses 20%, it'll complain about not having enough memory.

This seems to bother quite a few users in the V0 deprecation RFC: #18571 (comment).

So reattempt #14419 by leveraging V0's memory profiling util to address that as @youkaichao suggested.

Test Plan

1. Unit test -- it's already covered by the following test which is broken on main

pytest -vv tests/entrypoints/llm/test_gpu_utilization.py

2. E2E test
   Starting 2 servers as below:

CUDA_VISIBLE_DEVICES=0 wp vllm serve "meta-llama/Llama-3.2-1B" \
    --gpu-memory-utilization 0.7

VLLM_LOGGING_LEVEL=DEBUG CUDA_VISIBLE_DEVICES=0 wp python3 benchmarks/benchmark_latency.py \
    --model "meta-llama/Llama-3.2-1B" \
    --tensor-parallel-size 1 \
    --trust-remote-code \
    --gpu-memory-utilization 0.2

Test Result

before

DEBUG 06-09 00:54:11 [gpu_worker.py:237] Initial free memory: 40.32 GiB, free memory: 37.74 GiB, total GPU memory: 139.72 GiB
DEBUG 06-09 00:54:11 [gpu_worker.py:241] Peak torch memory: 7.07 GiB, non-torch forward-pass memory: 0.17 GiB, available KVCache memory: 20.70 GiB

after

DEBUG 06-09 01:03:28 [gpu_worker.py:221] Initial free memory: 40.32 GiB, free memory: 37.74 GiB, requested GPU memory: 27.94 GiB
DEBUG 06-09 01:03:28 [gpu_worker.py:225] Memory profiling takes 6.13 seconds. Total non KV cache memory: 7.16GiB; torch peak memory increase: 4.69GiB; non-torch forward increase memory: 0.15GiB; weights memory: 2.32GiB.
INFO 06-09 01:03:28 [gpu_worker.py:226] Available KV cache memory: 20.78 GiB

---

old test before rebase:

1. Unit test

before

ValueError: No available memory for the cache blocks. Try increasing `gpu_memory_utilization` when initializing the engine.
================================================ short test summary info =================================================
FAILED tests/entrypoints/llm/test_gpu_utilization.py::test_gpu_memory_utilization - RuntimeError: Engine core initialization failed. See root cause above. Failed core proc(s): {}

after

tests/entrypoints/llm/test_gpu_utilization.py::test_gpu_memory_utilization PASSED                                  [100%]

2. E2E test

before

ERROR 06-03 03:14:47 [core.py:508]     check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
ERROR 06-03 03:14:47 [core.py:508]   File "/data/users/yeq/gitrepos/vllm/vllm/v1/core/kv_cache_utils.py", line 532, in check_enough_kv_cache_memory
ERROR 06-03 03:14:47 [core.py:508]     raise ValueError("No available memory for the cache blocks. "
ERROR 06-03 03:14:47 [core.py:508] ValueError: No available memory for the cache blocks. Try increasing `gpu_memory_utilization` when initializing the engine.

after

INFO 06-06 23:28:31 [gpu_worker.py:224] Memory profiling takes 6.20 seconds
INFO 06-06 23:28:31 [gpu_worker.py:224] the current vLLM instance can use total_gpu_memory (139.72GiB) x gpu_memory_utilization (0.20) = 27.94GiB
INFO 06-06 23:28:31 [gpu_worker.py:224] model weights take 2.32GiB; non_torch_memory takes 0.15GiB; PyTorch activation peak memory takes 4.69GiB; the rest of the memory reserved for KV Cache is 20.78GiB.
INFO 06-06 23:28:31 [kv_cache_utils.py:638] GPU KV cache size: 681,040 tokens
INFO 06-06 23:28:31 [kv_cache_utils.py:641] Maximum concurrency for 131,072 tokens per request: 5.20x
...
Avg latency: 0.2802899726356069 seconds

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Hello @yeqcharlotte, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello! Gemini here, providing a summary of this pull request. This PR addresses a critical bug in the V1 memory profiling logic. Previously, V1 incorrectly included memory used by other processes on the same GPU as part of its own non_torch_allocations. This led to scenarios where attempting to start multiple vLLM instances on a single GPU, even with sufficient total memory available and appropriate gpu_memory_utilization settings, would fail with 'No available memory for the cache blocks' errors. This PR resolves this by adopting the more accurate memory profiling utility from V0, which correctly measures only the memory footprint of the current vLLM instance, allowing multiple instances to coexist on the same card as intended.

Highlights

• Bugfix: Incorrect Memory Profiling: Fixes a bug in V1's memory profiling where memory used by other processes on the same GPU was incorrectly counted towards the current instance's memory usage.
• Enables Multiple Instances on One GPU: Resolves the issue preventing multiple vLLM servers or benchmark processes from starting on the same GPU with appropriate memory utilization settings.
• Adopts V0 Profiling Method: Leverages the memory profiling utility from V0 (memory_profiling context manager and MemorySnapshot) for more accurate measurement of the current instance's memory footprint.
• Improved Memory Logging: Adds detailed logging during initialization showing the breakdown of memory usage (total GPU, utilization, weights, non-torch, activation peak, KV cache).

Changelog

• vllm/v1/worker/gpu_worker.py
  • Imported MemorySnapshot and memory_profiling from vllm.utils (lines 25).
  • Initialized self.baseline_snapshot = MemorySnapshot() after clearing cache in init_device (line 134).
  • Wrapped the self.model_runner.profile_run() call within the memory_profiling context manager, passing the baseline snapshot and weights memory (lines 187-191).
  • Replaced the old logic for calculating non_torch_allocations and peak_memory with the results obtained from the memory_profiling context manager (result.non_kv_cache_memory) (lines 202-205).
  • Added detailed logging using logger.info to print the memory breakdown during initialization (lines 207-223).

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively addresses a significant bug in V1 memory profiling that prevented multiple vLLM servers from efficiently sharing a GPU. The approach of leveraging the V0 memory profiling utility through the memory_profiling context manager is a smart way to reuse existing, robust logic.

The changes are clear, and the detailed PR description, including the test plan and results, is very helpful for understanding the impact and verifying the fix. The new detailed log message for memory breakdown is also a great addition for observability.

Overall, the code quality is good, and the fix appears to be correct and well-implemented. I don't have any major concerns.

Merge Readiness

The pull request seems to be in good shape. The bug fix is well-targeted, and the solution appears robust. The provided unit and E2E tests confirm the fix. I believe this PR is ready for merging after standard CI checks pass. As always, ensure any other relevant reviewers have a chance to look before merging. I am not authorized to approve pull requests.

[houseroad]

Overall looks pretty neat approach.

We can try to add some unittest to ensure the memory consumed by other process is not incorrectly counted as non-torch memory in the current case.

[houseroad]

We can add some comment to explain the motivation of baseline_snapshot.

[houseroad]

Wondering if @youkaichao or @WoosukKwon would like to give it a pass?

[mergify]

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

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

[ywang96]

The approach makes sense to me and I left a comment. Can you please also update this PR with main? Thanks!

[ywang96]

I don't think we need to always show this message to the end user so it's better to have this as debug rather than info.

It's also better to break up this long msg for code readability. See

vllm/vllm/v1/worker/gpu_worker.py

Lines 236 to 245 in 8335667

	GiB = lambda b: b / GiB_bytes
	logger.debug(
	"Initial free memory: %.2f GiB, free memory: %.2f GiB, "
	"total GPU memory: %.2f GiB", GiB(self.init_gpu_memory),
	GiB(free_gpu_memory), GiB(total_gpu_memory))
	logger.debug(
	"Peak torch memory: %.2f GiB, non-torch forward-pass memory: "
	"%.2f GiB, available KVCache memory: %.2f GiB",
	GiB(peak_torch_memory), GiB(non_torch_alloc_bytes),
	GiB(available_kv_cache_memory))

[yeqcharlotte]

@ywang96 kept 1 info which i think it's useful for users

[yeqcharlotte]

Overall looks pretty neat approach.

We can try to add some unittest to ensure the memory consumed by other process is not incorrectly counted as non-torch memory in the current case.

existing memory_profiling util already has coverage on that:

vllm/tests/test_utils.py

Line 323 in 12e5829

def test_memory_profiling():

speaking of that. good time to check how many more V0 utils can be reused in V1.

[houseroad]

We can create a refactor list for v0 deprecation, like more concrete plan of which part of v0 should be reused

[ProExpertProg]

I like this, thanks for cleaning this up. Could you also add the logging changes from #17122?

[ProExpertProg]

Could this use the existing data from the profile_result?

[ProExpertProg]

I think we can improve this message. We should say something along the lines of "other process freed up memory during profiling"

[Datta0]

Should we make this info as well? It would be very useful.

[yeqcharlotte]

i guess this message could be confusing to generic end users. keeping it as debug for now for developers to turn on through VLLM_LOGGING_LEVEL=DEBUG .
it's always easier to flip it if too many issues complain about this part.

[ywang96]

Generally speaking we now try to reduce the amount of startup server logs as much as possible so that it's less confusing to the end users, and IMO it makes sense to keep this kind of information to DEBUG level.

[Datta0]

gc.collect after log doesn't feel right

[yeqcharlotte]

it's some proactive final cleanup after the profile runs. probably wouldn't matter too much just in case we left around some objects. gc perf overhead here should not matter too.

[yeqcharlotte]

@ProExpertProg @Datta0 @ywang96 @maxdebayser - lemme know if you folks have more feedback! thanks!