From 064a3e1b74a6dc6b72d0be9fe89e38db0796d6e4 Mon Sep 17 00:00:00 2001 From: hid64306686 Date: Tue, 24 Feb 2026 23:36:48 +0800 Subject: [PATCH] dynamic import current_stream --- checkpoint_engine/distributed/vllm_nccl.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/checkpoint_engine/distributed/vllm_nccl.py b/checkpoint_engine/distributed/vllm_nccl.py index 2ffe253..b2eb1aa 100644 --- a/checkpoint_engine/distributed/vllm_nccl.py +++ b/checkpoint_engine/distributed/vllm_nccl.py @@ -12,11 +12,21 @@ ncclResult_t, ) from vllm.distributed.utils import StatelessProcessGroup -from vllm.utils import current_stream from checkpoint_engine.distributed.base import CommGroup, Distributed, _common_all_gather_object +try: + from vllm.utils.torch_utils import current_stream +except ImportError: + try: + from vllm.utils import current_stream + except ImportError: + raise ImportError( + "Could not find 'current_stream' in vllm. Please check your vllm version." + ) from None + + class NcclConfigT(ctypes.Structure): _fields_: ClassVar[list[tuple[str, Any]]] = [ ("size", ctypes.c_size_t),