diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 8e49cf7a6f2e..ef38f1650418 100644 --- a/tests/utils_/test_torch_utils.py +++ b/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 ( diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index b2096207d1e9..d808f36b0676 100644 --- a/vllm/utils/torch_utils.py +++ b/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: