[ZJY0516]

Overview

weak_ref_tensors only weak-refs one level. For a model output like (hidden_states, [aux_hidden_states, ...]), the inner list is returned as is, so every captured CUDA graph entry keeps its aux hidden states alive and the graph pool cannot reuse that memory for the next capture size. This PR makes weak_ref_tensors recurse into nested lists/tuples.

Related: #63406 lists this as "aux hidden states kept by every PIECEWISE size" in its memory breakdown. This PR is the standalone fix and also applies to main's capture sizes.

Claims

• Captured graphs no longer pin the aux hidden states used by EAGLE3, DSpark, DFlash and extract_hidden_states. Before, each captured size leaked num_aux × num_tokens × hidden_size × dtype_size bytes (3 × 8192 × 4096 × 2 B = 192 MiB for one 8192-token size of DeepSeek-V4-Flash).
• No output change: the outputs alias the same memory as before; only the references held by the graph entries become weak.
• weak_ref_tensors now raises on unknown types at any nesting level, not only at the top level, and passes ints and LayerNames through (piecewise subgraph outputs can contain ints, and Inductor graph partition outputs can contain LayerNames, see below).

Validation

Unit test (tests/utils_/test_torch_utils.py, run by the "Async Engine, Inputs, Utils, Worker" CI step), on GB200:

# with this PR
pytest tests/utils_/test_torch_utils.py -v
39 passed

# torch_utils.py reverted to main, new tests only
pytest tests/utils_/test_torch_utils.py -v -k weak_ref
FAILED test_weak_ref_tensors_does_not_keep_nested_tensors_alive
E       AssertionError: assert 4194304 == (8388608 - 8388608)
FAILED test_weak_ref_tensors_rejects_unknown_nested_types
E       Failed: DID NOT RAISE ValueError

What actually reaches weak_ref_tensors. I wrapped the function in vllm/compilation/cuda_graph.py and vllm/compilation/breakable_cudagraph.py, then started each engine and recorded the structure of every input during capture (T = tensor; [n:...] = list of length n; (n:...) = tuple of length n):

Model	CUDA graph path	Inputs seen
Qwen3-1.7B + EAGLE3	torch.compile piecewise	[1:T] [2:T] [4:T] [5:T] [6:T]
Qwen3-1.7B + EAGLE3	breakable	target (2:T,[3:T]), drafter (2:T)
gpt-oss-20b + EAGLE3	torch.compile piecewise	[1:T] [2:T] [5:T] [6:T]
gpt-oss-20b + EAGLE3	breakable	target (2:T,[3:T]), drafter (2:T)
Qwen3-30B-A3B-FP8	torch.compile piecewise	[1:T] [5:T] [6:T,int]
Qwen3-30B-A3B-FP8	breakable	T

The nested case is the breakable target output with aux hidden states. The int case is a piecewise subgraph output of the FP8 MoE model, which is why ints are passed through. No None, float or 0-size tensor showed up.

The table does not cover use_inductor_graph_partition=True. There, on torch >= 2.11, CUDAGraphWrapper wraps each Inductor partition, and a partition's outputs can include LayerName objects that torch.compile hoists as graph inputs. This PR's first CI run caught it (ValueError: Invalid type for tensors: <class 'vllm.utils.torch_utils.LayerName'> in test_full_cudagraph.py[*-partition=True] and test_full_graph.py::test_custom_compile_config), so LayerNames pass through as well. After the fix, on GB200:

pytest compile/fullgraph/test_full_cudagraph.py -k "True and not MLA"
30 passed, 20 skipped (FA3, RocmAttn)

pytest compile/fullgraph/test_full_graph.py -k "compilation_config4 or compilation_config5"
2 passed

End-to-end A/B: DeepSeek-V4-Flash-DSpark, TP4+EP, DSpark K=7, 4×GB200. Breakable CUDA graph is on by default for this model, and capture sizes are main's defaults. Main is this PR's parent (2c58be2). Memory numbers are from TP rank 0.

	main	this PR	Δ
CUDA graph pool, actual	5.86 GiB	5.39 GiB	−0.47 GiB
CUDA graph pool, profiler estimate	8.55 GiB	7.90 GiB	−0.65 GiB
Available KV cache memory	112.94 GiB	113.59 GiB	+0.65 GiB
GPU KV cache size	7,277,119 tokens	7,319,402 tokens	+0.58%
GSM8K 5-shot, flexible-extract	0.9560	0.9522	−5 / 1319 questions
DSpark acceptance (accepted / drafted)	57.74%	57.84%
Mean acceptance length	5.04	5.05

The leak on main is about 0.65 GiB here, and it grows with the number and size of PIECEWISE capture sizes (#63406 raises both).

The 5-question GSM8K gap is from a single run per arm; serving is not batch invariant, so some run-to-run variance is expected. It does not indicate corrupted aux hidden states: they only feed the drafter, so corruption would lower acceptance, and acceptance is unchanged. The target model alone decides the output tokens under greedy rejection sampling.

Serve command

vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \
  --tensor-parallel-size 4 --enable-expert-parallel \
  --trust-remote-code --kv-cache-dtype fp8 --block-size 256 \
  --tokenizer-mode deepseek_v4 --reasoning-parser deepseek_v4 \
  --speculative-config '{"method":"dspark","num_speculative_tokens":7,"draft_sample_method":"greedy"}'

GSM8K: 5-shot, temperature 0, chat completions, thinking off. Acceptance is read from vllm:spec_decode_* metrics after the GSM8K run. Memory numbers come from the CUDA graph pool memory, Available KV cache memory and GPU KV cache size log lines.

Details

Affected paths. weak_ref_tensors is called on capture outputs by:

• BreakableCUDAGraphWrapper._capture: captures the whole model forward, so EAGLE3/DSpark targets return (hidden_states, [aux, ...]). Breakable CUDA graph is on by default for DeepSeek-V4/V4.1, MiniMax-M3, GLM-5-Next, Kimi-K3 and others, under both model runners.
• CUDAGraphWrapper in FULL mode on MRV1, which also wraps the whole model.

Not affected: torch.compile piecewise subgraph outputs are flat, and MRV2 FULL graphs copy outputs into persistent buffers instead of weak-refing them.

Safety of the weak refs. Once the aux tensors are weak, later captures may reuse their memory, the same as for hidden_states today. So the aux tensors must be consumed before the next graph replay. Every consumer copies them first: MRV1 torch.cat/torch.stack in gpu_model_runner.py, MRV2 DFlash _store_hidden_states (concatenates into its staging buffer), MRV2 extract_hidden_states (torch.stack + copy_).

Type contract. Previously, non-tensor values inside a list/tuple passed through weak_ref_tensor silently, which is how the nested list slipped through. Now containers recurse, ints and LayerNames pass, and any other type raises with its type name. The tradeoff: a model whose capture output contains a leaf type not in the table above (e.g. None) fails at capture time instead of passing silently. The raise already caught one such case before (IntermediateTensors, #24275), and caught LayerName in this PR's first CI run.

Not a duplicate: no open PR touches weak_ref_tensors or this leak. #63406 changes capture sizes and only references the leak in its memory accounting.

The fix and analysis originate from Bo Pang (co-author). AI assistance (Claude Code) was used to write the change, the tests and the validation scripts.

---

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.

🤖 Generated with Claude Code

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

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-23T23:48:13.702248Z	a06d7fe	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a06d7fe559

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add an LLM-level regression for nested graph outputs

The new regression invokes only weak_ref_tensors, so it never traverses either CUDA-graph wrapper, captures multiple sizes, or consumes nested auxiliary outputs through serving—the user-visible path whose per-entry retention this fix targets. This leaves the wrapper integration unguarded; add an LLM/vllm_runner regression with real weights that exercises nested aux outputs and verifies the capture-memory behavior.

AGENTS.md reference: tests/AGENTS.md:L8-L10

Useful? React with 👍 / 👎.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97427 for commit f7443a509505.

[ZJY0516]

/ci run --allow-stale

[github-actions]

✅ Triggered Buildkite CI #97453 for commit 6ee3f40d6137.

⚠️ This PR is 3 commits behind upstream main. Running CI at your own risk because --allow-stale was requested; outdated CI configuration may cause failures. Before merging, merge or rebase onto the latest main, then rerun /ci run on the latest PR commit.