update embedding replace pattern

Signed-off-by: cascade812 <[email protected]>

Author: cascade812

vllm/compilation/collective_fusion.py

@@ -65,20 +65,20 @@ def replace_with_embedding_reduce_scatter_rmsnorm(
     reduce_scatter = torch.ops.vllm.reduce_scatter.default(
         where, dim=0, world_size=tp_size, group_name=tp.unique_name)
 
-    # rmsnorm_result = torch.empty_like(reduce_scatter)
+    rmsnorm_result = torch.empty_like(reduce_scatter)
     rmsnorm = torch.ops.higher_order.auto_functionalized(
         torch.ops._C.rms_norm.default,
-        result=permute,
+        result=rmsnorm_result,
         input=reduce_scatter,
         weight=arg3_1,
         epsilon=1e-5)
 
-    all_gather = torch.ops.vllm.all_gather.default(reduce_scatter,
+    all_gather = torch.ops.vllm.all_gather.default(rmsnorm[1],
                                                    dim=0,
                                                    world_size=tp_size,
                                                    group_name=tp.unique_name)
 
-    return rmsnorm[1], all_gather
+    return all_gather, reduce_scatter
 
 
 def search_gemm_allreduce_rmsnorm(