fix: correct final hidden normalization for MTP

lightllm/common/basemodel/basemodel.py

@@ -35,6 +35,7 @@
     get_added_mtp_kv_layer_num,
 )
 from lightllm.distributed.communication_op import dist_group_manager
+from lightllm.common.mtp_config import FinalHiddenPolicy
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput
 from lightllm.common.basemodel.hidden_collector import (
     NoopHiddenCollector,
@@ -349,6 +350,7 @@ def _select_mem_indexes(self, model_input: ModelInput):
     def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0):
         infer_state = self.infer_state_class()
         infer_state.hidden_collector = self.hidden_collector_prototype.new_instance()
+        infer_state.final_hidden_policy = self.mtp_manager.get_final_hidden_policy(self)
         infer_state.input_ids = model_input.input_ids
         infer_state.is_prefill = model_input.is_prefill
         infer_state.is_draft_model = self.is_mtp_draft_model
@@ -630,7 +632,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"
@@ -641,6 +642,7 @@ def _context_forward(self, infer_state: InferStateInfo):
         input_tensors = [input_embs]
         if Autotuner.is_autotune_warmup():
             infer_state.hidden_collector = NoopHiddenCollector()
+            infer_state.final_hidden_policy = FinalHiddenPolicy.NONE
 
         def prefill_func(input_tensors, _infer_state):
             hidden_collector = _infer_state.hidden_collector
@@ -686,7 +688,7 @@ def prefill_func(input_tensors, _infer_state):
 
         post_output = self.post_infer.token_forward(last_input_embs, infer_state, self.pre_post_weight)
         hidden_collector = infer_state.hidden_collector
-        hidden_collector.add_final_hidden(last_input_embs)
+        hidden_collector.add_final_hidden(post_output.final_hidden)
         model_output = self._create_model_output(post_output, infer_state)
         del post_output
 
@@ -712,7 +714,7 @@ def _token_forward(self, infer_state: InferStateInfo):
             last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight
         )
 
-        hidden_collector.add_final_hidden(last_input_embs)
+        hidden_collector.add_final_hidden(post_output.final_hidden)
         model_output = self._create_model_output(post_output, infer_state)
         del post_output
 
@@ -915,8 +917,8 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
         )
         g_cache_manager.cache_env_out()
 
-        hidden_collector0.add_final_hidden(last_input_embs)
-        hidden_collector1.add_final_hidden(last_input_embs1)
+        hidden_collector0.add_final_hidden(post_output.final_hidden)
+        hidden_collector1.add_final_hidden(post_output1.final_hidden)
         model_output = self._create_model_output(post_output, infer_state)
         model_output1 = self._create_model_output(post_output1, infer_state1)
         del post_output, post_output1
@@ -957,8 +959,8 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1:
             last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
 
-        hidden_collector0.add_final_hidden(last_input_embs)
-        hidden_collector1.add_final_hidden(last_input_embs1)
+        hidden_collector0.add_final_hidden(post_output.final_hidden)
+        hidden_collector1.add_final_hidden(post_output1.final_hidden)
         model_output = self._create_model_output(post_output, infer_state)
         model_output1 = self._create_model_output(post_output1, infer_state1)
         del post_output, post_output1

lightllm/common/basemodel/batch_objs.py

@@ -175,10 +175,11 @@ def unpad_prefill(self, origin_handle_token_num: int) -> "ModelMtpOutputCollecto
 
 @dataclass
 class PostLayerOutput:
-    """输出层 logits，以及可选的候选列到 token ID 的映射。"""
+    """输出层 logits、可选的候选 token ID 映射和供 collector 收集的最终 hidden。"""
 
     logits: torch.Tensor
     logits_token_ids: Optional[torch.Tensor] = None
+    final_hidden: Optional[torch.Tensor] = None
 
 
 @dataclass

lightllm/common/basemodel/hidden_collector.py

@@ -75,7 +75,7 @@ def add(self, layer_index: int, hidden: torch.Tensor) -> None:
         """
         return
 
-    def add_final_hidden(self, final_hidden: torch.Tensor) -> None:
+    def add_final_hidden(self, final_hidden: Optional[torch.Tensor]) -> None:
         """接收完成模型输出侧 gather 后的最终层 hidden tensor。
 
         BaseModel 在 logits 计算完成后、调用 :meth:`finish_output` 前调用该接口。基类
@@ -165,7 +165,8 @@ def __init__(self) -> None:
     def new_instance(self) -> HiddenCollector:
         return FinalHiddenCollector()
 
-    def add_final_hidden(self, final_hidden: torch.Tensor) -> None:
+    def add_final_hidden(self, final_hidden: Optional[torch.Tensor]) -> None:
+        assert final_hidden is not None, "post layer must return the requested final hidden"
         self.final_hidden = final_hidden
 
     def finish_output(self, infer_state) -> ModelMtpOutputCollector:

lightllm/common/basemodel/infer_struct.py

@@ -8,6 +8,7 @@
 from .triton_kernel.gen_prefill_params import gen_prefill_params
 from .triton_kernel.gen_decode_params import gen_decode_params
 from .triton_kernel.multimodal_emb import mark_multimodal_obj
+from lightllm.common.mtp_config import FinalHiddenPolicy
 from .batch_objs import ModelInput
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.dist_utils import get_global_dp_rank, get_dp_world_size
@@ -68,6 +69,7 @@ def __init__(self):
         # 当前 forward 独占的 hidden state 收集器。普通推理使用短生命周期实例；
         # Prefill CUDA Graph 使用随 graph infer state 长期保存的实例。
         self.hidden_collector = None
+        self.final_hidden_policy = FinalHiddenPolicy.NONE
 
         # 衍生使用的管理变量，为了方便扩展接入其他的高性能attention推理算子，在
         # inferstate 基类上添加下面的标记变量，用于扩展。

lightllm/common/basemodel/mtp_manager.py

@@ -1,5 +1,10 @@
+from functools import lru_cache
 from typing import ClassVar, Optional
 
+from transformers.configuration_utils import PretrainedConfig
+
+from lightllm.common.mtp_config import FinalHiddenPolicy, MtpPairConfig
+from lightllm.models.draft_registry import DraftModelRegistry
 from lightllm.common.basemodel.hidden_collector import (
     FinalHiddenCollector,
     HiddenCollector,
@@ -71,6 +76,18 @@ def get_decode_draft_step(self, is_draft_model: bool) -> int:
 
         return self.get_decode_tokens_per_request(is_draft_model) - 1
 
+    @lru_cache(maxsize=None)
+    def _get_pair_config(self, draft_model_dir: str) -> MtpPairConfig:
+        draft_config, _ = PretrainedConfig.get_config_dict(draft_model_dir)
+        return DraftModelRegistry.get_pair_config(draft_config, self.args.mtp_mode)
+
+    def get_final_hidden_policy(self, model) -> FinalHiddenPolicy:
+        if self.args.mtp_mode is None:
+            return FinalHiddenPolicy.NONE
+        draft_model_dir = model.weight_dir_ if model.is_mtp_draft_model else self.args.mtp_draft_model_dir[0]
+        config = self._get_pair_config(draft_model_dir)
+        return config.draft_final_hidden_policy if model.is_mtp_draft_model else config.main_final_hidden_policy
+
     def create_hidden_collector(
         self,
         model,

lightllm/common/mtp_config.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Type, Union
+
+
+class FinalHiddenPolicy(Enum):
+    NONE = auto()
+    PRE_NORM = auto()
+    POST_NORM = auto()
+
+
+@dataclass(frozen=True)
+class MtpPairConfig:
+    draft_model_class: Union[Type, str]
+    main_final_hidden_policy: FinalHiddenPolicy
+    draft_final_hidden_policy: FinalHiddenPolicy

lightllm/models/builtin.py

@@ -1,5 +1,7 @@
 """Built-in model registrations; keep implementation imports out of this file."""
 
+from lightllm.common.mtp_config import FinalHiddenPolicy
+
 from .registry import ModelRegistry, is_reward_model, llm_model_type_is
 from .draft_registry import DraftModelRegistry
 
@@ -151,42 +153,92 @@ def _tarsier_text_model_is(model_type):
     "deepseek_v3",
     ("vanilla_with_att", "eagle_with_att"),
     "lightllm.models.deepseek_mtp.model:Deepseek3MTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
 )
 DraftModelRegistry.register(
     "glm4_moe_lite",
     ("vanilla_with_att", "eagle_with_att"),
     "lightllm.models.glm4_moe_lite_mtp.model:Glm4MoeLiteMTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.PRE_NORM,
 )
 DraftModelRegistry.register(
     ("glm5_next", "glm5_next_text"),
     ("vanilla_with_att", "eagle_with_att"),
     "lightllm.models.glm5_next_mtp.model:Glm5NextMTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
 )
 DraftModelRegistry.register(
-    "mistral", ("vanilla_no_att", "eagle_no_att"), "lightllm.models.mistral_mtp.model:MistralMTPModel"
+    "mistral",
+    ("vanilla_no_att", "eagle_no_att"),
+    "lightllm.models.mistral_mtp.model:MistralMTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.PRE_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.PRE_NORM,
 )
 DraftModelRegistry.register(
-    ("qwen3_5", "qwen3_5_text"), "dflash", "lightllm.models.qwen3_5_dflash.model:Qwen3_5DFlashModel"
+    ("qwen3_5", "qwen3_5_text"),
+    "dflash",
+    "lightllm.models.qwen3_5_dflash.model:Qwen3_5DFlashModel",
+    main_final_hidden_policy=FinalHiddenPolicy.NONE,
+    draft_final_hidden_policy=FinalHiddenPolicy.NONE,
 )
 DraftModelRegistry.register(
-    ("qwen3_5", "qwen3_5_text"), "dspark", "lightllm.models.qwen3_5_dspark.model:Qwen3_5DSparkModel"
+    ("qwen3_5", "qwen3_5_text"),
+    "dspark",
+    "lightllm.models.qwen3_5_dspark.model:Qwen3_5DSparkModel",
+    main_final_hidden_policy=FinalHiddenPolicy.NONE,
+    draft_final_hidden_policy=FinalHiddenPolicy.NONE,
 )
 DraftModelRegistry.register(
     ("qwen3_5_moe", "qwen3_5_moe_text"),
     ("vanilla_with_att", "eagle_with_att"),
     "lightllm.models.qwen3_5_moe_mtp.model:Qwen3_5MoeMTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
 )
 DraftModelRegistry.register(
     ("qwen3_5", "qwen3_5_text"),
     ("vanilla_with_att", "eagle_with_att"),
     "lightllm.models.qwen3_5_mtp.model:Qwen3_5MTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.POST_NORM,
+)
+DraftModelRegistry.register(
+    "qwen3",
+    "dflash2",
+    "lightllm.models.qwen3_dflash2.model:Qwen3DFlash2Model",
+    main_final_hidden_policy=FinalHiddenPolicy.NONE,
+    draft_final_hidden_policy=FinalHiddenPolicy.NONE,
+)
+DraftModelRegistry.register(
+    "qwen3",
+    "dflash",
+    "lightllm.models.qwen3_dflash.model:Qwen3DFlashModel",
+    main_final_hidden_policy=FinalHiddenPolicy.NONE,
+    draft_final_hidden_policy=FinalHiddenPolicy.NONE,
+)
+DraftModelRegistry.register(
+    "qwen3",
+    "dspark",
+    "lightllm.models.qwen3_dspark.model:Qwen3DSparkModel",
+    main_final_hidden_policy=FinalHiddenPolicy.NONE,
+    draft_final_hidden_policy=FinalHiddenPolicy.NONE,
 )
-DraftModelRegistry.register("qwen3", "dflash", "lightllm.models.qwen3_dflash.model:Qwen3DFlashModel")
-DraftModelRegistry.register("qwen3", "dflash2", "lightllm.models.qwen3_dflash2.model:Qwen3DFlash2Model")
-DraftModelRegistry.register("qwen3", "dspark", "lightllm.models.qwen3_dspark.model:Qwen3DSparkModel")
-DraftModelRegistry.register("qwen3", "eagle3", "lightllm.models.qwen3_eagle.model:Qwen3EagleModel")
 DraftModelRegistry.register(
-    "qwen3_moe", ("vanilla_no_att", "eagle_no_att"), "lightllm.models.qwen3_moe_mtp.model:Qwen3MOEMTPModel"
+    "qwen3",
+    "eagle3",
+    "lightllm.models.qwen3_eagle.model:Qwen3EagleModel",
+    main_final_hidden_policy=FinalHiddenPolicy.NONE,
+    draft_final_hidden_policy=FinalHiddenPolicy.PRE_NORM,
+)
+DraftModelRegistry.register(
+    "qwen3_moe",
+    ("vanilla_no_att", "eagle_no_att"),
+    "lightllm.models.qwen3_moe_mtp.model:Qwen3MOEMTPModel",
+    main_final_hidden_policy=FinalHiddenPolicy.PRE_NORM,
+    draft_final_hidden_policy=FinalHiddenPolicy.PRE_NORM,
 )
 
 
@@ -196,4 +248,9 @@ def _tarsier_text_model_is(model_type):
     for configs in ModelRegistry._registry.values()
     for config in configs
 }
-MODEL_CLASS_PATHS.update({path.rsplit(":", 1)[1]: path for path in DraftModelRegistry._registry.values()})
+MODEL_CLASS_PATHS.update(
+    {
+        config.draft_model_class.rsplit(":", 1)[1]: config.draft_model_class
+        for config in DraftModelRegistry._registry.values()
+    }
+)

lightllm/models/draft_registry.py

@@ -2,6 +2,8 @@
 
 from typing import Callable, Dict, List, Tuple, Type, TypeVar, Union
 
+from lightllm.common.mtp_config import FinalHiddenPolicy, MtpPairConfig
+
 from .registry import load_model_class
 
 
@@ -10,15 +12,24 @@
 
 class _DraftModelRegistry:
     def __init__(self):
-        self._registry: Dict[Tuple[str, str], Union[Type, str]] = {}
+        self._registry: Dict[Tuple[str, str], MtpPairConfig] = {}
 
     def __call__(
         self,
         model_type: Union[str, List[str], Tuple[str, ...]],
         spec_modes: Union[str, List[str], Tuple[str, ...]],
+        *,
+        main_final_hidden_policy: FinalHiddenPolicy,
+        draft_final_hidden_policy: FinalHiddenPolicy,
     ) -> Callable[[T], T]:
         def decorator(model_class: T) -> T:
-            self.register(model_type, spec_modes, model_class)
+            self.register(
+                model_type,
+                spec_modes,
+                model_class,
+                main_final_hidden_policy=main_final_hidden_policy,
+                draft_final_hidden_policy=draft_final_hidden_policy,
+            )
             return model_class
 
         return decorator
@@ -28,25 +39,31 @@ def register(
         model_type: Union[str, List[str], Tuple[str, ...]],
         spec_modes: Union[str, List[str], Tuple[str, ...]],
         model_class: Union[Type, str],
+        *,
+        main_final_hidden_policy: FinalHiddenPolicy,
+        draft_final_hidden_policy: FinalHiddenPolicy,
     ) -> None:
+        config = MtpPairConfig(model_class, main_final_hidden_policy, draft_final_hidden_policy)
         model_types = (model_type,) if isinstance(model_type, str) else tuple(model_type)
         modes = (spec_modes,) if isinstance(spec_modes, str) else tuple(spec_modes)
         for current_model_type in model_types:
             for spec_mode in modes:
                 key = (current_model_type, spec_mode)
                 if key in self._registry:
                     raise ValueError(f"Duplicate draft model registration: {key}")
-                self._registry[key] = model_class
+                self._registry[key] = config
 
-    def get_model_class(self, model_cfg: dict, spec_mode: str) -> Type:
+    def get_pair_config(self, model_cfg: dict, spec_mode: str) -> MtpPairConfig:
         model_type = model_cfg.get("model_type", "")
         try:
-            model_class = self._registry[(model_type, spec_mode)]
+            return self._registry[(model_type, spec_mode)]
         except KeyError:
             raise ValueError(
                 f"Unsupported speculative draft model: mode={spec_mode}, model_type={model_type}"
             ) from None
-        return load_model_class(model_class)
+
+    def get_model_class(self, model_cfg: dict, spec_mode: str) -> Type:
+        return load_model_class(self.get_pair_config(model_cfg, spec_mode).draft_model_class)
 
 
 DraftModelRegistry = _DraftModelRegistry()

lightllm/models/glm4_moe_lite_mtp/model.py

@@ -9,7 +9,6 @@
 
 
 class Glm4MoeLiteMTPModel(Glm4MoeLiteTpPartModel):
-
     # MTP draft model marker (consumed by the decode CUDA-graph / padding paths).
     is_mtp_draft_model = True
 

lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py

@@ -13,7 +13,6 @@ def __init__(self, network_config):
     def _fuse_hidden(self, input_embeddings, infer_state, layer_weight):
         previous_hidden = infer_state.mtp_draft_input_hiddens
         assert input_embeddings.shape[0] == previous_hidden.shape[0]
-        layer_weight.main_norm_weight_(input=previous_hidden, eps=self.eps_, out=previous_hidden)
         layer_weight.enorm_weight_(input=input_embeddings, eps=self.eps_, out=input_embeddings)
         layer_weight.hnorm_weight_(input=previous_hidden, eps=self.eps_, out=previous_hidden)
         return layer_weight.eh_proj_weight_.mm(torch.cat((input_embeddings, previous_hidden), dim=-1))

lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py

@@ -24,8 +24,6 @@ def __init__(self, data_type, network_config, quant_cfg):
         )
         self.wte_weight_ = None
         self.lm_head_weight_ = None
-        # Shared with the target model and injected by the draft model.
-        self.main_norm_weight_: RMSNormWeight = None
 
     def load_hf_weights(self, weights):
         add_language_model_aliases(weights)

lightllm/models/glm5_next_mtp/model.py

@@ -33,7 +33,6 @@ def _init_weights(self, start_layer_index=None):
         self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, self.config, self.quant_cfg)
         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_
         self.trans_layers_weight = [
             self.transformer_weight_class(self.config["num_hidden_layers"], self.data_type, self.config, self.quant_cfg)
         ]