Enable MTP support for Qwen3.5

src/transformers/generation/candidate_generator.py

@@ -1512,15 +1512,15 @@ def get_candidates(
 
             # We need the last invalidated tokens, as well as the new one in a single passcfor efficiency
             mtp_input_ids = input_ids[:, -num_last_main_model_tokens - self.num_mtp_layers :]
-            mtp_position_ids = model_kwargs["position_ids"][:, -num_last_main_model_tokens - self.num_mtp_layers :]
+            mtp_position_ids = model_kwargs["position_ids"][..., -num_last_main_model_tokens - self.num_mtp_layers :]
             mtp_attention_mask = model_kwargs["attention_mask"][:, -num_last_main_model_tokens - self.num_mtp_layers :]
             last_hidden_states = self.full_seq_last_hidden_states[
                 :, -num_last_main_model_tokens - self.num_mtp_layers :, :
             ]
         # If we have a single layer, or validated enough, we can simply pass the new tokens
         else:
             mtp_input_ids = input_ids[:, -num_last_main_model_tokens:]
-            mtp_position_ids = model_kwargs["position_ids"][:, -num_last_main_model_tokens:]
+            mtp_position_ids = model_kwargs["position_ids"][..., -num_last_main_model_tokens:]
             mtp_attention_mask = model_kwargs["attention_mask"][:, -num_last_main_model_tokens:]
 
         candidate_ids, candidate_logits, _ = self.mtp_model(

src/transformers/modeling_layers.py

@@ -502,18 +502,27 @@ def forward(
         for i, mtp_layer in enumerate(self.layers):
             # We need to recompute those every layer since they change
             inputs_embeds = self.embed_tokens(input_ids).to(last_hidden_states.device)
+
+            # Support both 2D position_ids and MRoPE-style nD ids (e.g. Qwen3.5 uses [4, B, L])
+            if position_ids.ndim == 2:
+                text_position_ids = position_ids
+                rope_position_ids = position_ids
+            else:
+                text_position_ids = position_ids[0]  # first component is text positions
+                rope_position_ids = position_ids[1:] if position_ids.shape[0] > 1 else position_ids
+
             position_embeddings = (
-                self.rotary_emb(inputs_embeds, position_ids=position_ids) if self.rotary_emb is not None else None
+                self.rotary_emb(inputs_embeds, position_ids=rope_position_ids) if self.rotary_emb is not None else None
             )
 
             # In full generality, we may need to recompute masks for every layer due to the position offset of each layer
-            masks = self.create_masks_for_mtp_layer(i, inputs_embeds, mtp_cache, position_ids)
+            masks = self.create_masks_for_mtp_layer(i, inputs_embeds, mtp_cache, text_position_ids)
 
             last_hidden_states = mtp_layer(
                 inputs_embeds,
                 last_hidden_states,
                 position_embeddings=position_embeddings,
-                position_ids=position_ids,
+                position_ids=text_position_ids,
                 past_key_values=mtp_cache,
                 **masks,
                 **kwargs,
@@ -549,7 +558,8 @@ def forward(
             # Roll by 1 and append for next layer
             input_ids = torch.cat([input_ids[:, 1:], next_mtp_token], dim=-1)
             attention_mask = torch.cat([attention_mask[:, 1:], attention_mask.new_ones(batch_size, 1)], dim=-1)  # type: ignore
-            position_ids = torch.cat([position_ids[:, 1:], position_ids[:, -1:] + 1], dim=-1)
+            next_position_ids = position_ids[..., -1:] + 1
+            position_ids = torch.cat([position_ids[..., 1:], next_position_ids], dim=-1)
 
             # Need to cat ful_ids as well for the processors
             if full_input_ids is not None:
@@ -564,12 +574,15 @@ def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs)
         pretrained_model_name_or_path = main_model.config.name_or_path
         num_hidden_layers = main_model.config.get_text_config().num_hidden_layers
         # Heuristic: the main model should have the mtp layer patterns under `_keys_to_ignore_on_load_unexpected` to avoid
-        # loading them by default, so use it to later load the correct keys from the checkpoints
+        # loading them by default, so use it to later load the correct keys from the checkpoints.
+        # Filter out non-MTP patterns (e.g. vision tower weights for multimodal models like Qwen3.5).
         mtp_patterns = main_model._keys_to_ignore_on_load_unexpected.copy()  # type: ignore
-        # Due to different released checkpoints, only keep the ones with layer number >= num_hidden_layers - otherwise
-        # mtp layers in a smaller checkpoints could be wrongly added as a 2nd mtp layer of a bigger checkpoint
         final_mtp_patterns = []
         for pattern in mtp_patterns:
+            has_mtp_keyword = "mtp" in pattern.lower()
+            has_layer_pattern = re.search(r"layers?\.\d+", pattern, re.IGNORECASE) is not None
+            if not has_mtp_keyword and not has_layer_pattern:
+                continue
             match_object = re.search(r"\.(\d+)", pattern)
             if match_object is not None and int(match_object.group(1)) < num_hidden_layers:
                 continue

src/transformers/models/qwen3_5/configuration_qwen3_5.py

@@ -38,6 +38,9 @@ class Qwen3_5TextConfig(PreTrainedConfig):
         Number of key heads used in linear attention layers.
     linear_num_value_heads (`int`, *optional*, defaults to 32):
         Number of value heads used in linear attention layers.
+    mtp_num_hidden_layers (`int`, *optional*, defaults to 0):
+        Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled.
+        This field is exposed as `num_mtp_layers` via `attribute_map`.
 
     ```python
     >>> from transformers import Qwen3_5TextModel, Qwen3_5TextConfig
@@ -105,7 +108,10 @@ class Qwen3_5TextConfig(PreTrainedConfig):
     bos_token_id: int | None = None
     eos_token_id: int | list[int] | None = None
     base_config_key = "text_config"
+
+    attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"}
     ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"}
+    mtp_num_hidden_layers: int = 0
 
     def __post_init__(self, **kwargs):
         kwargs.setdefault("partial_rotary_factor", 0.25)  # assign default for BC
@@ -120,6 +126,13 @@ def __post_init__(self, **kwargs):
 
         super().__post_init__(**kwargs)
 
+    # MTP layers always use full attention
+    @property
+    def mtp_layer_types(self):
+        if self.num_mtp_layers is not None:
+            return ["full_attention"] * self.num_mtp_layers
+        return None
+
 
 @auto_docstring(checkpoint="Qwen/Qwen3.5-27B")
 @strict
@@ -152,6 +165,10 @@ class Qwen3_5VisionConfig(PreTrainedConfig):
 @strict
 class Qwen3_5Config(PreTrainedConfig):
     r"""
+    mtp_num_hidden_layers (`int`, *optional*, defaults to 0):
+        Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled.
+        This field is exposed as `num_mtp_layers` via `attribute_map`.
+
     Example:
 
     ```python
@@ -180,6 +197,10 @@ class Qwen3_5Config(PreTrainedConfig):
     vision_end_token_id: int = 248054
     tie_word_embeddings: bool = False
 
+    attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"}
+
+    mtp_num_hidden_layers: int = 0
+
     def __post_init__(self, **kwargs):
         if isinstance(self.vision_config, dict):
             # old ckpt with incorrect model type -> override manually
@@ -196,5 +217,12 @@ def __post_init__(self, **kwargs):
 
         super().__post_init__(**kwargs)
 
+    # MTP layers always use full attention
+    @property
+    def mtp_layer_types(self):
+        if self.num_mtp_layers is not None:
+            return ["full_attention"] * self.num_mtp_layers
+        return None
+
 
 __all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"]

src/transformers/models/qwen3_5/modular_qwen3_5.py

@@ -81,6 +81,9 @@ class Qwen3_5TextConfig(Qwen3NextConfig):
         Number of key heads used in linear attention layers.
     linear_num_value_heads (`int`, *optional*, defaults to 32):
         Number of value heads used in linear attention layers.
+    mtp_num_hidden_layers (`int`, *optional*, defaults to 0):
+        Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled.
+        This field is exposed as `num_mtp_layers` via `attribute_map`.
 
     ```python
     >>> from transformers import Qwen3_5TextModel, Qwen3_5TextConfig
@@ -99,6 +102,8 @@ class Qwen3_5TextConfig(Qwen3NextConfig):
     model_type = "qwen3_5_text"
     base_config_key = "text_config"
 
+    attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"}
+
     base_model_tp_plan = {
         "layers.*.self_attn.q_proj": "colwise",
         "layers.*.self_attn.k_proj": "colwise",
@@ -123,6 +128,7 @@ class Qwen3_5TextConfig(Qwen3NextConfig):
     intermediate_size: int = 12288
     num_hidden_layers: int = 32
     num_key_value_heads: int = 4
+    mtp_num_hidden_layers: int = 0
 
     decoder_sparse_step = AttributeError()
     norm_topk_prob = AttributeError()
@@ -138,6 +144,13 @@ def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
         del self.mlp_only_layers
 
+    # MTP layers always use full attention
+    @property
+    def mtp_layer_types(self):
+        if self.num_mtp_layers is not None:
+            return ["full_attention"] * self.num_mtp_layers
+        return None
+
 
 @auto_docstring(checkpoint="Qwen/Qwen3.5-27B")
 @strict
@@ -156,6 +169,10 @@ class Qwen3_5VisionConfig(Qwen3VLVisionConfig):
 @strict
 class Qwen3_5Config(Qwen3VLConfig):
     r"""
+    mtp_num_hidden_layers (`int`, *optional*, defaults to 0):
+        Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled.
+        This field is exposed as `num_mtp_layers` via `attribute_map`.
+
     Example:
 
     ```python
@@ -176,6 +193,17 @@ class Qwen3_5Config(Qwen3VLConfig):
     vision_start_token_id: int = 248053
     vision_end_token_id: int = 248054
 
+    attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"}
+
+    mtp_num_hidden_layers: int = 0
+
+    # MTP layers always use full attention
+    @property
+    def mtp_layer_types(self):
+        if self.num_mtp_layers is not None:
+            return ["full_attention"] * self.num_mtp_layers
+        return None
+
 
 class Qwen3_5VisionRotaryEmbedding(Qwen3VLVisionRotaryEmbedding):
     pass

tests/generation/test_utils.py

@@ -2488,6 +2488,11 @@ def test_generate_with_mtp(self):
             ):
                 self.skipTest("No MTP keys registered")
 
+            # Assisted generation (use_mtp) is blocked for stateful models, whose recurrent state
+            # cannot be rolled back during candidate verification
+            if getattr(model_class, "_is_stateful", False):
+                self.skipTest("MTP generation is not supported for stateful models")
+
             config.get_text_config().num_mtp_layers = 1
             model = model_class(config).to(torch_device).eval()
             mtp_model = MtpModel(model, num_mtp_layers=1)