[Bugfix] Fix unclean shutdown crash with AllReduce Fusion workspace (#36955)

Signed-off-by: Jarvis <brayden.stanley.0127@gmail.com>

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: