[Bugfix][Distributed] Delegate MNNVL allreduce one-shot selection (#47589)
Signed-off-by: jesco-absolut <team@srswti.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
jesco committed
Jul 6, 2026 at 10:47 UTC
598d51153a15ff4a803037dade1c4ea43d1ea9bc
2 files changed
+60
-13
tests/compile/passes/distributed/test_fusion_all_reduce.py
+30
@@ -13,6 +13,7 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
13
from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
14
AllReduceFusionPass,
15
RocmAiterAllReduceFusionPass,
16
+ _select_flashinfer_allreduce_use_oneshot,
17
)
18
from vllm.compilation.passes.fx_utils import find_op_nodes
19
from vllm.compilation.passes.utility.fix_functionalization import (
@@ -48,6 +49,35 @@ from vllm.utils.torch_utils import set_random_seed
49
DEVICE_TYPE = current_platform.device_type
50
51
52
+@pytest.mark.parametrize(
53
+ ("workspace_backend", "device_capability", "world_size", "tensor_size", "expected"),
54
+ [
55
+ ("mnnvl", 103, 8, 2 * 1024 * 1024, None),
56
+ ("trtllm", 103, 8, 2 * 1024 * 1024, True),
57
+ ("trtllm", 103, 8, 2 * 1024 * 1024 + 1, False),
58
+ ("trtllm", 100, 4, 4 * 1024 * 1024, True),
59
+ ("trtllm", 100, 4, 4 * 1024 * 1024 + 1, False),
60
+ ("trtllm", None, 8, 128 * 1024 * 1024, True),
61
+ ],
62
+)
63
+def test_select_flashinfer_allreduce_use_oneshot(
64
+ workspace_backend: str,
65
+ device_capability: int | None,
66
+ world_size: int,
67
+ tensor_size: int,
68
+ expected: bool | None,
69
+):
70
+ assert (
71
+ _select_flashinfer_allreduce_use_oneshot(
72
+ workspace_backend,
73
+ device_capability,
74
+ world_size,
75
+ tensor_size,
76
+ )
77
+ is expected
78
+ )
79
+
80
+
81
class TestAllReduceRMSNormModel(torch.nn.Module):
82
def __init__(
83
self,
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+30
-13
@@ -128,6 +128,29 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {
128
},
129
}
130
131
+MiB = 1024 * 1024
132
+
133
+
134
+def _select_flashinfer_allreduce_use_oneshot(
135
+ workspace_backend: str,
136
+ device_capability: int | None,
137
+ world_size: int,
138
+ current_tensor_size: int,
139
+) -> bool | None:
140
+ if workspace_backend == "mnnvl":
141
+ # FlashInfer sizes MNNVL workspaces around its own AUTO strategy.
142
+ # Forcing vLLM's per-rank threshold can request one-shot for tensors
143
+ # larger than the MNNVL one-shot workspace.
144
+ return None
145
+
146
+ if device_capability is None:
147
+ max_one_shot_size = None
148
+ else:
149
+ max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
150
+ device_capability, {}
151
+ ).get(world_size)
152
+ return max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB
153
+
154
155
if flashinfer_comm is not None:
156
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
@@ -138,8 +161,6 @@ if flashinfer_comm is not None:
161
162
ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern
163
141
- MiB = 1024 * 1024
142
-
164
def call_trtllm_fused_allreduce_norm(
165
allreduce_in: torch.Tensor,
166
residual: torch.Tensor,
@@ -174,16 +195,6 @@ if flashinfer_comm is not None:
195
)
196
curr_device = current_platform.get_device_capability()
197
device_capability = curr_device.to_int() if curr_device is not None else None
177
- # Get one shot input size limit for the current world size
178
- # for the current device capability
179
- max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
180
- device_capability, # type: ignore[arg-type, unused-ignore]
181
- {},
182
- ).get(world_size, None)
183
- # Use one shot if no max size is specified
184
- use_oneshot = (
185
- max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB
186
- )
198
199
# Select workspace based on pattern: quant patterns use the
200
# trtllm quant workspace, non-quant patterns use the primary workspace.
@@ -205,6 +216,12 @@ if flashinfer_comm is not None:
216
assert workspace is not None, (
217
"Flashinfer allreduce workspace must be initialized when using flashinfer"
218
)
219
+ use_oneshot = _select_flashinfer_allreduce_use_oneshot(
220
+ workspace.backend,
221
+ device_capability,
222
+ world_size,
223
+ current_tensor_size,
224
+ )
225
assert flashinfer_comm is not None
226
if norm_out is None:
227
norm_out = allreduce_in
@@ -248,7 +265,7 @@ if flashinfer_comm is not None:
265
# the end for the one-shot path; the two-shot path is synchronized
266
# and keeps the early completion. Related one-shot instability in
267
# the same kernel: flashinfer-ai/flashinfer#1223.
251
- trigger_completion_at_end=use_oneshot
268
+ trigger_completion_at_end=(use_oneshot is True)
269
or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
270
)
271