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 15, 01:00 UTC
Fix Flashinfer Allreduce+Norm enable disable calculation based on `fi_allreduce_fusion_max_token_num` (#21325)
Signed-off-by: XIn Li <xinli@nvidia.com>
Xin Li committed
Jul 22, 2025 at 15:42 UTC
ae268b6326033b9d883fc04078f3b623ca6a8f78
1 file changed
+13
-6
vllm/compilation/collective_fusion.py
+13
-6
@@ -159,6 +159,9 @@ if flashinfer_comm is not None:
159
6: MiB // 2, # 512KB
160
8: MiB // 2, # 512KB
161
}
162
+ # opt for a more conservative default value
163
+ # when world size is not in _FI_MAX_SIZES
164
+ _DEFAULT_FI_MAX_SIZE = MiB // 2
165
166
def call_trtllm_fused_allreduce_norm(
167
allreduce_in: torch.Tensor,
@@ -173,12 +176,16 @@ if flashinfer_comm is not None:
176
max_token_num: int,
177
norm_out: Optional[torch.Tensor] = None,
178
) -> None:
176
- use_flashinfer = allreduce_in.shape[0] * allreduce_in.shape[
177
- 1] * allreduce_in.element_size() <= min(
178
- _FI_MAX_SIZES[world_size],
179
- max_token_num * allreduce_in.shape[0] *
180
- allreduce_in.element_size(),
181
- )
179
+
180
+ num_tokens, hidden_size = allreduce_in.shape
181
+ element_size = allreduce_in.element_size()
182
+ current_tensor_size = num_tokens * hidden_size * element_size
183
+ max_fusion_size = max_token_num * hidden_size * element_size
184
+ use_flashinfer = current_tensor_size <= min(
185
+ _FI_MAX_SIZES.get(world_size, _DEFAULT_FI_MAX_SIZE),
186
+ max_fusion_size,
187
+ )
188
+
189
if use_flashinfer:
190
assert (_FI_WORKSPACE_TENSOR is not None
191
), "Flashinfer must be enabled when using flashinfer"