fix router_scaling_scale

Author: none

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_ep.py

@@ -95,7 +95,7 @@ def __init__(
         self.n_group = network_config["n_group"]
         network_config["topk_group"] = network_config.get("topk_group", 0)
         self.topk_group = network_config["topk_group"]
-        network_config["routed_scaling_factor"] = network_config.get("routed_scaling_factor", 0)
+        network_config["routed_scaling_factor"] = network_config.get("routed_scaling_factor", 1.0)
         self.routed_scaling_factor = network_config["routed_scaling_factor"]
 
         self.lock = threading.Lock()
@@ -126,6 +126,7 @@ def experts(
             num_expert_group=num_expert_group,
             scoring_func=self.scoring_func,
         )
+        topk_weights.mul_(self.routed_scaling_factor)
 
         if self.redundancy_expert_num > 0:
             redundancy_topk_ids_repair(
@@ -173,6 +174,7 @@ def low_latency_dispatch(
             num_expert_group=self.n_group,
             scoring_func=self.scoring_func,
         )
+        topk_weights.mul_(self.routed_scaling_factor)
 
         if self.redundancy_expert_num > 0:
             redundancy_topk_ids_repair(
@@ -213,6 +215,7 @@ def select_experts_and_quant_input(
             num_expert_group=self.n_group,
             scoring_func=self.scoring_func,
         )
+        topk_weights.mul_(self.routed_scaling_factor)
         if self.redundancy_expert_num > 0:
             redundancy_topk_ids_repair(
                 topk_ids=topk_idx,

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_tp.py

@@ -68,6 +68,7 @@ def experts(self, input_tensor, router_logits, top_k, renormalize, use_grouped_t
             num_expert_group=num_expert_group,
             scoring_func=self.scoring_func,
         )
+        topk_weights.mul_(self.routed_scaling_factor)
         if self.num_fused_shared_experts > 0:
             pad_topk_ids = torch.arange(
                          start=self.n_routed_experts - self.num_fused_shared_experts, 
@@ -76,7 +77,7 @@ def experts(self, input_tensor, router_logits, top_k, renormalize, use_grouped_t
                          dtype=topk_ids.dtype,
                          device="cuda").view(1, self.num_fused_shared_experts).repeat(topk_ids.shape[0], 1)
             pad_topk_weights = torch.full((topk_weights.shape[0], self.num_fused_shared_experts),
-                                          fill_value=1.0 / self.routed_scaling_factor,
+                                          fill_value=1.0,
                                           device="cuda",
                                           dtype=topk_weights.dtype)
 

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -56,7 +56,6 @@ def __init__(self, layer_num, network_config, mode=[]):
         self.norm_topk_prob = network_config["norm_topk_prob"]
         self.n_group = network_config["n_group"]
         self.topk_group = network_config["topk_group"]
-        self.routed_scaling_factor = network_config["routed_scaling_factor"]
 
         self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5)
         if network_config.get("rope_scaling", None) is not None:
@@ -680,8 +679,6 @@ def _moe_ffn(
             num_expert_group=self.n_group,
         )
 
-        hidden_states.mul_(self.routed_scaling_factor)
-
         if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0:
             hidden_states.add_(shared_output)
 
@@ -707,7 +704,6 @@ def _moe_ffn_edp(
             num_expert_group=self.n_group,
             is_prefill=infer_state.is_prefill,
         )
-        ep_output.mul_(self.routed_scaling_factor)
 
         if self.n_shared_experts is not None:
             ep_output.add_(shared_output)
@@ -819,7 +815,6 @@ def overlap_tpsp_token_forward(
         # 0 hook
         if getattr(infer_state, "hook", None) is not None:
             infer_state.hook()
-            _0_ffn_out *= self.routed_scaling_factor
             if self.n_shared_experts is not None:
                 _0_ffn_out.add_(_0_shared_output)
             input_embdings.add_(_0_ffn_out.view(-1, self.embed_dim_))
@@ -833,7 +828,6 @@ def overlap_tpsp_token_forward(
         def _1_hook_post():
             _1_hook()
             nonlocal _1_ffn_out
-            _1_ffn_out *= self.routed_scaling_factor
             if self.n_shared_experts is not None:
                 _1_ffn_out.add_(_1_shared_output)
             input_embdings1.add_(_1_ffn_out.view(-1, self.embed_dim_))
@@ -965,7 +959,6 @@ def overlap_tpsp_context_forward(
 
         _1_combine_event = Buffer.capture()
 
-        _0_ffn_out *= self.routed_scaling_factor
         if self.n_shared_experts is not None:
             _0_ffn_out.add_(_0_shared_output)
         input_embdings.add_(_0_ffn_out.view(-1, self.embed_dim_))
@@ -976,7 +969,6 @@ def overlap_tpsp_context_forward(
         def _1_hook_post():
             _1_hook()
             nonlocal _1_ffn_out
-            _1_ffn_out *= self.routed_scaling_factor
             if self.n_shared_experts is not None:
                 _1_ffn_out.add_(_1_shared_output)
             input_embdings1.add_(_1_ffn_out.view(-1, self.embed_dim_))