[Core] Run garbage collector after CUDA graph capture to fix throughput regression

[micah-wil]

Summary

We identified a regression in throughput caused by #21146 in some tests (most noticeably when using small input & output lengths). Specifically, on 8xMI300X:

VLLM_USE_V1=1 VLLM_V1_USE_PREFILL_DECODE_ATTENTION=1 python3 vllm/benchmarks/benchmark_throughput.py --model /models/Llama-3.1-70B-Instruct-FP8-KV -tp 8 --num-prompts 1024 --kv-cache-dtype "fp8" --input-len 128 --output-len 128 --dtype float16  --max-num-seqs 1024 --gpu-memory-utilization 0.9 --compilation-config '{"cudagraph_mode":"FULL"}'
...
Capturing CUDA graphs (mixed prefill-decode, FULL): 100%|██████████| 67/67 [00:06<00:00, 10.02it/s]
...

Result (avg over 5 runs): 15694.636 total tok/s

Whereas if we disable gc freeze during graph capture by setting VLLM_ENABLE_CUDAGRAPH_GC=1:

VLLM_ENABLE_CUDAGRAPH_GC=1 VLLM_USE_V1=1 VLLM_V1_USE_PREFILL_DECODE_ATTENTION=1 python3 vllm/benchmarks/benchmark_throughput.py --model /models/Llama-3.1-70B-Instruct-FP8-KV -tp 8 --num-prompts 1024 --kv-cache-dtype "fp8" --input-len 128 --output-len 128 --dtype float16  --max-num-seqs 1024 --gpu-memory-utilization 0.9 --compilation-config '{"cudagraph_mode":"FULL"}'
...
Capturing CUDA graphs (mixed prefill-decode, FULL): 100%|██████████| 67/67 [00:39<00:00,  1.68it/s]
...

Result (avg over 5 runs): 16416.11 total tok/s

The throughput is ~5% higher with VLLM_ENABLE_CUDAGRAPH_GC=1.
It seems that we can recover the throughput if we force a gc.collect() right after the CUDA graph capture is complete. With this PR:

VLLM_USE_V1=1 VLLM_V1_USE_PREFILL_DECODE_ATTENTION=1 python3 vllm/benchmarks/benchmark_throughput.py --model /models/Llama-3.1-70B-Instruct-FP8-KV -tp 8 --num-prompts 1024 --kv-cache-dtype "fp8" --input-len 128 --output-len 128 --dtype float16  --max-num-seqs 1024 --gpu-memory-utilization 0.9 --compilation-config '{"cudagraph_mode":"FULL"}'
...
Capturing CUDA graphs (mixed prefill-decode, FULL): 100%|██████████| 67/67 [00:06<00:00,  9.99it/s]
...

Result (avg over 5 runs): 16334.29283 total tok/s

That is, with this PR, it seems that we can recover the baseline throughput while getting the speedup from freezing gc during CUDA graph capture.

[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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request addresses a throughput regression by adding a gc.collect() call immediately after CUDA graph capture completes. The change is located within the freeze_gc context manager in GPUModelRunner. By explicitly triggering garbage collection after it has been re-enabled, this change reclaims memory that may have accumulated while the garbage collector was frozen during graph capture. The benchmarks provided in the pull request description demonstrate that this successfully recovers the lost throughput without reintroducing the graph capture slowdown. The implementation is correct and well-justified.

[robertgshaw2-redhat]

good find!

[njhill]

@micah-wil small question ... does it make a positive or negative difference if you move that gc.collect() right before the gc.unfreeze() instead of after it?

[micah-wil]

@micah-wil small question ... does it make a positive or negative difference if you move that gc.collect() right before the gc.unfreeze() instead of after it?

Good question. I ran 5 of the same test with this change, and the average tok/s was 15799.514. So, it is worse than calling gc.collect() after gc.unfreeze.

[njhill]

Thanks @micah-wil!