[ROCm] Work around DeepEP teardown SIGSEGV in MoE test harness (#51174)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rohan Potdar committed
Aug 5, 2026 at 14:44 UTC
c2d8009044abb174576dab6439dda56d66c83847
1 file changed
+12
tests/kernels/moe/parallel_utils.py
+12
@@ -6,6 +6,7 @@ DeepEP test utilities
6
7
import dataclasses
8
import os
9
+import sys
10
import traceback
11
from collections.abc import Callable
12
from typing import Concatenate
@@ -15,6 +16,7 @@ from torch.distributed import ProcessGroup
16
from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
17
from typing_extensions import ParamSpec
18
19
+from vllm.platforms import current_platform
20
from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
21
from vllm.utils.network_utils import get_open_port
22
@@ -72,6 +74,7 @@ def _worker_parallel_launch(
74
barrier = torch.tensor([rank], device=device)
75
torch.distributed.all_reduce(barrier)
76
77
+ exit_code = 0
78
try:
79
worker(
80
ProcessGroupInfo(
@@ -88,9 +91,18 @@ def _worker_parallel_launch(
91
except Exception as ex:
92
print(ex)
93
traceback.print_exc()
94
+ exit_code = 1
95
raise
96
finally:
97
torch.distributed.destroy_process_group()
98
+ if current_platform.is_rocm():
99
+ # Bypass a ROCm teardown use-after-free (SIGSEGV in HIP's atexit
100
+ # handler; fixed by https://github.com/ROCm/rocm-systems/pull/6942).
101
+ # Flush first since os._exit skips it, to keep a failing traceback.
102
+ # TODO (Rohan138): Remove after ROCm 10.0 lands
103
+ sys.stdout.flush()
104
+ sys.stderr.flush()
105
+ os._exit(exit_code)
106
107
108
def parallel_launch(