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

Add the MTP draft model packages and register them:

- qwen3_5_mtp: a forced single full-attn-layer draft model, with the
  MTP pre-layer infer (embed/hidden norm + fc fusion) and pre/post +
  transformer-layer weight loaders reading the mtp.* namespace.
- qwen3_5_moe_mtp: the MoE variant draft weight loaders + model.
- register qwen3_5 / qwen3_5_moe MTP draft models with per-block
  draft_idx, plus the qwen3_5 verify infer_struct.

Unit tests scaffold the MTP draft layer and the hybrid verify forward.

Author: sufubao

lightllm/models/qwen3_5/infer_struct.py

@@ -16,4 +16,39 @@ def init_some_extra_state(self, model):
         mtp_step = get_env_start_args().mtp_step
 
         self.b_buffer_idx = self.b_req_idx * (mtp_step + 1) + self.b_mtp_index
+        # conv buffer is now ONE widened slot per request (indexed by req_idx),
+        # dropping the *(S+1) + mtp_index addressing used by the SSM block.
+        self.b_conv_buffer_idx = self.b_req_idx
+        # MTP verify batch: decode-mode, S+1 expanded, and gated on the
+        # per-real-request accept tensor that decode_mtp threads in. Gating on
+        # b_num_accepted_tokens (vs only b_mtp_index, which is set for any decode)
+        # distinguishes the main-model verify forward from draft/plain decode.
+        self.is_mtp_verify = (
+            (mtp_step > 0)
+            and (not self.is_prefill)
+            and (self.b_mtp_index is not None)
+            and (self.b_num_accepted_tokens is not None)
+        )
+        self.b_gdn_verify_cu_seqlens = None
+        self.b_ssm_index_rows = None
+        # b_num_accepted_tokens is threaded onto the infer_state from ModelInput by
+        # _create_inferstate (mirrors b_mtp_index) BEFORE this runs; nothing to do here.
+        if self.is_mtp_verify:
+            step = mtp_step + 1
+            n_real = self.b_req_idx.shape[0] // step
+            self.b_gdn_verify_cu_seqlens = torch.arange(
+                0, (n_real + 1) * step, step, dtype=torch.int32, device=self.b_req_idx.device
+            )
+            req_first = self.b_req_idx.view(n_real, step)[:, 0]
+            base = (req_first * step).view(n_real, 1)
+            self.b_ssm_index_rows = base + torch.arange(step, device=base.device, dtype=base.dtype).view(1, step)
+            assert self.b_ssm_index_rows.shape == (n_real, step)
+            # The spec conv kernel is per-SEQUENCE (one program per real request),
+            # indexed by conv_state_indices[idx_seq] with idx_seq in [0, n_real),
+            # aligned 1:1 with b_gdn_verify_cu_seqlens / b_num_accepted_tokens. The
+            # default b_conv_buffer_idx = b_req_idx has the expanded length n_real*step,
+            # which launches n_real*step conv programs and reads num_accepted/
+            # query_start_loc out of bounds for idx_seq >= n_real, corrupting the
+            # committed conv slot. Narrow it to one widened conv slot per request.
+            self.b_conv_buffer_idx = req_first
         return

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"]

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"]

lightllm/models/qwen3_5_moe_mtp/layer_weights/transformer_layer_weight.py

@@ -0,0 +1,142 @@
+from lightllm.common.basemodel.layer_weights.meta_weights import (
+    COLMMWeight,
+    FusedMoeWeight,
+    QKVROWNMMWeight,
+    ROWMMWeight,
+)
+from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import (
+    Qwen35MOETransformerLayerWeight,
+)
+from lightllm.utils.envs_utils import get_env_start_args
+
+
+class Qwen3_5MoeMTPTransformerLayerWeight(Qwen35MOETransformerLayerWeight):
+    _MAIN_PREFIX = "model.layers."
+    _MTP_PREFIX = "mtp.layers."
+
+    def _retarget(self, name):
+        if name is None:
+            return None
+        return name.replace(self._MAIN_PREFIX, self._MTP_PREFIX, 1)
+
+    def _init_weight_names(self):
+        super()._init_weight_names()
+        self._q_weight_name = self._retarget(self._q_weight_name)
+        self._q_norm_name = self._retarget(self._q_norm_name)
+        self._q_bias_name = self._retarget(self._q_bias_name)
+        self._k_weight_name = self._retarget(self._k_weight_name)
+        self._k_norm_name = self._retarget(self._k_norm_name)
+        self._k_bias_name = self._retarget(self._k_bias_name)
+        self._v_weight_name = self._retarget(self._v_weight_name)
+        self._v_bias_name = self._retarget(self._v_bias_name)
+        self._kv_weight_name = self._retarget(self._kv_weight_name)
+        self._kv_bias_name = self._retarget(self._kv_bias_name)
+        self._o_weight_name = self._retarget(self._o_weight_name)
+        self._o_bias_name = self._retarget(self._o_bias_name)
+        self._att_norm_weight_name = self._retarget(self._att_norm_weight_name)
+        self._att_norm_bias_name = self._retarget(self._att_norm_bias_name)
+        self._ffn_norm_weight_name = self._retarget(self._ffn_norm_weight_name)
+        self._ffn_norm_bias_name = self._retarget(self._ffn_norm_bias_name)
+
+    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"),
+        )
+
+    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,
+        )

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

lightllm/models/qwen3_5_mtp/__init__.py

@@ -0,0 +1,41 @@
+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)

lightllm/models/qwen3_5_mtp/layer_infer/__init__.py

@@ -0,0 +1,46 @@
+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