[Bugfix][CUDA Graph] Weak-ref nested outputs of captured graphs

tests/utils_/test_torch_utils.py

@@ -6,6 +6,7 @@
 
 from vllm.utils.torch_utils import (
     OMP_NUM_THREADS_SET_BY_VLLM,
+    LayerName,
     async_tensor_h2d,
     available_cpu_count,
     common_broadcastable_dtype,
@@ -16,9 +17,34 @@
     set_random_seed,
     set_torch_threads_for_runtime,
     startup_omp_num_threads,
+    weak_ref_tensors,
 )
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_weak_ref_tensors_does_not_keep_nested_tensors_alive() -> None:
+    # CUDA graph wrappers weak-ref model outputs such as
+    # (hidden_states, [aux_hidden_states, ...]); no tensor may stay pinned.
+    # Inductor graph partition outputs can also carry hoisted LayerNames.
+    hidden, aux = (torch.empty(1 << 20, device="cuda") for _ in range(2))
+    layer_name = LayerName("model.layers.0.self_attn.attn")
+    out = weak_ref_tensors((hidden, [aux], 3, layer_name))
+    assert isinstance(out, tuple)
+    weak_hidden, weak_list, size, name = out
+    assert isinstance(weak_list, list) and size == 3 and name is layer_name
+    assert weak_hidden.data_ptr() == hidden.data_ptr()
+    assert weak_list[0].data_ptr() == aux.data_ptr()
+    nbytes = hidden.nbytes + aux.nbytes
+    allocated = torch.accelerator.memory_allocated()
+    del hidden, aux
+    assert torch.accelerator.memory_allocated() == allocated - nbytes
+
+
+def test_weak_ref_tensors_rejects_unknown_nested_types() -> None:
+    with pytest.raises(ValueError):
+        weak_ref_tensors((1, [None]))
+
+
 def test_nvfp4_4over6_cache_dtype() -> None:
     from vllm.config.cache import CacheConfig
     from vllm.v1.kv_cache_interface import (

vllm/utils/torch_utils.py

@@ -887,14 +887,17 @@ def weak_ref_tensors(
     | IntermediateTensors,
 ) -> torch.Tensor | list[Any] | tuple[Any] | Any:
     """Convenience function to create weak references to tensors,
-    for single tensor, list of tensors or tuple of tensors.
+    for single tensor, possibly nested list/tuple of tensors, or
+    IntermediateTensors. Ints and LayerNames are returned as is.
     """
     if isinstance(tensors, torch.Tensor):
         return weak_ref_tensor(tensors)
     if isinstance(tensors, list):
-        return [weak_ref_tensor(t) for t in tensors]
+        return [weak_ref_tensors(t) for t in tensors]
     if isinstance(tensors, tuple):
-        return tuple(weak_ref_tensor(t) for t in tensors)
+        return tuple(weak_ref_tensors(t) for t in tensors)
+    if isinstance(tensors, (int, LayerName)):
+        return tensors
 
     # For IntermediateTensors used in pipeline parallelism
     from vllm.sequence import IntermediateTensors
@@ -904,7 +907,7 @@ def weak_ref_tensors(
             {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}
         )
         return ret
-    raise ValueError("Invalid type for tensors")
+    raise ValueError(f"Invalid type for tensors: {type(tensors)}")
 
 
 def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor: