feat(qwen3_5_mtp): Qwen3.5 / Qwen3.5-MoE MTP draft models

sufubao · sufubao · commit f03a5b0790a0 · 2026-06-15T15:36:03.000+08:00
Self-contained dense (qwen3_5_mtp) and MoE (qwen3_5_moe_mtp) MTP draft packages: each carries its own draft wiring (reuse the main model's req/mem managers + rope caches, is_mtp_draft_model marker) and shares a weight-retarget mixin (mtp.* head, embeddings shared with the main model) plus the MTP pre-layer fuse. No shared model base class.
diff --git a/lightllm/models/qwen3_5/infer_struct.py b/lightllm/models/qwen3_5/infer_struct.py
@@ -1,8 +1,4 @@
-import torch
-from typing import List
-
 from lightllm.models.qwen2_vl.infer_struct import Qwen2VLInferStateInfo
-from lightllm.utils.envs_utils import get_env_start_args
 
 
 class Qwen35InferStateInfo(Qwen2VLInferStateInfo):
@@ -12,8 +8,7 @@ def __init__(self):
 
     def init_some_extra_state(self, model):
         super().init_some_extra_state(model)
-        self.b_att_seq_len = self.b_seq_len
-        mtp_step = get_env_start_args().mtp_step
+        from lightllm.common.basemodel.mtp_verify_extra_state import init_mtp_verify_extra_state
 
-        self.b_buffer_idx = self.b_req_idx * (mtp_step + 1) + self.b_mtp_index
+        init_mtp_verify_extra_state(self)
         return
diff --git a/lightllm/models/qwen3_5_moe_mtp/__init__.py b/lightllm/models/qwen3_5_moe_mtp/__init__.py
@@ -0,0 +1,3 @@
+from lightllm.models.qwen3_5_moe_mtp.model import Qwen3_5MoeMTPModel
+
+__all__ = ["Qwen3_5MoeMTPModel"]
diff --git a/lightllm/models/qwen3_5_moe_mtp/layer_weights/__init__.py b/lightllm/models/qwen3_5_moe_mtp/layer_weights/__init__.py
@@ -0,0 +1,5 @@
+from lightllm.models.qwen3_5_moe_mtp.layer_weights.transformer_layer_weight import (
+    Qwen3_5MoeMTPTransformerLayerWeight,
+)
+
+__all__ = ["Qwen3_5MoeMTPTransformerLayerWeight"]
diff --git a/lightllm/models/qwen3_5_moe_mtp/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen3_5_moe_mtp/layer_weights/transformer_layer_weight.py
@@ -0,0 +1,96 @@
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+    COLMMWeight,
+    FusedMoeWeight,
+    ROWMMWeight,
+)
+from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import (
+    Qwen35MOETransformerLayerWeight,
+)
+from lightllm.models.qwen3_5_mtp.layer_weights.mtp_retarget_mixin import MTPRetargetMixin
+from lightllm.utils.envs_utils import get_env_start_args
+
+
+class Qwen3_5MoeMTPTransformerLayerWeight(MTPRetargetMixin, Qwen35MOETransformerLayerWeight):
+    def _init_weight_names(self):
+        super()._init_weight_names()
+        self._retarget_attn_norm_names()
+
+    def _init_moe(self):
+        moe_intermediate_size = self.network_config_["moe_intermediate_size"]
+        self.moe_gate = ROWMMWeight(
+            in_dim=self.network_config_["hidden_size"],
+            out_dims=[self.n_routed_experts],
+            weight_names=f"{self._MTP_PREFIX}{self.layer_num_}.mlp.gate.weight",
+            data_type=self.data_type_,
+            quant_method=None,
+            tp_rank=0,
+            tp_world_size=1,
+        )
+        self.experts = FusedMoeWeight(
+            gate_proj_name="gate_proj",
+            down_proj_name="down_proj",
+            up_proj_name="up_proj",
+            e_score_correction_bias_name="",
+            weight_prefix=f"{self._MTP_PREFIX}{self.layer_num_}.mlp.experts",
+            n_routed_experts=self.n_routed_experts,
+            hidden_size=self.network_config_["hidden_size"],
+            moe_intermediate_size=moe_intermediate_size,
+            data_type=self.data_type_,
+            quant_method=self.quant_cfg.get_quant_method(self.layer_num_, "fused_moe"),
+            layer_num=self.layer_num_,
+            network_config=self.network_config_,
+        )
+        self._init_gated_ffn()
+
+    def _init_gated_ffn(self):
+        hidden_size = self.network_config_["hidden_size"]
+        if "shared_expert_intermediate_size" not in self.network_config_:
+            return
+
+        prefix = f"{self._MTP_PREFIX}{self.layer_num_}.mlp.shared_expert"
+        inter_size = self.network_config_["shared_expert_intermediate_size"]
+        if get_env_start_args().enable_ep_moe:
+            self.gate_up_proj = ROWMMWeight(
+                in_dim=hidden_size,
+                out_dims=[inter_size, inter_size],
+                weight_names=[f"{prefix}.gate_proj.weight", f"{prefix}.up_proj.weight"],
+                data_type=self.data_type_,
+                quant_method=self.get_quant_method("gate_up_proj"),
+                tp_rank=0,
+                tp_world_size=1,
+            )
+            self.down_proj = COLMMWeight(
+                in_dim=inter_size,
+                out_dims=[hidden_size],
+                weight_names=f"{prefix}.down_proj.weight",
+                data_type=self.data_type_,
+                quant_method=self.get_quant_method("down_proj"),
+                tp_rank=0,
+                tp_world_size=1,
+            )
+        else:
+            self.gate_up_proj = ROWMMWeight(
+                in_dim=hidden_size,
+                out_dims=[inter_size, inter_size],
+                weight_names=[f"{prefix}.gate_proj.weight", f"{prefix}.up_proj.weight"],
+                data_type=self.data_type_,
+                quant_method=self.get_quant_method("gate_up_proj"),
+            )
+            self.down_proj = COLMMWeight(
+                in_dim=inter_size,
+                out_dims=[hidden_size],
+                weight_names=f"{prefix}.down_proj.weight",
+                data_type=self.data_type_,
+                quant_method=self.get_quant_method("down_proj"),
+            )
+
+        self.ffn_gate = ROWMMWeight(
+            in_dim=hidden_size,
+            out_dims=[1],
+            weight_names=f"{self._MTP_PREFIX}{self.layer_num_}.mlp.shared_expert_gate.weight",
+            data_type=self.data_type_,
+            bias_names=None,
+            quant_method=None,
+            tp_rank=0,
+            tp_world_size=1,
+        )
diff --git a/lightllm/models/qwen3_5_moe_mtp/model.py b/lightllm/models/qwen3_5_moe_mtp/model.py
@@ -0,0 +1,8 @@
+from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel
+from lightllm.models.qwen3_5_moe_mtp.layer_weights.transformer_layer_weight import (
+    Qwen3_5MoeMTPTransformerLayerWeight,
+)
+
+
+class Qwen3_5MoeMTPModel(Qwen3_5MTPModel):
+    transformer_weight_class = Qwen3_5MoeMTPTransformerLayerWeight
diff --git a/lightllm/models/qwen3_5_mtp/__init__.py b/lightllm/models/qwen3_5_mtp/__init__.py
diff --git a/lightllm/models/qwen3_5_mtp/layer_infer/__init__.py b/lightllm/models/qwen3_5_mtp/layer_infer/__init__.py
diff --git a/lightllm/models/qwen3_5_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/qwen3_5_mtp/layer_infer/pre_layer_infer.py
@@ -0,0 +1,40 @@
+import torch
+
+from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer
+from lightllm.models.qwen3_5_mtp.layer_weights.pre_and_post_layer_weight import Qwen3_5MTPPreAndPostLayerWeight
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+
+class Qwen3_5MTPPreLayerInfer(Qwen3VLMultimodalPreLayerInfer):
+    def __init__(self, network_config):
+        super().__init__(network_config)
+        self.eps_ = network_config["rms_norm_eps"]
+        self.hidden_size = network_config["hidden_size"]
+        return
+
+    def _mtp_fuse(
+        self,
+        input_embdings: torch.Tensor,
+        infer_state: LlamaInferStateInfo,
+        layer_weight: Qwen3_5MTPPreAndPostLayerWeight,
+    ) -> torch.Tensor:
+        tgt_embdings = infer_state.mtp_draft_input_hiddens
+        assert (
+            input_embdings.shape[0] == tgt_embdings.shape[0]
+        ), f"shape {input_embdings.shape} != shape {tgt_embdings.shape}"
+
+        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)
+
+        return layer_weight.eh_proj_weight_.mm(cat_embdings)
+
+    def context_forward(
+        self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: Qwen3_5MTPPreAndPostLayerWeight
+    ):
+        input_embdings = super().context_forward(input_ids, infer_state, layer_weight)
+        return self._mtp_fuse(input_embdings, infer_state, layer_weight)
+
+    def token_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: Qwen3_5MTPPreAndPostLayerWeight):
+        input_embdings = super().token_forward(input_ids, infer_state, layer_weight)
+        return self._mtp_fuse(input_embdings, infer_state, layer_weight)
diff --git a/lightllm/models/qwen3_5_mtp/layer_weights/__init__.py b/lightllm/models/qwen3_5_mtp/layer_weights/__init__.py
diff --git a/lightllm/models/qwen3_5_mtp/layer_weights/mtp_retarget_mixin.py b/lightllm/models/qwen3_5_mtp/layer_weights/mtp_retarget_mixin.py
@@ -0,0 +1,61 @@
+from lightllm.common.basemodel.layer_weights.meta_weights import ROWMMWeight, QKVROWNMMWeight
+
+
+class MTPRetargetMixin:
+    """Shared MTP weight-name retargeting (model.layers.* -> mtp.layers.*) and qkv/o_gate wiring,
+    used by both the dense and MoE Qwen3.5 MTP layer-weight classes (#11). The dense subclass adds
+    its dense-MLP retargets on top; the MoE subclass must not (it uses fused experts)."""
+
+    _MAIN_PREFIX = "model.layers."
+    _MTP_PREFIX = "mtp.layers."
+
+    _ATTN_NORM_NAME_ATTRS = (
+        "_q_weight_name",
+        "_q_norm_name",
+        "_q_bias_name",
+        "_k_weight_name",
+        "_k_norm_name",
+        "_k_bias_name",
+        "_v_weight_name",
+        "_v_bias_name",
+        "_kv_weight_name",
+        "_kv_bias_name",
+        "_o_weight_name",
+        "_o_bias_name",
+        "_att_norm_weight_name",
+        "_att_norm_bias_name",
+        "_ffn_norm_weight_name",
+        "_ffn_norm_bias_name",
+    )
+
+    def _retarget(self, name):
+        if name is None:
+            return None
+        return name.replace(self._MAIN_PREFIX, self._MTP_PREFIX, 1)
+
+    def _retarget_attn_norm_names(self):
+        for attr in self._ATTN_NORM_NAME_ATTRS:
+            setattr(self, attr, self._retarget(getattr(self, attr)))
+
+    def _init_qkv(self):
+        in_dim = self.n_embed
+        q_out_dim = self.q_head_num_ * self.head_dim
+        self.qkv_proj = QKVROWNMMWeight(
+            in_dim=in_dim,
+            q_head_num=self.q_head_num_,
+            kv_head_num=self.k_head_num_,
+            head_dim=self.head_dim,
+            weight_names=[self._q_weight_name, self._k_weight_name, self._v_weight_name],
+            data_type=self.data_type_,
+            bias_names=[self._q_bias_name, self._k_bias_name, self._v_bias_name],
+            quant_method=self.get_quant_method("qkv_proj"),
+        )
+        self._o_gate_weight_name = f"{self._MTP_PREFIX}{self.layer_num_}.self_attn.o_gate_proj.weight"
+        self._o_gate_proj = ROWMMWeight(
+            in_dim=in_dim,
+            out_dims=[q_out_dim],
+            weight_names=[self._o_gate_weight_name],
+            data_type=self.data_type_,
+            bias_names=None,
+            quant_method=self.get_quant_method("o_gate_proj"),
+        )
diff --git a/lightllm/models/qwen3_5_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen3_5_mtp/layer_weights/pre_and_post_layer_weight.py
@@ -0,0 +1,45 @@
+from lightllm.common.basemodel import PreAndPostLayerWeight
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+    EmbeddingWeight,
+    LMHeadWeight,
+    NoTpGEMMANormWeight,
+    ROWMMWeight,
+)
+from lightllm.common.quantization import Quantcfg
+
+
+class Qwen3_5MTPPreAndPostLayerWeight(PreAndPostLayerWeight):
+    def __init__(self, data_type, network_config, quant_cfg: Quantcfg):
+        super().__init__(data_type, network_config)
+        self.quant_cfg: Quantcfg = quant_cfg
+        hidden_size = network_config["hidden_size"]
+
+        self.eh_proj_weight_ = ROWMMWeight(
+            in_dim=hidden_size * 2,
+            out_dims=[hidden_size],
+            weight_names="mtp.fc.weight",
+            data_type=self.data_type_,
+            quant_method=self.quant_cfg.get_quant_method(0, "eh_proj"),
+            tp_rank=0,
+            tp_world_size=1,
+        )
+        self.enorm_weight_ = NoTpGEMMANormWeight(
+            dim=hidden_size,
+            weight_name="mtp.pre_fc_norm_embedding.weight",
+            data_type=self.data_type_,
+        )
+        self.hnorm_weight_ = NoTpGEMMANormWeight(
+            dim=hidden_size,
+            weight_name="mtp.pre_fc_norm_hidden.weight",
+            data_type=self.data_type_,
+        )
+        self.final_norm_weight_ = NoTpGEMMANormWeight(
+            dim=hidden_size,
+            weight_name="mtp.norm.weight",
+            data_type=self.data_type_,
+        )
+
+        # 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
+        return
diff --git a/lightllm/models/qwen3_5_mtp/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen3_5_mtp/layer_weights/transformer_layer_weight.py
@@ -0,0 +1,23 @@
+from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import (
+    Qwen35TransformerLayerWeight,
+)
+from lightllm.models.qwen3_5_mtp.layer_weights.mtp_retarget_mixin import MTPRetargetMixin
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class Qwen3_5MTPTransformerLayerWeight(MTPRetargetMixin, Qwen35TransformerLayerWeight):
+    def _init_weight_names(self):
+        super()._init_weight_names()
+        # Retarget all main-model layer key names to the mtp.* namespace.
+        self._retarget_attn_norm_names()
+        # MLP (dense) projection names retargeted by Qwen35TransformerLayerWeight.
+        self._gate_weight_name = self._retarget(self._gate_weight_name)
+        self._gate_bias_name = self._retarget(self._gate_bias_name)
+        self._up_weight_name = self._retarget(self._up_weight_name)
+        self._up_bias_name = self._retarget(self._up_bias_name)
+        self._gate_up_weight_name = self._retarget(self._gate_up_weight_name)
+        self._gate_up_bias_name = self._retarget(self._gate_up_bias_name)
+        self._down_weight_name = self._retarget(self._down_weight_name)
+        self._down_bias_name = self._retarget(self._down_bias_name)
diff --git a/lightllm/models/qwen3_5_mtp/model.py b/lightllm/models/qwen3_5_mtp/model.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+from lightllm.models.qwen3_5_moe_mtp.model import Qwen3_5MoeMTPModel`
	`2`	`+`
	`3`	`+__all__ = ["Qwen3_5MoeMTPModel"]`