Disable fp4 test. Cleanup fusion. Move allreduce out of fused_moe custom op

Signed-off-by: ilmarkov <imarkov@redhat.com>

Author: ilmarkov

tests/compile/test_fusion_all_reduce.py

@@ -136,12 +136,14 @@ def ops_in_model_before(self):
 
 
 @multi_gpu_test(num_gpus=2)
-@pytest.mark.parametrize("test_model", [
-    TestAllReduceRMSNormModel,
-    TestAllReduceFusedAddRMSNormModel,
-    TestAllReduceFusedAddRMSNormStaticQuantFP8Model,
-    TestAllReduceFusedAddRMSNormStaticQuantFP4Model,
-])
+@pytest.mark.parametrize(
+    "test_model",
+    [
+        TestAllReduceRMSNormModel,
+        TestAllReduceFusedAddRMSNormModel,
+        TestAllReduceFusedAddRMSNormStaticQuantFP8Model,
+        # TestAllReduceFusedAddRMSNormStaticQuantFP4Model,
+    ])
 @pytest.mark.parametrize("batch_size", [8])
 @pytest.mark.parametrize("seq_len", [8])
 @pytest.mark.parametrize("hidden_size", [16])

vllm/compilation/collective_fusion.py

@@ -417,7 +417,6 @@ def call_trtllm_fused_allreduce_norm(
         fp32_acc: bool,
         max_token_num: int,
         pattern_code: int,
-        fuse_rms_quant: bool,
         norm_out: Optional[torch.Tensor] = None,
         quant_out: Optional[torch.Tensor] = None,
         scale_out: Optional[torch.Tensor] = None,
@@ -489,13 +488,8 @@ def call_trtllm_fused_allreduce_norm(
                     torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma,
                                           rms_eps)
                 if scale_factor is not None:
-                    assert scale_out is not None
                     torch.ops._C.scaled_fp4_quant(quant_out, norm_out,
                                                   scale_out, scale_factor)
-                    # if scale_out is not None:
-                    # else:
-                    #     torch.ops._C.static_scaled_fp8_quant(
-                    #         quant_out, norm_out, scale_factor)
             if scale_factor is None or norm_out is not None:
                 # we need to return allreduce outpput
                 # in cases of non quant fused AR + RMS norm
@@ -514,7 +508,6 @@ def call_trtllm_fused_allreduce_norm_fake(
             fp32_acc: bool,
             max_token_num: int,
             pattern_code: int,
-            fuse_rms_quant: bool,
             norm_out: Optional[torch.Tensor] = None,
             quant_out: Optional[torch.Tensor] = None,
             scale_out: Optional[torch.Tensor] = None,
@@ -547,17 +540,14 @@ def __init__(
         world_size: int,
         use_fp32_lamport: bool = False,
         max_token_num: int = 1024,
-        fuse_rms_quant: bool = False,
     ):
         self.rank = rank
         self.world_size = world_size
         self.use_fp32_lamport = use_fp32_lamport
         self.trigger_completion_at_end = True
         self.launch_with_pdl = True
         self.fp32_acc = True
-        self.use_oneshot = False
         self.max_token_num = max_token_num
-        self.fuse_rms_quant = fuse_rms_quant
 
     def get_trtllm_fused_allreduce_kwargs(self):
         return {
@@ -567,7 +557,6 @@ def get_trtllm_fused_allreduce_kwargs(self):
             "trigger_completion_at_end": self.trigger_completion_at_end,
             "fp32_acc": self.fp32_acc,
             "max_token_num": self.max_token_num,
-            "fuse_rms_quant": self.fuse_rms_quant,
         }
 
 
@@ -1103,10 +1092,7 @@ def __init__(self, config: VllmConfig):
             world_size=self.tp_size,
             use_fp32_lamport=use_fp32_lamport,
             max_token_num=max_num_token,
-            # fuse rms norm static fp8 quant fused op
-            # in fallback path, when we don't use flashinfer
-            fuse_rms_quant=config.compilation_config.pass_config.enable_fusion)
-
+        )
         for epsilon in [1e-5, 1e-6]:
             AllReduceFusedRMSNormStaticQuantFP8Pattern(
                 epsilon,

vllm/model_executor/layers/fused_moe/layer.py

@@ -1411,10 +1411,16 @@ def forward(self, hidden_states: torch.Tensor,
         # TODO: Once the OOM issue for the TPU backend is resolved, we will
         # switch to using the moe_forward custom op.
         if current_platform.is_tpu():
-            return self.forward_impl(hidden_states, router_logits)
+            final_hidden_states = self.forward_impl(hidden_states,
+                                                    router_logits)
         else:
-            return torch.ops.vllm.moe_forward(hidden_states, router_logits,
-                                              self.layer_name)
+            final_hidden_states = torch.ops.vllm.moe_forward(
+                hidden_states, router_logits, self.layer_name)
+        if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
+            # Default set to False. (May have to add shared expert outputs.
+            final_hidden_states = self.maybe_all_reduce_tensor_model_parallel(
+                final_hidden_states)
+        return final_hidden_states
 
     def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
                              full_router_logits: torch.Tensor):
@@ -1538,10 +1544,6 @@ def forward_impl(self, hidden_states: torch.Tensor,
 
         if do_naive_dispatch_combine:
             final_hidden_states = get_ep_group().combine(final_hidden_states)
-        if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
-            # Default set to False. (May have to add shared expert outputs.
-            final_hidden_states = self.maybe_all_reduce_tensor_model_parallel(
-                final_hidden_states)
 
         return final_hidden_states