[Bugfix] Don't sync-police or retry FlashInfer all-reduce workspace creation in eager mode (#58498)
Signed-off-by: khluu <khluu000@gmail.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Kevin H. Luu committed
Sep 28, 2026 at 12:14 UTC
764413559a1416511002268ac460bbf8e7ddd394
2 files changed
+56
-14
tests/distributed/test_comm_ops.py
+30
@@ -389,6 +389,36 @@ def test_flashinfer_standalone_workspace_size(
389
assert create_workspace.call_args.args[3] == expected
390
391
392
+def test_flashinfer_workspace_failure_is_not_retried(
393
+ monkeypatch: pytest.MonkeyPatch,
394
+) -> None:
395
+ """Creation is collective; retrying it per all-reduce desyncs ranks."""
396
+ create_workspace = Mock(return_value=None)
397
+ monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", None)
398
+ monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", None)
399
+ monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace_failed", False)
400
+ monkeypatch.setattr(
401
+ flashinfer_all_reduce,
402
+ "_resolve_fi_ar_backend",
403
+ Mock(return_value=("mnnvl", True)),
404
+ )
405
+ monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 1)
406
+ monkeypatch.setattr(
407
+ flashinfer_all_reduce, "_get_tuned_standalone_max_size", Mock(return_value=None)
408
+ )
409
+ monkeypatch.setattr(flashinfer_all_reduce, "_create_workspace", create_workspace)
410
+ args = (2, 0, 128, 4096, torch.bfloat16, Mock())
411
+
412
+ assert flashinfer_all_reduce.get_fi_ar_workspace(*args) is None
413
+ assert flashinfer_all_reduce.get_fi_ar_workspace(*args) is None
414
+ # One mnnvl attempt plus one trtllm fallback, never repeated.
415
+ assert create_workspace.call_count == 2
416
+
417
+ flashinfer_all_reduce.destroy_fi_ar_workspace()
418
+ flashinfer_all_reduce.get_fi_ar_workspace(*args)
419
+ assert create_workspace.call_count == 4
420
+
421
+
422
def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) -> None:
423
output = torch.empty(2)
424
fi_ar_comm = Mock(disabled=False)
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+26
-14
@@ -20,6 +20,7 @@ from vllm.distributed.device_communicators.all_reduce_utils import (
20
from vllm.distributed.parallel_state import _node_count, get_node_count
21
from vllm.logger import init_logger
22
from vllm.platforms import current_platform
23
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
24
25
logger = init_logger(__name__)
26
@@ -45,6 +46,10 @@ _fi_ar_workspace = None
46
# allreduce backend or a fallback backend when the primary workspace is not
47
# available on the current topology.
48
_fi_ar_quant_workspace = None
49
+# Creation is collective, so a failed attempt must not be retried on every
50
+# all-reduce call.
51
+_fi_ar_workspace_failed = False
52
+_fi_ar_quant_workspace_failed = False
53
_fi_ar_workspace_groups: dict[int, ProcessGroup] = {}
54
55
@@ -79,16 +84,19 @@ def _create_workspace(
84
rng_state = random.getstate()
85
try:
86
random.seed(int.from_bytes(os.urandom(16), byteorder="big"))
82
- workspace = flashinfer_comm.create_allreduce_fusion_workspace(
83
- backend=backend,
84
- world_size=world_size,
85
- rank=rank,
86
- max_token_num=max_token_num,
87
- hidden_dim=hidden_dim,
88
- dtype=dtype,
89
- comm_backend=comm_backend,
90
- group=group,
91
- )
87
+ # Creation may run lazily inside the first sync-checked forward (e.g.
88
+ # with enforce_eager, which skips warmup).
89
+ with gpu_sync_allowed(first_only=True):
90
+ workspace = flashinfer_comm.create_allreduce_fusion_workspace(
91
+ backend=backend,
92
+ world_size=world_size,
93
+ rank=rank,
94
+ max_token_num=max_token_num,
95
+ hidden_dim=hidden_dim,
96
+ dtype=dtype,
97
+ comm_backend=comm_backend,
98
+ group=group,
99
+ )
100
if backend == "mnnvl" and not getattr(workspace, "mc_ptr", 0):
101
workspace.destroy()
102
logger.warning_once(
@@ -175,8 +183,8 @@ def get_fi_ar_workspace(
183
for standalone allreduce. Backend is controlled by
184
VLLM_FLASHINFER_ALLREDUCE_BACKEND env var.
185
"""
178
- global _fi_ar_workspace
179
- if _fi_ar_workspace is not None:
186
+ global _fi_ar_workspace, _fi_ar_workspace_failed
187
+ if _fi_ar_workspace is not None or _fi_ar_workspace_failed:
188
return _fi_ar_workspace
189
190
backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
@@ -218,6 +226,7 @@ def get_fi_ar_workspace(
226
f"with backend={backend}"
227
)
228
else:
229
+ _fi_ar_workspace_failed = True
230
logger.error_once(
231
"Failed to initialize FlashInfer Allreduce norm fusion workspace "
232
f"with backend={backend}"
@@ -240,8 +249,8 @@ def get_fi_ar_quant_workspace(
249
non-quant fusion. With ``auto`` this prefers mnnvl and falls back to trtllm
250
only on single-node topologies where mnnvl multicast is unavailable.
251
"""
243
- global _fi_ar_quant_workspace
244
- if _fi_ar_quant_workspace is not None:
252
+ global _fi_ar_quant_workspace, _fi_ar_quant_workspace_failed
253
+ if _fi_ar_quant_workspace is not None or _fi_ar_quant_workspace_failed:
254
return _fi_ar_quant_workspace
255
256
backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
@@ -291,6 +300,7 @@ def get_fi_ar_quant_workspace(
300
f"fusion workspace with backend={backend}"
301
)
302
else:
303
+ _fi_ar_quant_workspace_failed = True
304
logger.error_once(
305
"Failed to initialize FlashInfer Allreduce norm quantization "
306
f"fusion workspace with backend={backend}"
@@ -304,6 +314,7 @@ _fi_ar_workspace_lock = threading.Lock()
314
315
def destroy_fi_ar_workspace():
316
global _fi_ar_workspace, _fi_ar_quant_workspace
317
+ global _fi_ar_workspace_failed, _fi_ar_quant_workspace_failed
318
with _fi_ar_workspace_lock:
319
is_alias = _fi_ar_workspace is _fi_ar_quant_workspace
320
@@ -313,6 +324,7 @@ def destroy_fi_ar_workspace():
324
_fi_ar_quant_workspace.destroy()
325
326
_fi_ar_workspace = _fi_ar_quant_workspace = None
327
+ _fi_ar_workspace_failed = _fi_ar_quant_workspace_failed = False
328
_fi_ar_workspace_groups.clear()
329
330