biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Jarvis <brayden.stanley.0127@gmail.com>
This commit is contained in:
Siew's Capital Jarvis
2026-03-17 22:22:09 +08:00
committed by GitHub
parent 77d2a5f17b
commit c25dbc2d27
1 changed files with 19 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
vllm/distributed/device_communicators/flashinfer_all_reduce.py
View File

@@ -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: