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, 09:30 UTC
[Bug] Fix FlashInfer allreduce fusion workspace uninitialized error (#37461)
Signed-off-by: root <root@prenyx0169.a51.clusters.nvidia.com> Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: <> Co-authored-by: root <root@prenyx0169.a51.clusters.nvidia.com> Co-authored-by: root <root@prenyx0042.a51.clusters.nvidia.com>
Wei Zhao committed
Mar 20, 2026 at 03:09 UTC
0140eafb1546c24d8486761abdbaa538d948bf42
2 files changed
+128
-125
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+38
-39
@@ -86,8 +86,6 @@ if flashinfer_comm is not None:
86
destroy_fi_ar_workspace,
87
get_fi_ar_quant_workspace,
88
get_fi_ar_workspace,
89
- initialize_fi_ar_quant_workspace,
90
- initialize_fi_ar_workspace,
89
)
90
91
ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern
@@ -133,15 +131,23 @@ if flashinfer_comm is not None:
131
132
# Select workspace based on pattern: quant patterns use the
133
# trtllm quant workspace, non-quant patterns use the primary workspace.
136
- if pattern_code in (
134
+ is_quant_pattern = pattern_code in (
135
ar_fusion_patterns.kARResidualRMSNormFP8Quant,
136
ar_fusion_patterns.kARResidualRMSNormFP4Quant,
139
- ):
140
- workspace = get_fi_ar_quant_workspace()
141
- else:
142
- workspace = get_fi_ar_workspace()
137
+ )
138
+ get_workspace_fn = (
139
+ get_fi_ar_quant_workspace if is_quant_pattern else get_fi_ar_workspace
140
+ )
141
+ workspace = get_workspace_fn(
142
+ world_size=world_size,
143
+ rank=get_tensor_model_parallel_rank(),
144
+ max_token_num=max_token_num,
145
+ hidden_dim=hidden_size,
146
+ dtype=allreduce_in.dtype,
147
+ group=get_tp_group().device_group,
148
+ )
149
assert workspace is not None, (
144
- "Flashinfer workspace must be initialized when using flashinfer"
150
+ "Flashinfer allreduce workspace must be initialized when using flashinfer"
151
)
152
assert flashinfer_comm is not None
153
if norm_out is None:
@@ -753,35 +759,29 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
759
scope="global",
760
)
761
756
- for workspace_init_fn in [
757
- initialize_fi_ar_workspace,
758
- initialize_fi_ar_quant_workspace,
759
- ]:
760
- try:
761
- workspace_init_fn(
762
- world_size=self.tp_size,
763
- rank=rank,
764
- max_token_num=self.max_token_num,
765
- hidden_dim=self.hidden_dim,
766
- dtype=self.model_dtype,
767
- group=self.group,
768
- )
769
- except Exception as e:
770
- if "multicast" in str(e).lower():
771
- logger.warning(
772
- "AllReduce fusion pass is disabled: flashinfer workspace "
773
- "creation failed: %s. This is expected on GPUs without "
774
- "NVSwitch (e.g., NVLink bridge-only or PCIe topologies). "
775
- "Falling back to non-fused allreduce.",
776
- str(e),
777
- )
778
- else:
779
- logger.warning(
780
- "Failed to initialize FlashInfer All Reduce workspace: %s. "
781
- "AllReduce fusion pass will be disabled.",
782
- e,
783
- )
784
- return
762
+ workspace_kwargs = dict(
763
+ world_size=self.tp_size,
764
+ rank=rank,
765
+ max_token_num=self.max_token_num,
766
+ hidden_dim=self.hidden_dim,
767
+ dtype=self.model_dtype,
768
+ group=self.group,
769
+ )
770
+ if get_fi_ar_workspace(**workspace_kwargs) is None:
771
+ logger.warning_once(
772
+ "Failed to initialize Flashinfer allreduce workspace. "
773
+ "Flashinfer allreduce-norm fusion will be disabled."
774
+ )
775
+ return
776
+
777
+ self.supports_quant_fusion = (
778
+ get_fi_ar_quant_workspace(**workspace_kwargs) is not None
779
+ )
780
+ if not self.supports_quant_fusion:
781
+ logger.warning_once(
782
+ "Failed to initialize Flashinfer allreduce workspace. "
783
+ "Flashinfer allreduce-norm-quant fusion will be disabled."
784
+ )
785
786
self.allreduce_params = FlashInferFusedAllReduceParams(
787
world_size=self.tp_size,
@@ -793,9 +793,8 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
793
794
@enable_fake_mode
795
def register_patterns(self) -> None:
796
- supports_quantization = get_fi_ar_quant_workspace() is not None
796
for epsilon in [1e-5, 1e-6]:
798
- if supports_quantization:
797
+ if self.supports_quant_fusion:
798
AllReduceFusedRMSNormStaticQuantFP8Pattern(
799
epsilon,
800
self.model_dtype,
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+90
-86
@@ -29,50 +29,27 @@ try:
29
except ImportError:
30
pass
31
32
-# Global workspace for standalone allreduce and non-quant ar+rms fusion
32
+# Workspace for standalone allreduce and non-quant ar+rms fusion
33
_fi_ar_workspace = None
34
# Extra workspace for quant fusion patterns (only supported by trtllm backend)
35
-# Only created if primary workspace is not already trtllm
35
_fi_ar_quant_workspace = None
36
37
39
-def get_fi_ar_workspace():
40
- return _fi_ar_workspace
41
-
42
-
43
-def get_fi_ar_quant_workspace():
44
- return _fi_ar_quant_workspace
45
-
46
-
47
-def initialize_fi_ar_workspace(
38
+def _create_workspace(
39
+ backend: str,
40
world_size: int,
41
rank: int,
42
max_token_num: int,
43
hidden_dim: int,
44
dtype: torch.dtype,
45
group: ProcessGroup,
54
-) -> None:
55
- """
56
- Initialize the workspace if not already initialized.
57
-
58
- Currently, this function is called by either the AllReduceFusionPass
59
- or the FlashInferAllReduce backend for standalone allreduce.
60
- If the fusion pass is enabled via
61
- --compilation-config.pass_config.fuse_allreduce_rms=true,
62
- it will create the workspace first, and the standalone backend
63
- will reuse the workspace. Otherwise, the standalone backend will
64
- create the workspace.
65
- """
66
- global _fi_ar_workspace
67
- if _fi_ar_workspace is not None:
68
- return
69
-
70
- backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
46
+):
47
+ """Create a flashinfer allreduce workspace, returning None on failure."""
48
comm_backend = TorchDistBackend(group=group)
49
rng_state = random.getstate()
50
try:
51
random.seed(int.from_bytes(os.urandom(16), byteorder="big"))
75
- _fi_ar_workspace = flashinfer_comm.create_allreduce_fusion_workspace(
52
+ workspace = flashinfer_comm.create_allreduce_fusion_workspace(
53
backend=backend,
54
world_size=world_size,
55
rank=rank,
@@ -81,9 +58,22 @@ def initialize_fi_ar_workspace(
58
dtype=dtype,
59
comm_backend=comm_backend,
60
)
61
+ except Exception as e:
62
+ if "multicast" in str(e).lower():
63
+ logger.warning_once(
64
+ "Failed to initialize FlashInfer All Reduce workspace: %s. "
65
+ "This is expected on GPUs without NVSwitch (e.g., NVLink "
66
+ "bridge-only or PCIe topologies).",
67
+ e,
68
+ )
69
+ else:
70
+ logger.warning_once(
71
+ "Failed to initialize FlashInfer All Reduce workspace: %s.",
72
+ e,
73
+ )
74
+ return None
75
finally:
76
random.setstate(rng_state)
86
- assert _fi_ar_workspace is not None
77
logger.debug(
78
"Initialized FlashInfer All Reduce workspace: backend=%s, "
79
"world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",
@@ -94,70 +84,84 @@ def initialize_fi_ar_workspace(
84
hidden_dim,
85
dtype,
86
)
87
+ return workspace
88
+
89
+
90
+def get_fi_ar_workspace(
91
+ world_size: int,
92
+ rank: int,
93
+ max_token_num: int,
94
+ hidden_dim: int,
95
+ dtype: torch.dtype,
96
+ group: ProcessGroup,
97
+):
98
+ """
99
+ Return the allreduce workspace for non-quant patterns, initializing if needed.
100
+
101
+ Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce
102
+ for standalone allreduce. Backend is controlled by
103
+ VLLM_FLASHINFER_ALLREDUCE_BACKEND env var.
104
+ """
105
+ global _fi_ar_workspace
106
+ if _fi_ar_workspace is not None:
107
+ return _fi_ar_workspace
108
+
109
+ backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
110
+
111
+ # Reuse the quant workspace if it was already created with the same backend
112
+ if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == backend:
113
+ _fi_ar_workspace = _fi_ar_quant_workspace
114
+ return _fi_ar_workspace
115
+
116
+ _fi_ar_workspace = _create_workspace(
117
+ backend, world_size, rank, max_token_num, hidden_dim, dtype, group
118
+ )
119
+ return _fi_ar_workspace
120
121
99
-def initialize_fi_ar_quant_workspace(
122
+def get_fi_ar_quant_workspace(
123
world_size: int,
124
rank: int,
125
max_token_num: int,
126
hidden_dim: int,
127
dtype: torch.dtype,
128
group: ProcessGroup,
106
-) -> None:
129
+):
130
"""
108
- Initialize the workspace used by quantization fusion patterns.
131
+ Return the allreduce workspace for quant patterns, initializing if needed.
132
110
- Currently this always creates a workspace for trtllm backend as only it
111
- supports quantization fusion (FP8/FP4). If the primary workspace
112
- is already trtllm, the quant workspace aliases to it.
133
+ Always uses trtllm backend as it is the only one supporting quantization
134
+ fusion (FP8/FP4).
135
"""
136
global _fi_ar_quant_workspace
137
if _fi_ar_quant_workspace is not None:
116
- return
138
+ return _fi_ar_quant_workspace
139
118
- # If primary workspace is already trtllm, reuse it
140
+ # Reuse the non-quant workspace if it was already created with trtllm
141
if _fi_ar_workspace is not None and _fi_ar_workspace.backend == "trtllm":
142
_fi_ar_quant_workspace = _fi_ar_workspace
121
- return
143
+ return _fi_ar_quant_workspace
144
123
- comm_backend = TorchDistBackend(group=group)
124
- _fi_ar_quant_workspace = flashinfer_comm.create_allreduce_fusion_workspace(
125
- backend="trtllm",
126
- world_size=world_size,
127
- rank=rank,
128
- max_token_num=max_token_num,
129
- hidden_dim=hidden_dim,
130
- dtype=dtype,
131
- comm_backend=comm_backend,
132
- )
133
- assert _fi_ar_quant_workspace is not None
134
- logger.debug(
135
- "Initialized FlashInfer All Reduce workspace: backend=trtllm, "
136
- "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",
137
- world_size,
138
- rank,
139
- max_token_num,
140
- hidden_dim,
141
- dtype,
145
+ _fi_ar_quant_workspace = _create_workspace(
146
+ "trtllm", world_size, rank, max_token_num, hidden_dim, dtype, group
147
)
148
+ return _fi_ar_quant_workspace
149
150
151
_fi_ar_workspace_lock = threading.Lock()
152
153
154
def destroy_fi_ar_workspace():
149
- global _fi_ar_workspace
150
- global _fi_ar_quant_workspace
155
+ global _fi_ar_workspace, _fi_ar_quant_workspace
156
with _fi_ar_workspace_lock:
152
- if (
153
- _fi_ar_quant_workspace is not None
154
- and _fi_ar_quant_workspace is not _fi_ar_workspace
155
- ):
156
- _fi_ar_quant_workspace.destroy()
157
- _fi_ar_quant_workspace = None
157
+ is_alias = _fi_ar_workspace is _fi_ar_quant_workspace
158
+
159
if _fi_ar_workspace is not None:
160
_fi_ar_workspace.destroy()
160
- _fi_ar_workspace = None
161
+ if _fi_ar_quant_workspace is not None and not is_alias:
162
+ _fi_ar_quant_workspace.destroy()
163
+
164
+ _fi_ar_workspace = _fi_ar_quant_workspace = None
165
166
167
atexit.register(destroy_fi_ar_workspace)
@@ -209,29 +213,21 @@ class FlashInferAllReduce:
213
214
def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
215
"""Ensure the all reduce workspace is initialized."""
212
- if get_fi_ar_workspace() is not None:
213
- return True
216
if self.max_num_tokens == 0:
217
element_size = torch.tensor([], dtype=dtype, device="cpu").element_size()
218
self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
217
- try:
218
- initialize_fi_ar_workspace(
219
- world_size=self.world_size,
220
- rank=self.rank,
221
- max_token_num=self.max_num_tokens,
222
- hidden_dim=hidden_dim,
223
- dtype=dtype,
224
- group=self.group,
225
- )
226
- return True
227
- except Exception as e:
228
- logger.warning(
229
- "Failed to initialize FlashInfer All Reduce workspace: %s. "
230
- "FlashInfer All Reduce will be disabled.",
231
- e,
232
- )
219
+ workspace = get_fi_ar_workspace(
220
+ world_size=self.world_size,
221
+ rank=self.rank,
222
+ max_token_num=self.max_num_tokens,
223
+ hidden_dim=hidden_dim,
224
+ dtype=dtype,
225
+ group=self.group,
226
+ )
227
+ if workspace is None:
228
self.disabled = True
229
return False
230
+ return True
231
232
def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool:
233
if self.disabled:
@@ -257,7 +253,15 @@ class FlashInferAllReduce:
253
return self._ensure_workspace(hidden_dim, input_tensor.dtype)
254
255
def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
260
- workspace = get_fi_ar_workspace()
256
+ _, hidden_dim = input_tensor.shape
257
+ workspace = get_fi_ar_workspace(
258
+ world_size=self.world_size,
259
+ rank=self.rank,
260
+ max_token_num=self.max_num_tokens,
261
+ hidden_dim=hidden_dim,
262
+ dtype=input_tensor.dtype,
263
+ group=self.group,
264
+ )
265
return flashinfer_comm.allreduce_fusion(
266
input=input_tensor,
267
workspace=workspace,