[Aravind0403]

[Bugfix][Memory] Accurately profile CUDA graph memory for quantized/FP8 KV caches (#60350)

When using quantized KV caches (such as fp8, int8), the block table and dequantization scale allocations per CUDA graph allocate additional memory during capture. Previously, _extrapolate_full_graph_memory used a 1 MiB per-graph floor, which under-budgeted full graph capture overhead by ~0.8 GiB across ~30 graphs and triggered out-of-memory errors during engine startup capture on 24GB GPUs.

This commit:

1. Introduces _MIN_QUANTIZED_PER_GRAPH_BYTES = 16 << 20 (16 MiB) floor when is_quantized_kv_cache(kv_cache_dtype) is true.
2. Propagates kv_cache_dtype from profile_cudagraph_memory to _extrapolate_full_graph_memory.
3. Adds regression unit tests covering quantized vs non-quantized extrapolation in test_gpu_model_runner_v2_cudagraph_profiling.py.

Fixes #60350.

Overview

Accurately budgets CUDA graph memory for quantized/FP8 KV caches during memory profiling, preventing startup OOM errors during capture_model(). Fixes #60350.

Claims

• Prevents Startup OOM: Correctly reserves ~0.5–0.8 GiB of CUDA graph memory headroom for FP8/INT8 KV caches across multi-size graph capture on 24 GiB GPUs (RTX 3090/4090).
• Zero Overhead on Non-Quantized Paths: Preserves the existing 1 MiB floor for FP16/BF16/unquantized models without reducing available KV cache block counts.
• Unit Test Coverage: Adds regression tests for _extrapolate_full_graph_memory across both standard and quantized (fp8, fp8_e4m3) cache dtypes.

Validation

• Regression Unit Tests:

  pytest tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py
  

• Verified that:

_extrapolate_full_graph_memory enforces the 16 MiB floor per graph when kv_cache_dtype="fp8" or "fp8_e4m3".
Measured sample deltas that exceed 16 MiB are preserved.
Default non-quantized paths continue using the 1 MiB floor.

Details

Root Cause

During worker initialization, profile_cudagraph_memory() samples the memory of the largest FULL CUDA graphs into a throwaway pool and extrapolates the total capture memory with _extrapolate_full_graph_memory().

Previously, _extrapolate_full_graph_memory floored per-graph overhead at _MIN_PER_GRAPH_BYTES = 1 << 20 (1 MiB). While sufficient for 16-bit unquantized KV caches, quantized caches (fp8, int8) hold smaller element sizes, creating
2×2× more total blocks. During graph capture, allocating auxiliary scale factors and binding larger block tables consumes ~16–30 MiB per captured graph. For a typical batch schedule with ~30 full graphs, the 1 MiB floor underestimated graph capture memory by ~0.8 GiB, leaving insufficient headroom and triggering torch.OutOfMemoryError during the real capture_model() phase.

Design Decisions

_MIN_QUANTIZED_PER_GRAPH_BYTES = 16 << 20 (16 MiB): Applied when is_quantized_kv_cache(kv_cache_dtype) evaluates to true.
Plumbing: profile_cudagraph_memory extracts kv_cache_dtype from runner.cache_config.cache_dtype and passes it to _extrapolate_full_graph_memory.

Pull Request Checklist

I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

AI assistance was used during the creation of this PR.

Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[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.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

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

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[lucamotz]

I’m not sure this is the right approach to fixing the underlying issue. If I understand correctly, increasing the minimum estimated memory per CUDA graph to 16 MiB would permanently reduce the available KV cache allocation for quantized KV cache configurations, even if the actual CUDA graphs don’t require that much additional memory.

My main concern is that we’re introducing a fairly arbitrary buffer here rather than fixing the underlying memory profiling. This effectively trades KV cache capacity for additional safety margins, which I don’t think is a great solution, especially since the actual overhead likely depends on the configuration.

I’d rather understand why we’re underestimating CUDA graph memory in the first place and fix the profiling to account for the actual allocations. Otherwise we’re potentially wasting a significant amount of GPU memory without knowing whether the 16 MiB is even sufficient or necessary.

Is there a reason we can’t address this in the profiling itself?

[Aravind0403]

Thanks @lucamotz, fair critique. I measured on an RTX 4090 (24 GB, CUDA 13.0, PyTorch 2.13) with Qwen/Qwen2.5-1.5B-Instruct (35 full graphs) to look at the numbers directly:

• Negative sample deltas: The sampled graph deltas came back negative (-22 MiB on FP8, -24 MiB on BF16), so the current estimate becomes first_sample + (total_graphs - 1) * floor. With the 1 MiB floor that evaluates to ~12 MiB for FP8 and ~10 MiB for BF16.
• Actual capture memory (measured):
  • FP8 KV cache: Real capture consumed 0.51 GiB (~510 MiB deficit against the 12 MiB estimate).
  • BF16 KV cache: Real capture consumed 0.20 GiB (~195 MiB deficit against the 10 MiB estimate).

On this model, the 1 MiB floor under-budgets for both dtypes, not just FP8. That supports your point: a static 16 MiB floor is an arbitrary constant and doesn't address why the samples collapse.

Two ways forward:

1. Full capture during profiling (_max_full_descs_to_capture = None): Measures all FULL graphs directly, so no heuristic is needed. Capturing all 35 graphs took ~4s on this model during engine warmup.
2. Floor scaled by layer count and KV dtype: Keeps profiling fast, but it is still a heuristic, and I'd want several models measured across different depths before trusting any formula.

If you prefer Option 1, I'm happy to rework this PR into that and validate it thoroughly on the RTX 4090. Let me know which direction you'd like!

[lucamotz]

I don't think those are the right approaches. #56492 already covers this long term, and if we really wanted a fix today, I think this would be the way to go:

--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -488,4 +488,6 @@
                         if self._capture_mem_samples is not None:
                             torch.accelerator.synchronize()
+                            # torch.cuda.graph() empties the cache on entry.
+                            torch.accelerator.empty_cache()
                             free_before = torch.accelerator.get_memory_info()[0]
                         with (

torch.cuda.graph() calls empty_cache() on entry, so the memory cached by the warmup forward is currently released inside the sampled window. That would explain the negative samples, for BF16 as well. MRV1 reads the free memory before the warmup and doesn't have this.

But in general, I'm not sure this issue even exists in the way it's being described. It was run on v0.28.0, where MRV2 didn't reserve any CUDA graph memory yet (#53306 added that in v0.29.0), and it seems to be a bit of an edge case anyways. I would want to see the OOM reproduced on current main first. Otherwise I think closing this in favour of #56492 is the right way to go here.

[Aravind0403]

Thanks @lucamotz! I just verified your hypothesis on the 24GB RTX 4090:

1. empty_cache() fix verified: Adding torch.accelerator.empty_cache() right before reading free_before flipped mem_samples from negative (-22 MiB) to positive (+4.00 MiB), confirming that torch.cuda.graph()'s internal cache flush was distorting the measurement window.
2. OOM reproduction on current main: Tested Qwen/Qwen2.5-0.5B-Instruct with kv_cache_dtype="fp8" at gpu_memory_utilization=0.95 on current main, and it initialized 3.52M KV tokens and completed inference with zero OOM.

Given that the original issue was tied to v0.28.0 before graph memory reservation existed, and #56492 (extensible KV cache) dynamically manages graph and cache allocation, I'll close this in favor of #56492.

Thanks for the great review and the deep dive into the allocator mechanics!