[Core] Pre-size cudagraph output staging buffers to the max capture descriptor

[matteso1]

Purpose

Capture order came up in the #feat-startup-ux discussions about parallelizing and reordering CUDA-graph capture. @galv raised that capture order affects pool fragmentation. The in-tree comment in CudaGraphManager.capture() already orders PIECEWISE before FULL for pool-reuse reasons. So ordering is a real memory-tuning knob.

Today it is also a silent correctness invariant. You cannot measure fragmentation under a different order, because any non-descending order crashes the boot.

The mechanism: ModelCudaGraphManager allocates its output staging buffers (hidden_states, aux_hidden_states, intermediate_tensors) lazily inside the first warmup forward, through torch.empty_like. Their capacity becomes whatever num_tokens the first descriptor happens to have. The descending sort in _init_candidates masks this, because the first descriptor is the max and everything fits. Under any smaller-first order, a later and larger warmup crashes at the staging copy self.hidden_states[:num_tokens] = hidden_states with a shape mismatch.

Measured on a vLLM 0.24.0 wheel, H100 NVL, Qwen3-8B. The lazy-allocation pattern is unchanged on current main.

• Stock: ascending order and three random shuffled orders crash the boot, 8 of 8 attempts, with a deterministic shape mismatch at the staging copy. This reproduces on the vllm serve path and cross-model on Qwen3-0.6B.
• Falsification check on the mechanism: a max-first then ascending order boots fine. So the invariant is that the first descriptor must be the largest. That points at this buffer, not at the memory pool.
• With max-descriptor pre-sizing, all 12 formerly-crashing boots succeed (4 orders, N=3). That enabled the first capture-order fragmentation measurement: a 2.91 GiB pool gap on the default descending order, 3.1 to 3.4 GiB on shuffles, and 7.52 GiB on strict ascending, which is 2.6 times the default. Capture time is order-independent at about 3.9 s in every order. Raw receipts (per-boot JSON, logs, hardware fingerprint) are available if useful.

So a bad order costs memory, not seconds. With this fix it costs memory measurably instead of crashing.

Changes

• CudaGraphManager._staging_buffer_tokens() returns the max num_tokens across all registered capture descriptors. All of them are known at __init__ time.
• The three lazy staging allocations size their leading dimension from it. The dtype, device, and trailing dimensions still come from the live tensor through new_empty.

Allocations are byte-identical under the current descending order. The sizing rule just stops depending on iteration order. This holds in every reachable config, not only by luck of the sort. PIECEWISE descriptors get the raw cudagraph_capture_sizes. Decode FULL descriptors are filtered by rounded_num_tokens > max_cg_capture_size. compilation.py asserts cudagraph_capture_sizes[-1] == max_cudagraph_capture_size. So the first descriptor to reach the staging allocation is already the global max today, and there is zero capture-time memory delta.

What this removes is the hidden coupling. Future capture-order changes, such as fragmentation tuning or parallel-capture experiments, then fail by measurement rather than by boot crash.

V1 and V2 note: this is the Model Runner V2 manager in vllm/v1/worker/gpu/. The V1 path stages outputs per-graph inside CUDAGraphWrapper and does not share the first-descriptor sizing pattern.

Test Plan

Folded into tests/v1/spec_decode/test_dynamic_sd_cug.py, which exercises the MRv2 CudaGraphManager directly. Correction to an earlier version of this description: that is not the only test file importing vllm.v1.worker.gpu.cudagraph_utils. tests/v1/cudagraph/test_breakable_cudagraph.py also does, and tests/v1/cudagraph/test_cudagraph_manager.py arrived on 2026-07-21 with #48843. tests/v1/cudagraph/test_cudagraph_manager.py is arguably the better home now, and I am happy to move these cases there if reviewers prefer it.

The added cases cover three things: capacity equals the global max regardless of which descriptor allocates first, a descriptor larger than every registered one still fits, and the empty-descriptor fallback.

Test Result

I verified the logic locally against the extracted helper, because this machine has no GPU dependencies available. The folded test is CPU-only and runs in CI.

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

[matteso1]

Closing this as a standalone change. Current main deliberately captures the largest descriptors first, so the existing order already avoids the under-allocation. This patch makes alternate orders safe, but it does not improve latency or memory under the current order. I will bring it back with the capture-order optimization if that work produces a measured user-visible benefit.