A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 14:30 UTC
[Bugfix] Fix unclean shutdown crash with AllReduce Fusion workspace (#36955)
Signed-off-by: Jarvis <brayden.stanley.0127@gmail.com>
Siew's Capital Jarvis committed
Mar 17, 2026 at 22:22 UTC
c25dbc2d2728621385760d1c98bda6200f545900
1 file changed
+19
-9
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+19
-9
@@ -2,6 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
+import atexit
6
+import threading
7
+
8
import torch
9
import torch.distributed as dist
10
from torch.distributed import ProcessGroup
@@ -132,18 +135,25 @@ def initialize_fi_ar_quant_workspace(
135
)
136
137
138
+_fi_ar_workspace_lock = threading.Lock()
139
+
140
+
141
def destroy_fi_ar_workspace():
142
global _fi_ar_workspace
143
global _fi_ar_quant_workspace
138
- if (
139
- _fi_ar_quant_workspace is not None
140
- and _fi_ar_quant_workspace is not _fi_ar_workspace
141
- ):
142
- _fi_ar_quant_workspace.destroy()
143
- _fi_ar_quant_workspace = None
144
- if _fi_ar_workspace is not None:
145
- _fi_ar_workspace.destroy()
146
- _fi_ar_workspace = None
144
+ with _fi_ar_workspace_lock:
145
+ if (
146
+ _fi_ar_quant_workspace is not None
147
+ and _fi_ar_quant_workspace is not _fi_ar_workspace
148
+ ):
149
+ _fi_ar_quant_workspace.destroy()
150
+ _fi_ar_quant_workspace = None
151
+ if _fi_ar_workspace is not None:
152
+ _fi_ar_workspace.destroy()
153
+ _fi_ar_workspace = None
154
+
155
+
156
+atexit.register(destroy_fi_ar_workspace)
157
158
159
class FlashInferAllReduce: