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, 02:15 UTC
[Misc] Refactor AllReduceFusionPass. Remove parameter (#20918)
Signed-off-by: ilmarkov <imarkov@redhat.com> Co-authored-by: ilmarkov <imarkov@redhat.com>
Ilya Markov committed
Jul 15, 2025 at 08:57 UTC
37a7d5d74a9eddae3265bb1118efbb0f5ce10a93
3 files changed
+7
-10
tests/compile/test_fusion_all_reduce.py
+1
-3
@@ -132,9 +132,7 @@ def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
132
dtype=dtype,
133
seed=42)
134
135
- all_reduce_fusion_pass = AllReduceFusionPass(
136
- vllm_config, vllm_config.compilation_config.pass_config.
137
- fi_allreduce_fusion_max_token_num)
135
+ all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)
136
backend = TestBackend(all_reduce_fusion_pass)
137
138
model = test_model_cls(hidden_size)
vllm/compilation/collective_fusion.py
+5
-3
@@ -397,7 +397,7 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
397
398
class AllReduceFusionPass(VllmInductorPass):
399
400
- def __init__(self, config: VllmConfig, max_token_num: int):
400
+ def __init__(self, config: VllmConfig):
401
super().__init__(config)
402
self.disabled = True
403
self.tp_size = get_tensor_model_parallel_world_size()
@@ -429,7 +429,8 @@ class AllReduceFusionPass(VllmInductorPass):
429
flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
430
tp_rank=rank,
431
tp_size=self.tp_size,
432
- max_token_num=max_token_num,
432
+ max_token_num=config.compilation_config.pass_config.
433
+ fi_allreduce_fusion_max_token_num,
434
hidden_dim=self.hidden_dim,
435
group=self.group,
436
use_fp32_lamport=use_fp32_lamport,
@@ -441,7 +442,8 @@ class AllReduceFusionPass(VllmInductorPass):
442
rank=rank,
443
world_size=self.tp_size,
444
use_fp32_lamport=use_fp32_lamport,
444
- max_token_num=max_token_num,
445
+ max_token_num=config.compilation_config.pass_config.
446
+ fi_allreduce_fusion_max_token_num,
447
)
448
449
for epsilon in [1e-5, 1e-6]:
vllm/compilation/pass_manager.py
+1
-4
@@ -63,10 +63,7 @@ class PostGradPassManager(CustomGraphPass):
63
if self.pass_config.enable_attn_fusion:
64
self.passes += [AttnFusionPass(config)]
65
if self.pass_config.enable_fi_allreduce_fusion:
66
- self.passes += [
67
- AllReduceFusionPass(
68
- config, self.pass_config.fi_allreduce_fusion_max_token_num)
69
- ]
66
+ self.passes += [AllReduceFusionPass(config)]
67
self.fix_functionalization = FixFunctionalizationPass(config)
68
69
def add(self, pass_: InductorPass):