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 14, 05:30 UTC
[Misc] Add override for allreduce fusion thresholds (#23639)
Signed-off-by: Julien Lin <jullin@nvidia.com>
nvjullin committed
Aug 26, 2025 at 23:53 UTC
7ea22e42d5f666a26b3ce4117724dadfdb4d3887
2 files changed
+24
vllm/compilation/collective_fusion.py
+13
@@ -10,6 +10,7 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
10
from torch._inductor.pattern_matcher import PatternMatcherPass
11
from torch.distributed._symmetric_memory import enable_symm_mem_for_group
12
13
+import vllm.envs as envs
14
from vllm.config import VllmConfig
15
from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
16
from vllm.distributed.parallel_state import (
@@ -401,6 +402,18 @@ if flashinfer_comm is not None:
402
6: MiB // 2, # 512KB
403
8: MiB // 2, # 512KB
404
}
405
+
406
+ try:
407
+ _FI_MAX_SIZES.update({
408
+ int(k): int(float(v) * MiB)
409
+ for k, v in
410
+ envs.VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB.items()
411
+ })
412
+ except Exception as e:
413
+ raise ValueError(
414
+ "Failed to parse VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB: "
415
+ + str(e)) from e
416
+
417
# opt for a more conservative default value
418
# when world size is not in _FI_MAX_SIZES
419
_DEFAULT_FI_MAX_SIZE = MiB // 2
vllm/envs.py
+11
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import hashlib
5
+import json
6
import os
7
import sys
8
import tempfile
@@ -1046,6 +1047,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
1047
"VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE":
1048
lambda: int(os.getenv("VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE", "163840")),
1049
1050
+ # Specifies the thresholds of the communicated tensor sizes under which
1051
+ # vllm should use flashinfer fused allreduce. The variable should be a
1052
+ # JSON with the following format:
1053
+ # { <world size>: <max size in mb> }
1054
+ # Unspecified world sizes will fallback to
1055
+ # { 2: 64, 4: 1, <everything else>: 0.5 }
1056
+ "VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB":
1057
+ lambda: json.loads(os.getenv(
1058
+ "VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB", "{}")),
1059
+
1060
# MoE routing strategy selector.
1061
# See `RoutingSimulator.get_available_strategies()` # for available
1062
# strategies.