fix(mtp): stack per-expert MTP weights for MoE models in sanitize()

MTP layers in MoE models (35B-A3B, 122B-A10B) ship unfused per-expert weights (mtp.layers.{l}.mlp.experts.{i}.gate_proj.weight) whereas the backbone uses pre-fused switch_mlp format. Conversion was failing with ~768 parameters not in model.

Add a stacking loop in qwen3_5_moe.py sanitize() after the backbone expert loop, mirroring the same pattern for MTP prefixes.

Co-authored-by: Thump604 <thump604@users.noreply.github.com>

Author: AirRunner

mlx_lm/models/qwen3_5_moe.py

@@ -2,6 +2,8 @@
 
 from dataclasses import dataclass
 
+import mlx.core as mx
+
 from .base import BaseModelArgs
 from .qwen3_5 import Model as Qwen3_5Model
 
@@ -49,4 +51,20 @@ def sanitize(self, weights):
                     f"{prefix}.experts.down_proj"
                 )
 
+        # Stack per-expert MTP weights into switch_mlp format.
+        # MTP layers use unfused per-expert weights (experts.{i}.gate_proj etc)
+        # unlike backbone layers which use fused gate_up_proj.
+        mtp_num = getattr(self.language_model.args, "mtp_num_hidden_layers", 0)
+        num_experts = self.language_model.args.num_experts
+        for l in range(mtp_num):
+            prefix = f"language_model.mtp.layers.{l}.mlp"
+            test_key = f"{prefix}.experts.0.gate_proj.weight"
+            if test_key in new_weights:
+                for n in ["gate_proj", "up_proj", "down_proj"]:
+                    to_join = [
+                        new_weights.pop(f"{prefix}.experts.{e}.{n}.weight")
+                        for e in range(num_experts)
+                    ]
+                    new_weights[f"{prefix}.switch_mlp.{n}.weight"] = mx.stack(to_join)
+
         return self.language_model.sanitize(new_weights)