[Memory] Measure complete CUDA graph capture footprint for KV budgeting

[xiaohuguo2023]

Purpose

Fix CUDA graph memory undercounting in profile_cudagraph_memory().

The old path profiled only two descriptors per graph mode and extrapolated the rest. It also summed per-mode deltas even though FULL and PIECEWISE share one runtime pool, and it did not budget memory allocated during profiling setup or retained after profiling cleanup. The shortfall was assigned to the KV cache, so --gpu-memory-utilization could be exceeded during capture or serving while the KV pool was still mostly empty.

This PR:

• Profiles every capture descriptor in one end-to-end span across all decoder modes instead of extrapolating from two samples per mode.
• Includes profiling KV-cache / attention-metadata setup and cudagraph_profiling_retained_memory in the estimate (scratch the real capture reuses).
• Skips the profiling capture entirely when VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 (previously ran and discarded the result).
• Guards the unaccounted-graph-memory warning so it only fires when graphs will actually be captured (--enforce-eager / cudagraph_mode=NONE no longer trigger it).
• Preserves cleanup and graph-pool restoration when profiling raises.

On Kimi-K3 (8× MI355X, TP8), the estimate rose from ~1.3 GiB/GPU to ~4.7–16.9 GiB/GPU depending on config — matching what capture actually needs. This is accounting only; no kernel or graph-capture logic changes.

Fixes agentic OOMs where HSA_STATUS_ERROR_OUT_OF_RESOURCES occurred at 3–16% KV usage because graph/setup memory was budgeted to the KV cache.

Test Plan

Unit tests (CPU-only; fake device allocator, no GPU required):

pytest tests/v1/cudagraph/test_cudagraph_memory_profiling.py \
       tests/v1/cudagraph/test_cudagraph_memory_worker.py -q

Coverage:

• test_cudagraph_memory_profiling.py — all descriptors and modes; shared-pool span accounting; negative deltas and zero clamping; encoder graphs; profiling-KV setup; retained memory; cleanup on capture failure.
• test_cudagraph_memory_worker.py — estimate profiled and subtracted from KV budget; VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 skips profiling; no warning when cudagraph_mode=NONE.

Serving smoke (ROCm, optional for reviewers):

• Kimi-K3, TP8, FP8 KV, FULL_AND_PIECEWISE, gpu-memory-utilization=0.95
• Confirm log line Estimated CUDA graph memory: … GiB total reflects full capture footprint and serve completes without HSA OOM at previously failing batch/seq limits.

Test Result

Unit tests (this branch):

13 passed in 5.89s

Negative control: 8 behavior-changing tests fail on pristine upstream/main.

Serving (Kimi-K3, MI355X TP8):

config	metric	before (old sampler)	after (this PR)
cm16k s128, c16, util 0.95	graph memory estimate	~1.3 GiB/GPU	4.7–7.1 GiB/GPU
cm16k s128, c16, util 0.95	agentic output tok/s	61	227 (+272%)
cm16k s128, c16, util 0.95	TTFT / ITL avg	56s / 187ms	2.9s / 50ms
cm8k s128, c24, util 0.95	graph memory estimate	~1.3 GiB/GPU	16.90 GiB/GPU
cm8k s128, c24, util 0.95	run health	OOM (transient workspace)	440 reqs, 0.23% err
cm8k s128, c24, util 0.93	GPU KV cache size	—	3,660,123 tokens, clean run

Startup: profiling every descriptor adds ~8s on Kimi-K3 TP8 (21s → 29s for profiling + capture). Warmup reuse across profiling and real capture is intentionally deferred to a follow-up PR.

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