fix qwen3.5 mtp accept rate

Author: sufubao

lightllm/models/qwen3_5_mtp/layer_infer/pre_layer_infer.py

@@ -23,6 +23,7 @@ def _mtp_fuse(
             input_embdings.shape[0] == tgt_embdings.shape[0]
         ), f"shape {input_embdings.shape} != shape {tgt_embdings.shape}"
 
+        layer_weight.main_norm_weight_(input=tgt_embdings, eps=self.eps_, out=tgt_embdings)
         layer_weight.enorm_weight_(input=input_embdings, eps=self.eps_, out=input_embdings)
         layer_weight.hnorm_weight_(input=tgt_embdings, eps=self.eps_, out=tgt_embdings)
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)

lightllm/models/qwen3_5_mtp/layer_weights/pre_and_post_layer_weight.py

@@ -42,4 +42,5 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg):
         # Shared with the main Qwen3.5 model, injected by the model class (not loaded here).
         self.wte_weight_: EmbeddingWeight = None
         self.lm_head_weight_: LMHeadWeight = None
+        self.main_norm_weight_: NoTpGEMMANormWeight = None
         return

lightllm/models/qwen3_5_mtp/model.py

@@ -69,6 +69,7 @@ def _init_weights(self, start_layer_index=None):
         # Shared with the main Qwen3.5 model (mtp_use_dedicated_embeddings: false).
         self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_
         self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_
+        self.pre_post_weight.main_norm_weight_ = self.main_model.pre_post_weight.final_norm_weight_
         return
 
     def _init_infer_layer(self, start_layer_index=None):