diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index ea16c9376..1152277f7 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import atexit +import threading + import torch import torch.distributed as dist from torch.distributed import ProcessGroup @@ -132,18 +135,25 @@ def initialize_fi_ar_quant_workspace( ) +_fi_ar_workspace_lock = threading.Lock() + + def destroy_fi_ar_workspace(): global _fi_ar_workspace global _fi_ar_quant_workspace - if ( - _fi_ar_quant_workspace is not None - and _fi_ar_quant_workspace is not _fi_ar_workspace - ): - _fi_ar_quant_workspace.destroy() - _fi_ar_quant_workspace = None - if _fi_ar_workspace is not None: - _fi_ar_workspace.destroy() - _fi_ar_workspace = None + with _fi_ar_workspace_lock: + if ( + _fi_ar_quant_workspace is not None + and _fi_ar_quant_workspace is not _fi_ar_workspace + ): + _fi_ar_quant_workspace.destroy() + _fi_ar_quant_workspace = None + if _fi_ar_workspace is not None: + _fi_ar_workspace.destroy() + _fi_ar_workspace = None + + +atexit.register(destroy_fi_ar_workspace) class FlashInferAllReduce: